Merge pull request #160 from dhalperi/backport-52

Backport apache/beam#52

Author: davorbonaci

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/BoundedReadEvaluatorFactory.java

@@ -15,6 +15,8 @@
  */
 package com.google.cloud.dataflow.sdk.runners.inprocess;
 
+import com.google.cloud.dataflow.sdk.io.BoundedSource;
+import com.google.cloud.dataflow.sdk.io.BoundedSource.BoundedReader;
 import com.google.cloud.dataflow.sdk.io.Read.Bounded;
 import com.google.cloud.dataflow.sdk.io.Source.Reader;
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.CommittedBundle;
@@ -78,8 +80,7 @@ private <OutputT> TransformEvaluator<?> getTransformEvaluator(
   @SuppressWarnings("unchecked")
   private <OutputT> Queue<BoundedReadEvaluator<OutputT>> getTransformEvaluatorQueue(
       final AppliedPTransform<?, PCollection<OutputT>, Bounded<OutputT>> transform,
-      final InProcessEvaluationContext evaluationContext)
-      throws IOException {
+      final InProcessEvaluationContext evaluationContext) {
     // Key by the application and the context the evaluation is occurring in (which call to
     // Pipeline#run).
     EvaluatorKey key = new EvaluatorKey(transform, evaluationContext);
@@ -101,39 +102,51 @@ private <OutputT> Queue<BoundedReadEvaluator<OutputT>> getTransformEvaluatorQueu
     return evaluatorQueue;
   }
 
+  /**
+   * A {@link BoundedReadEvaluator} produces elements from an underlying {@link BoundedSource},
+   * discarding all input elements. Within the call to {@link #finishBundle()}, the evaluator
+   * creates the {@link BoundedReader} and consumes all available input.
+   *
+   * <p>A {@link BoundedReadEvaluator} should only be created once per {@link BoundedSource}, and
+   * each evaluator should only be called once per evaluation of the pipeline. Otherwise, the source
+   * may produce duplicate elements.
+   */
   private static class BoundedReadEvaluator<OutputT> implements TransformEvaluator<Object> {
     private final AppliedPTransform<?, PCollection<OutputT>, Bounded<OutputT>> transform;
     private final InProcessEvaluationContext evaluationContext;
-    private final Reader<OutputT> reader;
     private boolean contentsRemaining;
 
     public BoundedReadEvaluator(
         AppliedPTransform<?, PCollection<OutputT>, Bounded<OutputT>> transform,
-        InProcessEvaluationContext evaluationContext)
-        throws IOException {
+        InProcessEvaluationContext evaluationContext) {
       this.transform = transform;
       this.evaluationContext = evaluationContext;
-      reader =
-          transform.getTransform().getSource().createReader(evaluationContext.getPipelineOptions());
-      contentsRemaining = reader.start();
     }
 
     @Override
     public void processElement(WindowedValue<Object> element) {}
 
     @Override
     public InProcessTransformResult finishBundle() throws IOException {
-      UncommittedBundle<OutputT> output = evaluationContext.createRootBundle(transform.getOutput());
-      while (contentsRemaining) {
-        output.add(
-            WindowedValue.timestampedValueInGlobalWindow(
-                reader.getCurrent(), reader.getCurrentTimestamp()));
-        contentsRemaining = reader.advance();
+      try (final Reader<OutputT> reader =
+              transform
+                  .getTransform()
+                  .getSource()
+                  .createReader(evaluationContext.getPipelineOptions());) {
+        contentsRemaining = reader.start();
+        UncommittedBundle<OutputT> output =
+            evaluationContext.createRootBundle(transform.getOutput());
+        while (contentsRemaining) {
+          output.add(
+              WindowedValue.timestampedValueInGlobalWindow(
+                  reader.getCurrent(), reader.getCurrentTimestamp()));
+          contentsRemaining = reader.advance();
+        }
+        reader.close();
+        return StepTransformResult.withHold(transform, BoundedWindow.TIMESTAMP_MAX_VALUE)
+            .addOutput(output)
+            .build();
       }
-      return StepTransformResult
-          .withHold(transform, BoundedWindow.TIMESTAMP_MAX_VALUE)
-          .addOutput(output)
-          .build();
     }
   }
 }

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/UnboundedReadEvaluatorFactory.java

@@ -99,6 +99,16 @@ private <OutputT> Queue<UnboundedReadEvaluator<OutputT>> getTransformEvaluatorQu
     return evaluatorQueue;
   }
 
