Fix FanOutStreamingEngineWorkerHarnessTest - testOnNewWorkerMetadata_redistributesBudget (#34531)

* Fix Resources weren't released in time issue in testOnNewWorkerMetadata_redistributesBudget

* Import Assert.fail separately

* Refactor FanOutStreamingEngineWorkerHarness related components and tests to avoid cleanUp hanging indefinitely

* Run spotlessJavaApply

Author: akashorabek

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

@@ -298,7 +298,7 @@ private synchronized void consumeWindmillWorkerEndpoints(WindmillEndpoints newWi
         newWindmillEndpoints,
         activeMetadataVersion,
         newWindmillEndpoints.version());
-    closeStreamsNotIn(newWindmillEndpoints);
+    closeStreamsNotIn(newWindmillEndpoints).join();
     ImmutableMap<Endpoint, WindmillStreamSender> newStreams =
         createAndStartNewStreams(newWindmillEndpoints.windmillEndpoints()).join();
     StreamingEngineBackends newBackends =

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSender.java

@@ -22,6 +22,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
@@ -40,6 +41,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudgetSpender;
 import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.FixedStreamHeartbeatSender;
 import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
 
 /**
@@ -58,6 +60,7 @@
 @ThreadSafe
 final class WindmillStreamSender implements GetWorkBudgetSpender, StreamSender {
   private static final String STREAM_STARTER_THREAD_NAME = "StartWindmillStreamThread-%d";
+  private static final int TERMINATION_TIMEOUT_SECONDS = 5;
   private final AtomicBoolean started;
   private final AtomicReference<GetWorkBudget> getWorkBudget;
   private final GetWorkStream getWorkStream;
@@ -142,11 +145,26 @@ synchronized void start() {
 
   @Override
   public synchronized void close() {
-    streamStarter.shutdownNow();
+    streamStarter.shutdown();
     getWorkStream.shutdown();
     getDataStream.shutdown();
     workCommitter.stop();
     commitWorkStream.shutdown();
+    try {
+      if (!Preconditions.checkNotNull(streamStarter)
+          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+        streamStarter.shutdownNow();
+      }
+      Preconditions.checkNotNull(getWorkStream)
+          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+      Preconditions.checkNotNull(getDataStream)
+          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+      Preconditions.checkNotNull(commitWorkStream)
+          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new RuntimeException("Interrupted while waiting for streams to terminate", e);
+    }
   }
 
   @Override

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java

@@ -80,6 +80,8 @@
 @RunWith(JUnit4.class)
 public class FanOutStreamingEngineWorkerHarnessTest {
   private static final String CHANNEL_NAME = "FanOutStreamingEngineWorkerHarnessTest";
+  private static final long WAIT_FOR_METADATA_INJECTIONS_SECONDS = 5;
+  private static final long SERVER_SHUTDOWN_TIMEOUT_SECONDS = 30;
   private static final WindmillServiceAddress DEFAULT_WINDMILL_SERVICE_ADDRESS =
       WindmillServiceAddress.create(HostAndPort.fromParts(WindmillChannels.LOCALHOST, 443));
   private static final ImmutableMap<String, WorkerMetadataResponse.Endpoint> DEFAULT =
@@ -101,7 +103,7 @@ public class FanOutStreamingEngineWorkerHarnessTest {
           .build();
 
   @Rule
-  public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule().setTimeout(1, TimeUnit.MINUTES);
+  public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule().setTimeout(3, TimeUnit.MINUTES);
 
   private final GrpcWindmillStreamFactory streamFactory =
       spy(GrpcWindmillStreamFactory.of(JOB_HEADER).build());
@@ -167,10 +169,15 @@ public void setUp() throws IOException {
   }
 
   @After
-  public void cleanUp() {
+  public void cleanUp() throws InterruptedException {
     Preconditions.checkNotNull(fanOutStreamingEngineWorkProvider).shutdown();
     stubFactory.shutdown();
     fakeStreamingEngineServer.shutdown();
+    if (!Preconditions.checkNotNull(
+        fakeStreamingEngineServer.awaitTermination(
+            SERVER_SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS))) {
+      fakeStreamingEngineServer.shutdownNow();
+    }
   }
 
   private FanOutStreamingEngineWorkerHarness newFanOutStreamingEngineWorkerHarness(
@@ -329,8 +336,10 @@ public void testOnNewWorkerMetadata_redistributesBudget() throws InterruptedExce
 
     fakeGetWorkerMetadataStub.injectWorkerMetadata(firstWorkerMetadata);
     verify(getWorkBudgetDistributor, times(1)).distributeBudget(any(), any());
+    TimeUnit.SECONDS.sleep(WAIT_FOR_METADATA_INJECTIONS_SECONDS);
     fakeGetWorkerMetadataStub.injectWorkerMetadata(secondWorkerMetadata);
     verify(getWorkBudgetDistributor, times(2)).distributeBudget(any(), any());
+    TimeUnit.SECONDS.sleep(WAIT_FOR_METADATA_INJECTIONS_SECONDS);
   }
 
   private static class WindmillServiceFakeStub