Let shuffle reader and writer update counters.

To this end, provide a reference to CounterSet.AddCounterMutator to ApplianceShuffleWriter and ApplianceShuffleReader.

----Release Notes----

[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=112828809

Author: mwegiel

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/ApplianceShuffleReader.java

@@ -16,6 +16,8 @@
 
 package com.google.cloud.dataflow.sdk.runners.worker;
 
+import com.google.cloud.dataflow.sdk.util.common.CounterSet;
+
 import java.io.IOException;
 
 import javax.annotation.concurrent.ThreadSafe;
@@ -37,12 +39,21 @@ public final class ApplianceShuffleReader implements ShuffleReader {
    */
   private long nativePointer;
 
+  /**
+   * Mutator that can be used to update counters.
+   */
+  private CounterSet.AddCounterMutator addCounterMutator;
+
   /**
    * @param shuffleReaderConfig opaque configuration for creating a
    * shuffle reader
+   * @param addCounterMutator mutator that can be used to update counters
    */
-  public ApplianceShuffleReader(byte[] shuffleReaderConfig) {
+  public ApplianceShuffleReader(
+      byte[] shuffleReaderConfig,
+      CounterSet.AddCounterMutator addCounterMutator) {
     this.nativePointer = createFromConfig(shuffleReaderConfig);
+    this.addCounterMutator = addCounterMutator;
   }
 
   @Override

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/ApplianceShuffleWriter.java

@@ -16,6 +16,8 @@
 
 package com.google.cloud.dataflow.sdk.runners.worker;
 
+import com.google.cloud.dataflow.sdk.util.common.CounterSet;
+
 import java.io.IOException;
 import javax.annotation.concurrent.ThreadSafe;
 
@@ -35,14 +37,23 @@ public final class ApplianceShuffleWriter implements ShuffleWriter {
    */
   private long nativePointer;
 
+  /**
+   * Mutator that can be used to update counters.
+   */
+  private final CounterSet.AddCounterMutator addCounterMutator;
+
   /**
    * @param shuffleWriterConfig opaque configuration for creating a
    * shuffle writer
    * @param bufferSize the writer buffer size
+   * @param addCounterMutator mutator that can be used to update counters
    */
-  public ApplianceShuffleWriter(byte[] shuffleWriterConfig,
-                                long bufferSize) {
+  public ApplianceShuffleWriter(
+      byte[] shuffleWriterConfig,
+      long bufferSize,
+      CounterSet.AddCounterMutator addCounterMutator) {
     this.nativePointer = createFromConfig(shuffleWriterConfig, bufferSize);
+    this.addCounterMutator = addCounterMutator;
   }
 
   @Override

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/GroupingShuffleReader.java

@@ -127,7 +127,8 @@ protected StateKind getStateSamplerStateKind() {
   @Override
   public GroupingShuffleReaderIterator<K, V> iterator() throws IOException {
     Preconditions.checkArgument(shuffleReaderConfig != null);
-    ApplianceShuffleReader asr = new ApplianceShuffleReader(shuffleReaderConfig);
+    ApplianceShuffleReader asr =
+        new ApplianceShuffleReader(shuffleReaderConfig, addCounterMutator);
     String datasetId = asr.getDatasetId();
     initCounter(datasetId);
 

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/PartitioningShuffleReader.java

@@ -22,6 +22,7 @@
 import com.google.cloud.dataflow.sdk.util.CoderUtils;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
 import com.google.cloud.dataflow.sdk.util.WindowedValue.WindowedValueCoder;
+import com.google.cloud.dataflow.sdk.util.common.CounterSet;
 import com.google.cloud.dataflow.sdk.util.common.worker.BatchingShuffleEntryReader;
 import com.google.cloud.dataflow.sdk.util.common.worker.NativeReader;
 import com.google.cloud.dataflow.sdk.util.common.worker.ShuffleEntry;
@@ -43,15 +44,18 @@ public class PartitioningShuffleReader<K, V> extends NativeReader<WindowedValue<
   final byte[] shuffleReaderConfig;
   final String startShufflePosition;
   final String stopShufflePosition;
+  final CounterSet.AddCounterMutator addCounterMutator;
   Coder<K> keyCoder;
   WindowedValueCoder<V> windowedValueCoder;
 
   public PartitioningShuffleReader(PipelineOptions options, byte[] shuffleReaderConfig,
-      String startShufflePosition, String stopShufflePosition, Coder<WindowedValue<KV<K, V>>> coder)
+      String startShufflePosition, String stopShufflePosition, Coder<WindowedValue<KV<K, V>>> coder,
+      CounterSet.AddCounterMutator addCounterMutator)
       throws Exception {
     this.shuffleReaderConfig = shuffleReaderConfig;
     this.startShufflePosition = startShufflePosition;
     this.stopShufflePosition = stopShufflePosition;
+    this.addCounterMutator = addCounterMutator;
     initCoder(coder);
   }
 
@@ -79,7 +83,10 @@ private void initCoder(Coder<WindowedValue<KV<K, V>>> coder) throws Exception {
   public NativeReaderIterator<WindowedValue<KV<K, V>>> iterator() throws IOException {
     Preconditions.checkArgument(shuffleReaderConfig != null);
     return iterator(new BatchingShuffleEntryReader(
-        new ChunkingShuffleBatchReader(new ApplianceShuffleReader(shuffleReaderConfig))));
+        new ChunkingShuffleBatchReader(
+            new ApplianceShuffleReader(
+                shuffleReaderConfig,
+                addCounterMutator))));
   }
 
   PartitioningShuffleReaderIterator iterator(ShuffleEntryReader reader) {

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/PartitioningShuffleReaderFactory.java

@@ -50,19 +50,21 @@ public NativeReader<?> create(
     @SuppressWarnings({"unchecked", "rawtypes"})
     Coder<WindowedValue<KV<Object, Object>>> typedCoder =
         (Coder<WindowedValue<KV<Object, Object>>>) coder;
-    return createTyped(spec, options, typedCoder);
+    return createTyped(spec, options, typedCoder, addCounterMutator);
   }
 
   public <K, V> PartitioningShuffleReader<K, V> createTyped(
       CloudObject spec,
       PipelineOptions options,
-      Coder<WindowedValue<KV<K, V>>> coder)
+      Coder<WindowedValue<KV<K, V>>> coder,
+      CounterSet.AddCounterMutator addCounterMutator)
       throws Exception {
     return new PartitioningShuffleReader<K, V>(
         options,
         decodeBase64(getString(spec, PropertyNames.SHUFFLE_READER_CONFIG)),
         getString(spec, PropertyNames.START_SHUFFLE_POSITION, null),
         getString(spec, PropertyNames.END_SHUFFLE_POSITION, null),
-        coder);
+        coder,
+        addCounterMutator);
   }
 }

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/ShuffleSink.java

@@ -283,7 +283,7 @@ public void close() throws IOException {
   public SinkWriter<WindowedValue<T>> writer() throws IOException {
     Preconditions.checkArgument(shuffleWriterConfig != null);
     ApplianceShuffleWriter applianceWriter = new ApplianceShuffleWriter(
-        shuffleWriterConfig, SHUFFLE_WRITER_BUFFER_SIZE);
+        shuffleWriterConfig, SHUFFLE_WRITER_BUFFER_SIZE, addCounterMutator);
     String datasetId = applianceWriter.getDatasetId();
     return writer(new ChunkingShuffleEntryWriter(applianceWriter), datasetId);
   }

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/UngroupedShuffleReader.java

@@ -19,6 +19,7 @@
 import com.google.cloud.dataflow.sdk.coders.Coder;
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
 import com.google.cloud.dataflow.sdk.util.CoderUtils;
+import com.google.cloud.dataflow.sdk.util.common.CounterSet;
 import com.google.cloud.dataflow.sdk.util.common.worker.BatchingShuffleEntryReader;
 import com.google.cloud.dataflow.sdk.util.common.worker.NativeReader;
 import com.google.cloud.dataflow.sdk.util.common.worker.ShuffleEntry;
@@ -41,21 +42,27 @@ public class UngroupedShuffleReader<T> extends NativeReader<T> {
   final String startShufflePosition;
   final String stopShufflePosition;
   final Coder<T> coder;
+  final CounterSet.AddCounterMutator addCounterMutator;
 
   public UngroupedShuffleReader(
       @SuppressWarnings("unused") PipelineOptions options, byte[] shuffleReaderConfig,
-      @Nullable String startShufflePosition, @Nullable String stopShufflePosition, Coder<T> coder) {
+      @Nullable String startShufflePosition, @Nullable String stopShufflePosition, Coder<T> coder,
+      @Nullable CounterSet.AddCounterMutator addCounterMutator) {
     this.shuffleReaderConfig = shuffleReaderConfig;
     this.startShufflePosition = startShufflePosition;
     this.stopShufflePosition = stopShufflePosition;
     this.coder = coder;
+    this.addCounterMutator = addCounterMutator;
   }
 
   @Override
   public NativeReaderIterator<T> iterator() throws IOException {
     Preconditions.checkArgument(shuffleReaderConfig != null);
     return iterator(new BatchingShuffleEntryReader(
-        new ChunkingShuffleBatchReader(new ApplianceShuffleReader(shuffleReaderConfig))));
+        new ChunkingShuffleBatchReader(
+            new ApplianceShuffleReader(
+                shuffleReaderConfig,
+                addCounterMutator))));
   }
 
   UngroupedShuffleReaderIterator iterator(ShuffleEntryReader reader) {

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/UngroupedShuffleReaderFactory.java

@@ -43,19 +43,27 @@ public NativeReader<?> create(
       @Nullable CounterSet.AddCounterMutator addCounterMutator,
       @Nullable String operationName)
       throws Exception {
-    return create(spec, coder, options);
+    return create(spec, coder, options, addCounterMutator);
   }
 
   public <T> UngroupedShuffleReader<T> create(
-      CloudObject spec, Coder<T> coder, PipelineOptions options) throws Exception {
-    return create(options, spec, coder);
+      CloudObject spec,
+      Coder<T> coder,
+      PipelineOptions options,
+      CounterSet.AddCounterMutator addCounterMutator) throws Exception {
+    return create(options, spec, coder, addCounterMutator);
   }
 
   <T> UngroupedShuffleReader<T> create(
-      PipelineOptions options, CloudObject spec, Coder<T> coder) throws Exception {
+      PipelineOptions options,
+      CloudObject spec,
+      Coder<T> coder,
+      CounterSet.AddCounterMutator addCounterMutator) throws Exception {
     return new UngroupedShuffleReader<>(options,
         decodeBase64(getString(spec, PropertyNames.SHUFFLE_READER_CONFIG)),
         getString(spec, PropertyNames.START_SHUFFLE_POSITION, null),
-        getString(spec, PropertyNames.END_SHUFFLE_POSITION, null), coder);
+        getString(spec, PropertyNames.END_SHUFFLE_POSITION, null),
+        coder,
+        addCounterMutator);
   }
 }

sdk/src/test/java/com/google/cloud/dataflow/sdk/runners/worker/PartitioningShuffleReaderTest.java

@@ -94,7 +94,8 @@ private void runTestReadFromShuffle(List<WindowedValue<KV<Integer, String>>> exp
     // Read from shuffle with PartitioningShuffleReader.
     PartitioningShuffleReader<Integer, String> partitioningShuffleReader =
         new PartitioningShuffleReader<>(
-            PipelineOptionsFactory.create(), null, null, null, elemCoder);
+            PipelineOptionsFactory.create(), null, null, null, elemCoder,
+            addCounterMutator);
     ExecutorTestUtils.TestReaderObserver observer =
         new ExecutorTestUtils.TestReaderObserver(partitioningShuffleReader);
 

sdk/src/test/java/com/google/cloud/dataflow/sdk/runners/worker/UngroupedShuffleReaderTest.java

@@ -79,7 +79,7 @@ void runTestReadFromShuffle(List<Integer> expected) throws Exception {
         new UngroupedShuffleReader<>(
             PipelineOptionsFactory.create(),
             null, null, null,
-            elemCoder);
+            elemCoder, addCounterMutator);
     ExecutorTestUtils.TestReaderObserver observer =
         new ExecutorTestUtils.TestReaderObserver(ungroupedShuffleReader);