change datasampling to modify the consumers and not graph for sampling

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java

@@ -17,7 +17,6 @@
  */
 package org.apache.beam.fn.harness;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.EnumMap;
 import java.util.List;
@@ -33,7 +32,6 @@
 import org.apache.beam.fn.harness.control.ProcessBundleHandler;
 import org.apache.beam.fn.harness.data.BeamFnDataGrpcClient;
 import org.apache.beam.fn.harness.debug.DataSampler;
-import org.apache.beam.fn.harness.debug.DataSamplingDescriptorModifier;
 import org.apache.beam.fn.harness.logging.BeamFnLoggingClient;
 import org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache;
 import org.apache.beam.fn.harness.status.BeamFnStatusClient;
@@ -96,7 +94,7 @@ public class FnHarness {
   private static final String RUNNER_CAPABILITIES = "RUNNER_CAPABILITIES";
   private static final String ENABLE_DATA_SAMPLING_EXPERIMENT = "enable_data_sampling";
   private static final Logger LOG = LoggerFactory.getLogger(FnHarness.class);
-  private static final DataSampler dataSampler = new DataSampler();
+  private static DataSampler dataSampler = new DataSampler();
 
   private static Endpoints.ApiServiceDescriptor getApiServiceDescriptor(String descriptor)
       throws TextFormat.ParseException {
@@ -255,16 +253,12 @@ public static void main(
 
       FinalizeBundleHandler finalizeBundleHandler = new FinalizeBundleHandler(executorService);
 
-      // Add any graph modifications.
-      List<ProcessBundleDescriptorModifier> modifierList = new ArrayList<>();
+      // Create the sampler, if the experiment is enabled.
       Optional<List<String>> experimentList =
           Optional.ofNullable(options.as(ExperimentalOptions.class).getExperiments());
-
-      // If data sampling is enabled, then modify the graph to add any DataSampling Operations.
-      if (experimentList.isPresent()
-          && experimentList.get().contains(ENABLE_DATA_SAMPLING_EXPERIMENT)) {
-        modifierList.add(new DataSamplingDescriptorModifier());
-      }
+      boolean shouldSample =
+          experimentList.isPresent()
+              && experimentList.get().contains(ENABLE_DATA_SAMPLING_EXPERIMENT);
 
       // Retrieves the ProcessBundleDescriptor from cache. Requests the PBD from the Runner if it
       // doesn't exist. Additionally, runs any graph modifications.
@@ -274,29 +268,16 @@ public static void main(
             private final Cache<String, BeamFnApi.ProcessBundleDescriptor> cache =
                 Caches.subCache(processWideCache, PROCESS_BUNDLE_DESCRIPTORS);
 
-            private final List<ProcessBundleDescriptorModifier> modifiers = modifierList;
-
             @Override
             public BeamFnApi.ProcessBundleDescriptor apply(String id) {
               return cache.computeIfAbsent(id, this::loadDescriptor);
             }
 
             private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) {
-              ProcessBundleDescriptor descriptor =
-                  blockingControlStub.getProcessBundleDescriptor(
-                      BeamFnApi.GetProcessBundleDescriptorRequest.newBuilder()
-                          .setProcessBundleDescriptorId(id)
-                          .build());
-              for (ProcessBundleDescriptorModifier modifier : modifiers) {
-                try {
-                  LOG.debug("Modifying graph with " + modifier);
-                  descriptor = modifier.modifyProcessBundleDescriptor(descriptor);
-                } catch (ProcessBundleDescriptorModifier.GraphModificationException e) {
-                  LOG.warn("Could not modify graph with " + modifier + ": " + e.getMessage());
-                }
-              }
-
-              return descriptor;
+              return blockingControlStub.getProcessBundleDescriptor(
+                  BeamFnApi.GetProcessBundleDescriptorRequest.newBuilder()
+                      .setProcessBundleDescriptorId(id)
+                      .build());
             }
           };
 
@@ -313,7 +294,7 @@ private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) {
               metricsShortIds,
               executionStateSampler,
               processWideCache,
-              dataSampler);
+              shouldSample ? dataSampler : null);
       logging.setProcessBundleHandler(processBundleHandler);
 
       BeamFnStatusClient beamFnStatusClient = null;

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java

