fix tests

Author: scwhittle

sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregator.java

@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.fn.data;
 
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -98,15 +99,31 @@ public void prepareForInstruction(
       String instructionId, StreamObserver<Elements> outboundObserver) {
     if (timeLimit > 0) {
       synchronized (flushLock) {
+        checkState(this.instructionId == null && this.outboundObserver == null);
         this.instructionId = instructionId;
         this.outboundObserver = outboundObserver;
       }
     } else {
+      checkState(this.instructionId == null && this.outboundObserver == null);
       this.instructionId = instructionId;
       this.outboundObserver = outboundObserver;
     }
   }
 
+  public void finishInstruction() {
+    if (timeLimit > 0) {
+      synchronized (flushLock) {
+        checkState(this.instructionId != null && this.outboundObserver != null);
+        this.instructionId = null;
+        this.outboundObserver = null;
+      }
+    } else {
+      checkState(this.instructionId != null && this.outboundObserver != null);
+      this.instructionId = null;
+      this.outboundObserver = null;
+    }
+  }
+
   /** Starts the flushing daemon thread if data_buffer_time_limit_ms is set. */
   public void start() {
     if (timeLimit > 0 && this.flushFuture == null) {
@@ -198,7 +215,6 @@ public Elements sendOrCollectBufferedDataAndFinishOutboundStreams() {
     } else {
       bufferedElements = convertBufferForTransmission();
     }
-    checkNotNull(instructionId);
     LOG.debug(
         "Closing streams for instruction {} and outbound data {} and timers {}.",
         instructionId,
@@ -224,13 +240,10 @@ public Elements sendOrCollectBufferedDataAndFinishOutboundStreams() {
       entry.getValue().resetStats();
     }
     // This is the end of the bundle so we reset state to prepare for future bundles.
-    instructionId = null;
     if (collectElementsIfNoFlushes && !hasFlushedForBundle) {
-      outboundObserver = null;
       return bufferedElements.build();
     }
     checkNotNull(outboundObserver).onNext(bufferedElements.build());
-    outboundObserver = null;
     hasFlushedForBundle = false;
     return null;
   }
@@ -247,7 +260,6 @@ public void discard() {
   }
 
   private Elements.Builder convertBufferForTransmission() {
-    checkNotNull(instructionId);
     Elements.Builder bufferedElements = Elements.newBuilder();
     for (Map.Entry<String, Receiver<?>> entry : outputDataReceivers.entrySet()) {
       if (!entry.getValue().hasBufferedOutput()) {

sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataOutboundAggregatorTest.java

@@ -20,6 +20,7 @@
 import static org.hamcrest.Matchers.empty;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.fail;
 
 import java.io.IOException;
@@ -323,6 +324,37 @@ public void testConfiguredBufferLimitMultipleEndpoints() throws Exception {
     checkEqualInAnyOrder(builder.build(), values.get(1));
   }
 
+  @Test
+  public void testInstructionLifecycle() {
+    BeamFnDataOutboundAggregator aggregator =
+        new BeamFnDataOutboundAggregator(PipelineOptionsFactory.create(), false);
+    assertThrows(
+        NullPointerException.class, () -> aggregator.sendElements(Elements.getDefaultInstance()));
+    aggregator.prepareForInstruction(
+        "testInstruction",
+        TestStreams.withOnNext(
+                (Consumer<Elements>)
+                    e -> {
+                      throw new RuntimeException("");
+                    })
+            .build());
+    assertThrows(
+        IllegalStateException.class,
+        () ->
+            aggregator.prepareForInstruction(
+                "testInstruction",
+                TestStreams.withOnNext(
+                        (Consumer<Elements>)
+                            e -> {
+                              throw new RuntimeException("");
+                            })
+                    .build()));
+    aggregator.finishInstruction();
+    assertThrows(
+        NullPointerException.class, () -> aggregator.sendElements(Elements.getDefaultInstance()));
+    assertThrows(IllegalStateException.class, aggregator::finishInstruction);
+  }
+
   private void checkEqualInAnyOrder(Elements first, Elements second) {
     MatcherAssert.assertThat(
         first.getDataList(), Matchers.containsInAnyOrder(second.getDataList().toArray()));

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

@@ -637,6 +637,10 @@ private void embedOutboundElementsIfApplicable(
       collectedElements.add(elements);
     }
     if (!hasFlushedAggregator) {
+      for (BeamFnDataOutboundAggregator aggregator :
+          bundleProcessor.getOutboundAggregators().values()) {
+        aggregator.finishInstruction();
+      }
       Elements.Builder elementsToEmbed = Elements.newBuilder();
       for (Elements collectedElement : collectedElements) {
         elementsToEmbed.mergeFrom(collectedElement);
@@ -653,6 +657,7 @@ private void embedOutboundElementsIfApplicable(
         if (elements != null) {
           aggregator.sendElements(elements);
         }
+        aggregator.finishInstruction();
       }
     }
   }