+  /**
+   * A {@link UnboundedReadEvaluator} produces elements from an underlying {@link UnboundedSource},
+   * discarding all input elements. Within the call to {@link #finishBundle()}, the evaluator
+   * creates the {@link UnboundedReader} and consumes some currently available input.
+   *
+   * <p>Calls to {@link UnboundedReadEvaluator} are not internally thread-safe, and should only be
+   * used by a single thread at a time. Each {@link UnboundedReadEvaluator} maintains its own
+   * checkpoint, and constructs its reader from the current checkpoint in each call to
+   * {@link #finishBundle()}.
+   */
   private static class UnboundedReadEvaluator<OutputT> implements TransformEvaluator<Object> {
     private static final int ARBITRARY_MAX_ELEMENTS = 10;
     private final AppliedPTransform<?, PCollection<OutputT>, Unbounded<OutputT>> transform;
@@ -122,28 +132,29 @@ public void processElement(WindowedValue<Object> element) {}
     @Override
     public InProcessTransformResult finishBundle() throws IOException {
       UncommittedBundle<OutputT> output = evaluationContext.createRootBundle(transform.getOutput());
-      UnboundedReader<OutputT> reader =
-          createReader(
-              transform.getTransform().getSource(), evaluationContext.getPipelineOptions());
-      int numElements = 0;
-      if (reader.start()) {
-        do {
-          output.add(
-              WindowedValue.timestampedValueInGlobalWindow(
-                  reader.getCurrent(), reader.getCurrentTimestamp()));
-          numElements++;
-        } while (numElements < ARBITRARY_MAX_ELEMENTS && reader.advance());
+      try (UnboundedReader<OutputT> reader =
+              createReader(
+                  transform.getTransform().getSource(), evaluationContext.getPipelineOptions());) {
+        int numElements = 0;
+        if (reader.start()) {
+          do {
+            output.add(
+                WindowedValue.timestampedValueInGlobalWindow(
+                    reader.getCurrent(), reader.getCurrentTimestamp()));
+            numElements++;
+          } while (numElements < ARBITRARY_MAX_ELEMENTS && reader.advance());
+        }
+        checkpointMark = reader.getCheckpointMark();
+        checkpointMark.finalizeCheckpoint();
+        // TODO: When exercising create initial splits, make this the minimum watermark across all
+        // existing readers
+        StepTransformResult result =
+            StepTransformResult.withHold(transform, reader.getWatermark())
+                .addOutput(output)
+                .build();
+        evaluatorQueue.offer(this);
+        return result;
       }
-      checkpointMark = reader.getCheckpointMark();
-      checkpointMark.finalizeCheckpoint();
-      // TODO: When exercising create initial splits, make this the minimum watermark across all
-      // existing readers
-      StepTransformResult result =
-          StepTransformResult.withHold(transform, reader.getWatermark())
-              .addOutput(output)
-              .build();
-      evaluatorQueue.offer(this);
-      return result;
     }
 
     private <CheckpointMarkT extends CheckpointMark> UnboundedReader<OutputT> createReader(

sdk/src/test/java/com/google/cloud/dataflow/sdk/runners/inprocess/BoundedReadEvaluatorFactoryTest.java

@@ -18,24 +18,39 @@
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.emptyIterable;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.cloud.dataflow.sdk.coders.BigEndianLongCoder;
+import com.google.cloud.dataflow.sdk.coders.Coder;
 import com.google.cloud.dataflow.sdk.io.BoundedSource;
+import com.google.cloud.dataflow.sdk.io.BoundedSource.BoundedReader;
 import com.google.cloud.dataflow.sdk.io.CountingSource;
 import com.google.cloud.dataflow.sdk.io.Read;
 import com.google.cloud.dataflow.sdk.io.Read.Bounded;
+import com.google.cloud.dataflow.sdk.options.PipelineOptions;
+import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.CommittedBundle;
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.UncommittedBundle;
 import com.google.cloud.dataflow.sdk.testing.TestPipeline;
+import com.google.cloud.dataflow.sdk.transforms.AppliedPTransform;
 import com.google.cloud.dataflow.sdk.transforms.windowing.BoundedWindow;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
 import com.google.cloud.dataflow.sdk.values.PCollection;
+import com.google.common.collect.ImmutableList;
 
+import org.joda.time.Instant;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.NoSuchElementException;
 
 /**
  * Tests for {@link BoundedReadEvaluatorFactory}.
@@ -45,7 +60,7 @@ public class BoundedReadEvaluatorFactoryTest {
   private BoundedSource<Long> source;
   private PCollection<Long> longs;
   private TransformEvaluatorFactory factory;
-  private InProcessEvaluationContext context;
+  @Mock private InProcessEvaluationContext context;
 
   @Before
   public void setup() {
@@ -146,6 +161,125 @@ public void boundedSourceEvaluatorSimultaneousEvaluations() throws Exception {
             gw(1L), gw(2L), gw(4L), gw(8L), gw(9L), gw(7L), gw(6L), gw(5L), gw(3L), gw(0L)));
   }
 
+  @Test
+  public void boundedSourceEvaluatorClosesReader() throws Exception {
+    TestSource<Long> source = new TestSource<>(BigEndianLongCoder.of(), 1L, 2L, 3L);
+
+    TestPipeline p = TestPipeline.create();
+    PCollection<Long> pcollection = p.apply(Read.from(source));
+    AppliedPTransform<?, ?, ?> sourceTransform = pcollection.getProducingTransformInternal();
+
+    UncommittedBundle<Long> output = InProcessBundle.unkeyed(longs);
+    when(context.createRootBundle(pcollection)).thenReturn(output);
+
+    TransformEvaluator<?> evaluator = factory.forApplication(sourceTransform, null, context);
+    evaluator.finishBundle();
+    CommittedBundle<Long> committed = output.commit(Instant.now());
+    assertThat(committed.getElements(), containsInAnyOrder(gw(2L), gw(3L), gw(1L)));
+    assertThat(TestSource.readerClosed, is(true));
+  }
+
+  @Test
+  public void boundedSourceEvaluatorNoElementsClosesReader() throws Exception {
+    TestSource<Long> source = new TestSource<>(BigEndianLongCoder.of());
+
+    TestPipeline p = TestPipeline.create();
+    PCollection<Long> pcollection = p.apply(Read.from(source));
+    AppliedPTransform<?, ?, ?> sourceTransform = pcollection.getProducingTransformInternal();
+
+    UncommittedBundle<Long> output = InProcessBundle.unkeyed(longs);
+    when(context.createRootBundle(pcollection)).thenReturn(output);
+
+    TransformEvaluator<?> evaluator = factory.forApplication(sourceTransform, null, context);
+    evaluator.finishBundle();
+    CommittedBundle<Long> committed = output.commit(Instant.now());
+    assertThat(committed.getElements(), emptyIterable());
+    assertThat(TestSource.readerClosed, is(true));
+  }
+
+  private static class TestSource<T> extends BoundedSource<T> {
+    private static boolean readerClosed;
+    private final Coder<T> coder;
+    private final T[] elems;
+
+    public TestSource(Coder<T> coder, T... elems) {
+      this.elems = elems;
+      this.coder = coder;
+      readerClosed = false;
+    }
+
+    @Override
+    public List<? extends BoundedSource<T>> splitIntoBundles(
+        long desiredBundleSizeBytes, PipelineOptions options) throws Exception {
+      return ImmutableList.of(this);
+    }
+
+    @Override
+    public long getEstimatedSizeBytes(PipelineOptions options) throws Exception {
+      return 0;
+    }
+
+    @Override
+    public boolean producesSortedKeys(PipelineOptions options) throws Exception {
+      return false;
+    }
+
+    @Override
+    public BoundedSource.BoundedReader<T> createReader(PipelineOptions options) throws IOException {
+      return new TestReader<>(this, elems);
+    }
+
+    @Override
+    public void validate() {
+    }
+
+    @Override
+    public Coder<T> getDefaultOutputCoder() {
+      return coder;
+    }
+  }
+
+  private static class TestReader<T> extends BoundedReader<T> {
+    private final BoundedSource<T> source;
+    private final List<T> elems;
+    private int index;
+
+    public TestReader(BoundedSource<T> source, T... elems) {
+      this.source = source;
+      this.elems = Arrays.asList(elems);
+      this.index = -1;
+    }
+
+    @Override
+    public BoundedSource<T> getCurrentSource() {
+      return source;
+    }
+
+    @Override
+    public boolean start() throws IOException {
+      return advance();
+    }
+
+    @Override
+    public boolean advance() throws IOException {
+      if (elems.size() > index + 1) {
+        index++;
+        return true;
+      }
+      return false;
+    }
+
+    @Override
+    public T getCurrent() throws NoSuchElementException {
+      return elems.get(index);
+    }
+
+    @Override
+    public void close() throws IOException {
+      TestSource.readerClosed = true;
+    }
+  }
+
   private static WindowedValue<Long> gw(Long elem) {
     return WindowedValue.valueInGlobalWindow(elem);
   }