@@ -20,13 +20,11 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
 import org.apache.beam.fn.harness.control.BundleProgressReporter;
 import org.apache.beam.fn.harness.control.BundleSplitListener;
 import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.debug.DataSampler;
 import org.apache.beam.fn.harness.state.BeamFnStateClient;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest;
 import org.apache.beam.model.pipeline.v1.Endpoints;
@@ -61,9 +59,6 @@ interface Context {
     /** A client for handling state requests. */
     BeamFnStateClient getBeamFnStateClient();
 
-    /** The id of the parent ProcessBundleDescriptor. */
-    String getProcessBundleDescriptorId();
-
     /** The id of the PTransform. */
     String getPTransformId();
 
@@ -168,13 +163,6 @@ <T> void addIncomingTimerEndpoint(
      * instant provides the timeout on how long the finalization callback is valid for.
      */
     DoFn.BundleFinalizer getBundleFinalizer();
-
-    /**
-     * A DataSampler can be used to sample in-flight elements. This is used to plumb a global
-     * DataSampler to DataSampler operations in order to perform said sampling. Only present when
-     * using the "enable_data_sampling" experiment.
-     */
-    Optional<DataSampler> getDataSampler();
   }
 
   /**

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java

@@ -28,7 +28,6 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.WeakHashMap;
@@ -167,7 +166,7 @@ public class ProcessBundleHandler {
   @VisibleForTesting final BundleProcessorCache bundleProcessorCache;
   private final Set<String> runnerCapabilities;
 
-  private final DataSampler dataSampler;
+  private final @Nullable DataSampler dataSampler;
 
   public ProcessBundleHandler(
       PipelineOptions options,
@@ -179,7 +178,7 @@ public ProcessBundleHandler(
       ShortIdMap shortIds,
       ExecutionStateSampler executionStateSampler,
       Cache<Object, Object> processWideCache,
-      DataSampler dataSampler) {
+      @Nullable DataSampler dataSampler) {
     this(
         options,
         runnerCapabilities,
@@ -208,7 +207,7 @@ public ProcessBundleHandler(
       Map<String, PTransformRunnerFactory> urnToPTransformRunnerFactoryMap,
       Cache<Object, Object> processWideCache,
       BundleProcessorCache bundleProcessorCache,
-      DataSampler dataSampler) {
+      @Nullable DataSampler dataSampler) {
     this.options = options;
     this.fnApiRegistry = fnApiRegistry;
     this.beamFnDataClient = beamFnDataClient;
@@ -328,11 +327,6 @@ public BeamFnStateClient getBeamFnStateClient() {
                       return beamFnStateClient;
                     }
 
-                    @Override
-                    public String getProcessBundleDescriptorId() {
-                      return processBundleDescriptor.getId();
-                    }
-
                     @Override
                     public String getPTransformId() {
                       return pTransformId;
@@ -494,11 +488,6 @@ public BundleSplitListener getSplitListener() {
                     public BundleFinalizer getBundleFinalizer() {
                       return bundleFinalizer;
                     }
-
-                    @Override
-                    public Optional<DataSampler> getDataSampler() {
-                      return Optional.ofNullable(dataSampler);
-                    }
                   });
       if (runner instanceof BeamFnDataReadRunner) {
         channelRoots.add((BeamFnDataReadRunner) runner);
@@ -789,7 +778,7 @@ private BundleProcessor createBundleProcessor(
     bundleProgressReporterAndRegistrar.register(stateTracker);
     PCollectionConsumerRegistry pCollectionConsumerRegistry =
         new PCollectionConsumerRegistry(
-            stateTracker, shortIds, bundleProgressReporterAndRegistrar, bundleDescriptor);
+            stateTracker, shortIds, bundleProgressReporterAndRegistrar, bundleDescriptor, dataSampler);
     HashSet<String> processedPTransformIds = new HashSet<>();
 
     PTransformFunctionRegistry startFunctionRegistry =

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java

@@ -24,13 +24,16 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
+import javax.annotation.Nullable;
 import org.apache.beam.fn.harness.HandlesSplits;
 import org.apache.beam.fn.harness.control.BundleProgressReporter;
 import org.apache.beam.fn.harness.control.ExecutionStateSampler.ExecutionState;
 import org.apache.beam.fn.harness.control.ExecutionStateSampler.ExecutionStateTracker;
 import org.apache.beam.fn.harness.control.Metrics;
 import org.apache.beam.fn.harness.control.Metrics.BundleCounter;
 import org.apache.beam.fn.harness.control.Metrics.BundleDistribution;
+import org.apache.beam.fn.harness.debug.DataSampler;
+import org.apache.beam.fn.harness.debug.OutputSampler;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor;
 import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
 import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -87,12 +90,27 @@ public static ConsumerAndMetadata forConsumer(
   private final BundleProgressReporter.Registrar bundleProgressReporterRegistrar;
   private final ProcessBundleDescriptor processBundleDescriptor;
   private final RehydratedComponents rehydratedComponents;
+  private final @Nullable DataSampler dataSampler;
 
   public PCollectionConsumerRegistry(
       ExecutionStateTracker stateTracker,
       ShortIdMap shortIdMap,
       BundleProgressReporter.Registrar bundleProgressReporterRegistrar,
       ProcessBundleDescriptor processBundleDescriptor) {
+    this(
+        stateTracker,
+        shortIdMap,
+        bundleProgressReporterRegistrar,
+        processBundleDescriptor,
+        null);
+  }
+
+  public PCollectionConsumerRegistry(
+      ExecutionStateTracker stateTracker,
+      ShortIdMap shortIdMap,
+      BundleProgressReporter.Registrar bundleProgressReporterRegistrar,
+      ProcessBundleDescriptor processBundleDescriptor,
+      @Nullable DataSampler dataSampler) {
     this.stateTracker = stateTracker;
     this.shortIdMap = shortIdMap;
     this.pCollectionIdsToConsumers = new HashMap<>();
@@ -106,6 +124,7 @@ public PCollectionConsumerRegistry(
                 .putAllPcollections(processBundleDescriptor.getPcollectionsMap())
                 .putAllWindowingStrategies(processBundleDescriptor.getWindowingStrategiesMap())
                 .build());
+    this.dataSampler = dataSampler;
   }
 
   /**
@@ -201,16 +220,16 @@ public FnDataReceiver<WindowedValue<?>> getMultiplexingConsumer(String pCollecti
           if (consumerAndMetadatas.size() == 1) {
             ConsumerAndMetadata consumerAndMetadata = consumerAndMetadatas.get(0);
             if (consumerAndMetadata.getConsumer() instanceof HandlesSplits) {
-              return new SplittingMetricTrackingFnDataReceiver(pcId, coder, consumerAndMetadata);
+              return new SplittingMetricTrackingFnDataReceiver(pcId, coder, consumerAndMetadata, dataSampler);
             }
-            return new MetricTrackingFnDataReceiver(pcId, coder, consumerAndMetadata);
+            return new MetricTrackingFnDataReceiver(pcId, coder, consumerAndMetadata, dataSampler);
           } else {
             /* TODO(SDF), Consider supporting splitting each consumer individually. This would never
             come up in the existing SDF expansion, but might be useful to support fused SDF nodes.
             This would require dedicated delivery of the split results to each of the consumers
             separately. */
             return new MultiplexingMetricTrackingFnDataReceiver(
-                pcId, coder, ImmutableList.copyOf(consumerAndMetadatas));
+                pcId, coder, ImmutableList.copyOf(consumerAndMetadatas), dataSampler);
           }
         });
   }
