[Dataflow Streaming] Replace HashBasedTable with HashMap in CachingStateTable (#36743)

CachingStateTable 1 not using the additional functionality provided by HashBasedTable
and HashMap has lower (1 instead of 2) lookup overhead and reduce allocations.

Author: arunpandianp

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/CachingStateTable.java

@@ -17,38 +17,41 @@
  */
 package org.apache.beam.runners.dataflow.worker.windmill.state;
 
+import com.google.auto.value.AutoValue;
 import java.io.Closeable;
+import java.util.HashMap;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import org.apache.beam.runners.core.StateNamespace;
-import org.apache.beam.runners.core.StateTable;
 import org.apache.beam.runners.core.StateTag;
 import org.apache.beam.runners.core.StateTags;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.InternedByteString;
 import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache.ForKeyAndFamily;
 import org.apache.beam.sdk.coders.BooleanCoder;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.state.*;
-import org.apache.beam.sdk.transforms.Combine;
-import org.apache.beam.sdk.transforms.CombineWithContext;
+import org.apache.beam.sdk.transforms.Combine.CombineFn;
+import org.apache.beam.sdk.transforms.CombineWithContext.CombineFnWithContext;
 import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
 import org.apache.beam.sdk.util.CombineFnUtil;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier;
 
-final class CachingStateTable extends StateTable {
+final class CachingStateTable {
 
+  private final HashMap<StateTableKey, WindmillState> stateTable;
   private final String stateFamily;
   private final WindmillStateReader reader;
   private final WindmillStateCache.ForKeyAndFamily cache;
   private final boolean isSystemTable;
   private final Supplier<Closeable> scopedReadStateSupplier;
-  private final @Nullable StateTable derivedStateTable;
+  private final @Nullable CachingStateTable derivedStateTable;
   private final boolean isNewKey;
   private final boolean mapStateViaMultimapState;
   private final WindmillStateTagUtil windmillStateTagUtil;
 
   private CachingStateTable(Builder builder) {
+    this.stateTable = new HashMap<>();
     this.stateFamily = builder.stateFamily;
     this.reader = builder.reader;
     this.cache = builder.cache;
@@ -65,20 +68,45 @@ private CachingStateTable(Builder builder) {
     }
   }
 
-  static CachingStateTable.Builder builder(
+  static Builder builder(
       String stateFamily,
       WindmillStateReader reader,
       ForKeyAndFamily cache,
       boolean isNewKey,
       Supplier<Closeable> scopedReadStateSupplier,
       WindmillStateTagUtil windmillStateTagUtil) {
-    return new CachingStateTable.Builder(
+    return new Builder(
         stateFamily, reader, cache, scopedReadStateSupplier, isNewKey, windmillStateTagUtil);
   }
 
-  @Override
+  /**
+   * Gets the {@link State} in the specified {@link StateNamespace} with the specified {@link
+   * StateTag}, binding it using the {@link #binderForNamespace} if it is not already present in
+   * this {@link CachingStateTable}.
+   */
+  public <StateT extends State> StateT get(
+      StateNamespace namespace, StateTag<StateT> tag, StateContext<?> c) {
+
+    StateTableKey stateTableKey = StateTableKey.create(namespace, tag);
+    @SuppressWarnings("unchecked")
+    StateT storage =
+        (StateT)
+            stateTable.computeIfAbsent(
+                stateTableKey,
+                unusedKey -> (WindmillState) tag.bind(binderForNamespace(namespace, c)));
+    return storage;
+  }
+
+  public void clear() {
+    stateTable.clear();
+  }
+
+  public Iterable<WindmillState> values() {
+    return stateTable.values();
+  }
+
   @SuppressWarnings("deprecation")
-  protected StateTag.StateBinder binderForNamespace(StateNamespace namespace, StateContext<?> c) {
+  private StateTag.StateBinder binderForNamespace(StateNamespace namespace, StateContext<?> c) {
     // Look up state objects in the cache or create new ones if not found.  The state will
     // be added to the cache in persist().
     return new StateTag.StateBinder() {
@@ -190,7 +218,7 @@ public WatermarkHoldState bindWatermark(
       public <InputT, AccumT, OutputT> CombiningState<InputT, AccumT, OutputT> bindCombiningValue(
           StateTag<CombiningState<InputT, AccumT, OutputT>> address,
           Coder<AccumT> accumCoder,
-          Combine.CombineFn<InputT, AccumT, OutputT> combineFn) {
+          CombineFn<InputT, AccumT, OutputT> combineFn) {
         StateTag<CombiningState<InputT, AccumT, OutputT>> addressOrInternalTag =
             addressOrInternalTag(address);
 
@@ -214,7 +242,7 @@ public <InputT, AccumT, OutputT> CombiningState<InputT, AccumT, OutputT> bindCom
           CombiningState<InputT, AccumT, OutputT> bindCombiningValueWithContext(
               StateTag<CombiningState<InputT, AccumT, OutputT>> address,
               Coder<AccumT> accumCoder,
-              CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {
+              CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {
         return bindCombiningValue(
             addressOrInternalTag(address), accumCoder, CombineFnUtil.bindContext(combineFn, c));
       }
@@ -239,6 +267,21 @@ private <T extends State> StateTag<T> addressOrInternalTag(StateTag<T> address)
     };
   }
 
+  @AutoValue
+  abstract static class StateTableKey {
+
+    public abstract StateNamespace getStateNamespace();
+
+    public abstract String getId();
+
+    public static StateTableKey create(StateNamespace namespace, StateTag<?> stateTag) {
+      // TODO(https://github.com/apache/beam/issues/36753): stateTag.getId() returns only the
+      // string tag without system/user prefix. This could cause a collision between system and
+      // user tag with the same id. Consider adding the prefix to state table key.
+      return new AutoValue_CachingStateTable_StateTableKey(namespace, stateTag.getId());
+    }
+  }
+
   static class Builder {
 
     private final String stateFamily;
@@ -248,7 +291,7 @@ static class Builder {
     private final boolean isNewKey;
     private final WindmillStateTagUtil windmillStateTagUtil;
     private boolean isSystemTable;
-    private @Nullable StateTable derivedStateTable;
+    private @Nullable CachingStateTable derivedStateTable;
     private boolean mapStateViaMultimapState = false;
 
     private Builder(
@@ -268,7 +311,7 @@ private Builder(
       this.windmillStateTagUtil = windmillStateTagUtil;
     }
 
-    Builder withDerivedState(StateTable derivedStateTable) {
+    Builder withDerivedState(CachingStateTable derivedStateTable) {
       this.isSystemTable = false;
       this.derivedStateTable = derivedStateTable;
       return this;

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/IdTracker.java

@@ -24,7 +24,6 @@
 import java.util.concurrent.ExecutionException;
 import java.util.function.BiConsumer;
 import org.apache.beam.runners.core.StateNamespace;
-import org.apache.beam.runners.core.StateTable;
 import org.apache.beam.runners.core.StateTag;
 import org.apache.beam.runners.core.StateTags;
 import org.apache.beam.sdk.coders.InstantCoder;
@@ -95,7 +94,7 @@ final class IdTracker {
   // here.
   private final ValueState<Map<Range<Instant>, RangeSet<Instant>>> subRangeDeletionsValue;
 
-  IdTracker(StateTable stateTable, StateNamespace namespace, StateTag<?> spec) {
+  IdTracker(CachingStateTable stateTable, StateNamespace namespace, StateTag<?> spec) {
     StateTag<ValueState<Map<Range<Instant>, RangeSet<Long>>>> idsAvailableTag =
         StateTags.makeSystemTagInternal(
             StateTags.value(spec.getId() + IDS_AVAILABLE_STR, IDS_AVAILABLE_CODER));

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java

@@ -25,7 +25,6 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import org.apache.beam.runners.core.StateNamespace;
-import org.apache.beam.runners.core.StateTable;
 import org.apache.beam.runners.core.StateTag;
 import org.apache.beam.runners.dataflow.worker.WindmillTimeUtils;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.InternedByteString;
@@ -68,7 +67,7 @@ public class WindmillOrderedList<T> extends SimpleWindmillState implements Order
   private boolean cleared = false;
 
   WindmillOrderedList(
-      StateTable derivedStateTable,
+      CachingStateTable derivedStateTable,
       StateNamespace namespace,
       InternedByteString encodeKey,
       StateTag<OrderedListState<T>> spec,

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternals.java

@@ -26,7 +26,6 @@
 import java.util.concurrent.Future;
 import org.apache.beam.runners.core.StateInternals;
 import org.apache.beam.runners.core.StateNamespace;
-import org.apache.beam.runners.core.StateTable;
 import org.apache.beam.runners.core.StateTag;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest;
@@ -52,8 +51,8 @@ public class WindmillStateInternals<K> implements StateInternals {
   private final @Nullable K key;
 
   private final WindmillStateCache.ForKeyAndFamily cache;
-  private final StateTable workItemState;
-  private final StateTable workItemDerivedState;
+  private final CachingStateTable workItemState;
+  private final CachingStateTable workItemDerivedState;
   private final Supplier<Closeable> scopedReadStateSupplier;
 
   public WindmillStateInternals(
@@ -82,17 +81,11 @@ public WindmillStateInternals(
     return key;
   }
 
-  private void persist(List<Future<WorkItemCommitRequest>> commitsToMerge, StateTable stateTable) {
-    for (State location : stateTable.values()) {
-      if (!(location instanceof WindmillState)) {
-        throw new IllegalStateException(
-            String.format(
-                "%s wasn't created by %s -- unable to persist it",
-                location.getClass().getSimpleName(), getClass().getSimpleName()));
-      }
-
+  private void persist(
+      List<Future<WorkItemCommitRequest>> commitsToMerge, CachingStateTable stateTable) {
+    for (WindmillState location : stateTable.values()) {
       try {
-        commitsToMerge.add(((WindmillState) location).persist(cache));
+        commitsToMerge.add(location.persist(cache));
       } catch (IOException e) {
         throw new RuntimeException("Unable to persist state", e);
       }
@@ -102,8 +95,8 @@ private void persist(List<Future<WorkItemCommitRequest>> commitsToMerge, StateTa
     // Clear any references to the underlying reader to prevent space leaks.
     // The next work unit to use these cached State objects will reset the
     // reader to a current reader in case those values are modified.
-    for (State location : stateTable.values()) {
-      ((WindmillState) location).cleanupAfterWorkItem();
+    for (WindmillState location : stateTable.values()) {
+      location.cleanupAfterWorkItem();
     }
 
     // Clear out the map of already retrieved state instances.