[BEAM-141] Remove StateSampler from the SDK

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/inprocess/InProcessExecutionContext.java

@@ -22,7 +22,6 @@
 import org.apache.beam.sdk.util.BaseExecutionContext;
 import org.apache.beam.sdk.util.ExecutionContext;
 import org.apache.beam.sdk.util.TimerInternals;
-import org.apache.beam.sdk.util.common.worker.StateSampler;
 import org.apache.beam.sdk.util.state.CopyOnAccessInMemoryStateInternals;
 
 /**
@@ -47,8 +46,7 @@ public InProcessExecutionContext(Clock clock, Object key,
   }
 
   @Override
-  protected InProcessStepContext createStepContext(
-      String stepName, String transformName, StateSampler stateSampler) {
+  protected InProcessStepContext createStepContext(String stepName, String transformName) {
     return new InProcessStepContext(this, stepName, transformName);
   }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/inprocess/ParDoInProcessEvaluator.java

@@ -54,7 +54,7 @@ public static <InputT, OutputT> ParDoInProcessEvaluator<InputT> create(
         evaluationContext.getExecutionContext(application, inputBundle.getKey());
     String stepName = evaluationContext.getStepName(application);
     InProcessStepContext stepContext =
-        executionContext.getOrCreateStepContext(stepName, stepName, null);
+        executionContext.getOrCreateStepContext(stepName, stepName);
 
     CounterSet counters = evaluationContext.createCounterSet();
 

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java

@@ -483,15 +483,16 @@ void initializeState() {
       runnerSideInputs = runnerSideInputs.and(entry.getKey().getTagInternal(), entry.getValue());
     }
     outputManager = new DoFnRunnerBase.ListOutputManager();
-    fnRunner = DoFnRunners.createDefault(
-        options,
-        fn,
-        DirectSideInputReader.of(runnerSideInputs),
-        outputManager,
-        mainOutputTag,
-        sideOutputTags,
-        DirectModeExecutionContext.create().getOrCreateStepContext(STEP_NAME, TRANSFORM_NAME, null),
-        counterSet.getAddCounterMutator(),
-        WindowingStrategy.globalDefault());
+    fnRunner =
+        DoFnRunners.createDefault(
+            options,
+            fn,
+            DirectSideInputReader.of(runnerSideInputs),
+            outputManager,
+            mainOutputTag,
+            sideOutputTags,
+            DirectModeExecutionContext.create().getOrCreateStepContext(STEP_NAME, TRANSFORM_NAME),
+            counterSet.getAddCounterMutator(),
+            WindowingStrategy.globalDefault());
   }
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java

@@ -1200,7 +1200,7 @@ private static <InputT, OutputT, ActualInputT extends InputT> void evaluateHelpe
             outputManager,
             mainOutputTag,
             sideOutputTags,
-            executionContext.getOrCreateStepContext(stepName, stepName, null),
+            executionContext.getOrCreateStepContext(stepName, stepName),
             context.getAddCounterMutator(),
             input.getWindowingStrategy());
 

sdks/java/core/src/main/java/org/apache/beam/sdk/util/BaseExecutionContext.java

@@ -19,9 +19,9 @@
 
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.util.common.worker.StateSampler;
 import org.apache.beam.sdk.util.state.StateInternals;
 import org.apache.beam.sdk.values.TupleTag;
+import com.google.common.base.Supplier;
 
 import java.io.IOException;
 import java.util.Collection;
@@ -37,7 +37,7 @@
  * be cached for the lifetime of this {@link ExecutionContext}.
  *
  * <p>BaseExecutionContext is generic to allow implementing subclasses to return a concrete subclass
- * of {@link StepContext} from {@link #getOrCreateStepContext(String, String, StateSampler)} and
+ * of {@link StepContext} from {@link #getOrCreateStepContext(String, String)} and
  * {@link #getAllStepContexts()} without forcing each subclass to override the method, e.g.
  * <pre>
  * @Override
@@ -47,8 +47,8 @@
  * </pre>
  *
  * <p>When a subclass of {@code BaseExecutionContext} has been downcast, the return types of
- * {@link #createStepContext(String, String, StateSampler)},
- * {@link #getOrCreateStepContext(String, String, StateSampler}, and {@link #getAllStepContexts()}
+ * {@link #createStepContext(String, String)},
+ * {@link #getOrCreateStepContext(String, String)}, and {@link #getAllStepContexts()}
  * will be appropriately specialized.
  */
 public abstract class BaseExecutionContext<T extends ExecutionContext.StepContext>
@@ -60,21 +60,32 @@ public abstract class BaseExecutionContext<T extends ExecutionContext.StepContex
    * Implementations should override this to create the specific type
    * of {@link StepContext} they need.
    */
-  protected abstract T createStepContext(
-      String stepName, String transformName, StateSampler stateSampler);
-
+  protected abstract T createStepContext(String stepName, String transformName);
 
   /**
    * Returns the {@link StepContext} associated with the given step.
    */
   @Override
-  public T getOrCreateStepContext(
-      String stepName, String transformName, StateSampler stateSampler) {
+  public T getOrCreateStepContext(String stepName, String transformName) {
+    final String finalStepName = stepName;
+    final String finalTransformName = transformName;
+    return getOrCreateStepContext(
+        stepName,
+        new Supplier<T>() {
+          @Override
+          public T get() {
+            return createStepContext(finalStepName, finalTransformName);
+          }
+        });
+  }
+
+  protected final T getOrCreateStepContext(String stepName, Supplier<T> createContextFunc) {
     T context = cachedStepContexts.get(stepName);
     if (context == null) {
-      context = createStepContext(stepName, transformName, stateSampler);
+      context = createContextFunc.get();
       cachedStepContexts.put(stepName, context);
     }
+
     return context;
   }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/util/DirectModeExecutionContext.java

@@ -20,7 +20,6 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import org.apache.beam.sdk.runners.DirectPipelineRunner.ValueWithMetadata;
-import org.apache.beam.sdk.util.common.worker.StateSampler;
 import org.apache.beam.sdk.util.state.InMemoryStateInternals;
 import org.apache.beam.sdk.util.state.StateInternals;
 import org.apache.beam.sdk.values.TupleTag;
@@ -48,8 +47,7 @@ public static DirectModeExecutionContext create() {
   }
 
   @Override
-  protected StepContext createStepContext(
-      String stepName, String transformName, StateSampler stateSampler) {
+  protected StepContext createStepContext(String stepName, String transformName) {
     return new StepContext(this, stepName, transformName);
   }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/util/ExecutionContext.java

@@ -19,7 +19,6 @@
 
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.util.common.worker.StateSampler;
 import org.apache.beam.sdk.util.state.StateInternals;
 import org.apache.beam.sdk.values.TupleTag;
 
@@ -34,8 +33,7 @@ public interface ExecutionContext {
   /**
    * Returns the {@link StepContext} associated with the given step.
    */
-  StepContext getOrCreateStepContext(
-      String stepName, String transformName, StateSampler stateSampler);
+  StepContext getOrCreateStepContext(String stepName, String transformName);
 
   /**
    * Returns a collection view of all of the {@link StepContext}s.