@@ -228,9 +247,11 @@ private class MetricTrackingFnDataReceiver<T> implements FnDataReceiver<Windowed
     private final BundleCounter elementCountCounter;
     private final SampleByteSizeDistribution<T> sampledByteSizeDistribution;
     private final Coder<T> coder;
+    private OutputSampler<T> outputSampler;
 
     public MetricTrackingFnDataReceiver(
-        String pCollectionId, Coder<T> coder, ConsumerAndMetadata consumerAndMetadata) {
+        String pCollectionId, Coder<T> coder, ConsumerAndMetadata consumerAndMetadata,
+        @Nullable DataSampler dataSampler) {
       this.delegate = consumerAndMetadata.getConsumer();
       this.executionState = consumerAndMetadata.getExecutionState();
 
@@ -266,6 +287,9 @@ public MetricTrackingFnDataReceiver(
       bundleProgressReporterRegistrar.register(sampledByteSizeUnderlyingDistribution);
 
       this.coder = coder;
+      if (dataSampler != null) {
+        this.outputSampler = dataSampler.sampleOutput(pCollectionId, coder);
+      }
     }
 
     @Override
@@ -276,6 +300,10 @@ public void accept(WindowedValue<T> input) throws Exception {
       // we have window optimization.
       this.sampledByteSizeDistribution.tryUpdate(input.getValue(), this.coder);
 
+      if (outputSampler != null) {
+        outputSampler.sample(input.getValue());
+      }
+
       // Use the ExecutionStateTracker and enter an appropriate state to track the
       // Process Bundle Execution time metric and also ensure user counters can get an appropriate
       // metrics container.
@@ -302,9 +330,11 @@ private class MultiplexingMetricTrackingFnDataReceiver<T>
     private final BundleCounter elementCountCounter;
     private final SampleByteSizeDistribution<T> sampledByteSizeDistribution;
     private final Coder<T> coder;
+    private @Nullable OutputSampler<T> outputSampler = null;
 
     public MultiplexingMetricTrackingFnDataReceiver(
-        String pCollectionId, Coder<T> coder, List<ConsumerAndMetadata> consumerAndMetadatas) {
+        String pCollectionId, Coder<T> coder, List<ConsumerAndMetadata> consumerAndMetadatas,
+        @Nullable DataSampler dataSampler) {
       this.consumerAndMetadatas = consumerAndMetadatas;
 
       HashMap<String, String> labels = new HashMap<>();
@@ -339,6 +369,9 @@ public MultiplexingMetricTrackingFnDataReceiver(
       bundleProgressReporterRegistrar.register(sampledByteSizeUnderlyingDistribution);
 
       this.coder = coder;
+      if (dataSampler != null) {
+        this.outputSampler = dataSampler.sampleOutput(pCollectionId, coder);
+      }
     }
 
     @Override
@@ -349,6 +382,10 @@ public void accept(WindowedValue<T> input) throws Exception {
       // when we have window optimization.
       this.sampledByteSizeDistribution.tryUpdate(input.getValue(), coder);
 
+      if (outputSampler != null) {
+        outputSampler.sample(input.getValue());
+      }
+
       // Use the ExecutionStateTracker and enter an appropriate state to track the
       // Process Bundle Execution time metric and also ensure user counters can get an appropriate
       // metrics container.
@@ -377,8 +414,9 @@ private class SplittingMetricTrackingFnDataReceiver<T> extends MetricTrackingFnD
     private final HandlesSplits delegate;
 
     public SplittingMetricTrackingFnDataReceiver(
-        String pCollection, Coder<T> coder, ConsumerAndMetadata consumerAndMetadata) {
-      super(pCollection, coder, consumerAndMetadata);
+        String pCollection, Coder<T> coder, ConsumerAndMetadata consumerAndMetadata,
+        @Nullable DataSampler dataSampler) {
+      super(pCollection, coder, consumerAndMetadata, dataSampler);
       this.delegate = (HandlesSplits) consumerAndMetadata.getConsumer();
     }