Fix many flaky tests

examples/java/src/test/java/org/apache/beam/examples/complete/kafkatopubsub/KafkaToPubsubE2ETest.java

@@ -62,11 +62,15 @@ public class KafkaToPubsubE2ETest {
 
   @ClassRule
   public static final PubSubEmulatorContainer PUB_SUB_EMULATOR_CONTAINER =
-      new PubSubEmulatorContainer(DockerImageName.parse(PUBSUB_EMULATOR_IMAGE));
+      new PubSubEmulatorContainer(DockerImageName.parse(PUBSUB_EMULATOR_IMAGE))
+          .withStartupAttempts(3)
+          .withStartupTimeout(java.time.Duration.ofMinutes(5));
 
   @ClassRule
   public static final KafkaContainer KAFKA_CONTAINER =
-      new KafkaContainer(DockerImageName.parse(KAFKA_IMAGE_NAME));
+      new KafkaContainer(DockerImageName.parse(KAFKA_IMAGE_NAME))
+          .withStartupAttempts(3)
+          .withStartupTimeout(java.time.Duration.ofMinutes(5));
 
   @Rule public final transient TestPipeline pipeline = TestPipeline.fromOptions(OPTIONS);
   @Rule public final transient TestPubsub testPubsub = TestPubsub.fromOptions(OPTIONS);

runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapperTest.java

@@ -635,6 +635,8 @@ private static void testSourceDoesNotShutdown(boolean shouldHaveReaders) throws
       ProcessingTimeService timerService = Mockito.mock(ProcessingTimeService.class);
       Mockito.when(timerService.getCurrentProcessingTime()).thenReturn(Long.MAX_VALUE);
       Mockito.when(mock.getProcessingTimeService()).thenReturn(timerService);
+      Mockito.when(mock.getMetricGroup())
+          .thenReturn(UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup());
 
       sourceWrapper.setRuntimeContext(mock);
       sourceWrapper.open(new Configuration());

runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismRunner.java

@@ -18,10 +18,10 @@
 package org.apache.beam.runners.prism;
 
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import java.io.IOException;
 import java.net.ServerSocket;
+import java.net.Socket;
 import java.util.Arrays;
 import org.apache.beam.runners.portability.PortableRunner;
 import org.apache.beam.sdk.Pipeline;
@@ -46,6 +46,7 @@
 public class PrismRunner extends PipelineRunner<PipelineResult> {
 
   private static final Logger LOG = LoggerFactory.getLogger(PrismRunner.class);
+  private static final int MAX_START_ATTEMPTS = 5;
 
   private final PrismPipelineOptions prismPipelineOptions;
 
@@ -82,19 +83,37 @@ public PipelineResult run(Pipeline pipeline) {
         prismPipelineOptions.getDefaultEnvironmentType(),
         prismPipelineOptions.getJobEndpoint());
 
-    try {
-      PrismExecutor executor = startPrism();
-      PortableRunner delegate = PortableRunner.fromOptions(prismPipelineOptions);
-      return new PrismPipelineResult(delegate.run(pipeline), executor::stop);
-    } catch (IOException e) {
-      throw new RuntimeException(e);
+    for (int attempt = 1; ; attempt++) {
+      PrismExecutor executor;
+      try {
+        executor = startPrism();
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+      try {
+        PortableRunner delegate = PortableRunner.fromOptions(prismPipelineOptions);
+        return new PrismPipelineResult(delegate.run(pipeline), executor::stop);
+      } catch (RuntimeException e) {
+        boolean alive = executor.isAlive();
+        executor.stop();
+        if (attempt < MAX_START_ATTEMPTS
+            && (!alive
+                || (e.getMessage() != null && e.getMessage().contains("JobService/Prepare")))) {
+          LOG.warn(
+              "Prism job submission failed on attempt {} (alive={}), retrying on a new port...",
+              attempt,
+              alive,
+              e);
+          continue;
+        }
+        throw e;
+      }
     }
   }
 
   PrismExecutor startPrism() throws IOException {
     PrismLocator locator = new PrismLocator(prismPipelineOptions);
-    int port = findAvailablePort();
-    String portFlag = String.format(PrismExecutor.JOB_PORT_FLAG_TEMPLATE, port);
+    String command = locator.resolve();
     String serveHttpFlag =
         String.format(
             PrismExecutor.SERVE_HTTP_FLAG_TEMPLATE, prismPipelineOptions.getEnableWebUI());
@@ -104,18 +123,52 @@ PrismExecutor startPrism() throws IOException {
     String logLevelFlag =
         String.format(
             PrismExecutor.LOG_LEVEL_FLAG_TEMPLATE, prismPipelineOptions.getPrismLogLevel());
-    String endpoint = "localhost:" + port;
-    prismPipelineOptions.setJobEndpoint(endpoint);
-    String command = locator.resolve();
-    PrismExecutor executor =
-        PrismExecutor.builder()
-            .setCommand(command)
-            .setArguments(
-                Arrays.asList(portFlag, serveHttpFlag, idleShutdownTimeoutFlag, logLevelFlag))
-            .build();
-    executor.execute();
-    checkState(executor.isAlive());
-    return executor;
+
+    for (int attempt = 1; ; attempt++) {
+      int port = findAvailablePort();
+      String portFlag = String.format(PrismExecutor.JOB_PORT_FLAG_TEMPLATE, port);
+      String endpoint = "localhost:" + port;
+      prismPipelineOptions.setJobEndpoint(endpoint);
+      PrismExecutor executor =
+          PrismExecutor.builder()
+              .setCommand(command)
+              .setArguments(
+                  Arrays.asList(portFlag, serveHttpFlag, idleShutdownTimeoutFlag, logLevelFlag))
+              .build();
+      executor.execute();
+      if (waitForPort(executor, port)) {
+        return executor;
+      }
+      executor.stop();
+      String message =
+          String.format(
+              "Prism failed to start listening on port %d (attempt %d/%d)",
+              port, attempt, MAX_START_ATTEMPTS);
+      LOG.warn("{}", message);
+      if (attempt >= MAX_START_ATTEMPTS) {
+        throw new IOException(message);
+      }
+    }
+  }
+
+  private static boolean waitForPort(PrismExecutor executor, int port) {
+    long deadlineMillis = System.currentTimeMillis() + 10_000L;
+    while (System.currentTimeMillis() < deadlineMillis) {
+      if (!executor.isAlive()) {
+        return false;
+      }
+      try (Socket socket = new Socket("localhost", port)) {
+        return executor.isAlive();
+      } catch (IOException ignored) {
+        try {
+          Thread.sleep(10L);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          return false;
+        }
+      }
+    }
+    return executor.isAlive();
   }
 
   private static void assignDefaultsIfNeeded(PrismPipelineOptions prismPipelineOptions) {

runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java

@@ -177,8 +177,13 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) {
                 ctxRef.set(ctx);
                 if (!cancelRequested.get()) {
                   ctx.evaluate();
+                } else {
+                  cancelSparkJobs.run();
                 }
               } finally {
+                if (cancelRequested.get()) {
+                  cancelSparkJobs.run();
+                }
                 if (releaseSession) {
                   SparkSessionFactory.release(sparkSession);
                 }

runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java

@@ -26,7 +26,6 @@
 import com.esotericsoftware.kryo.serializers.JavaSerializer;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -125,14 +124,16 @@ public static synchronized SparkSession acquire(SparkStructuredStreamingPipeline
     if (options.getUseActiveSparkSession()) {
       return SparkSession.active();
     }
-    boolean noUsableSession =
-        !isUsable(SparkSession.getActiveSession()) && !isUsable(SparkSession.getDefaultSession());
+    SparkSession existingActive =
+        isUsable(SparkSession.getActiveSession()) ? SparkSession.getActiveSession().get() : null;
+    SparkSession existingDefault =
+        isUsable(SparkSession.getDefaultSession()) ? SparkSession.getDefaultSession().get() : null;
     SparkSession session = sessionBuilder(options.getSparkMaster(), options).getOrCreate();
     Integer count = OWNED_SESSIONS.get(session);
     if (count != null) {
       OWNED_SESSIONS.put(session, count + 1);
       LOG.info("Pipeline options will not be applied to the shared SparkSession");
-    } else if (noUsableSession) {
+    } else if (session != existingActive && session != existingDefault) {
       OWNED_SESSIONS.put(session, 1);
     }
     return session;
@@ -154,6 +155,8 @@ public static synchronized void release(SparkSession session) {
     OWNED_SESSIONS.remove(session);
     LOG.info("Stopping SparkSession created by the runner");
     session.stop();
+    SparkSession.clearActiveSession();
+    SparkSession.clearDefaultSession();
   }
 
   private static boolean isUsable(Option<SparkSession> session) {
@@ -176,19 +179,22 @@ private static SparkSession.Builder sessionBuilder(
         sparkConf.setAppName(options.getAppName());
       }
 
-      if (options.getFilesToStage() != null && !options.getFilesToStage().isEmpty()) {
-        // Append the files to stage provided by the user to `spark.jars`.
-        PipelineResources.prepareFilesForStaging(options);
-        String[] filesToStage = filterFilesToStage(options, Collections.emptyList());
-        String[] jars = getSparkJars(sparkConf);
-        sparkConf.setJars(jars.length > 0 ? ArrayUtils.addAll(jars, filesToStage) : filesToStage);
-      } else if (!sparkConf.contains("spark.jars") && !master.startsWith("local[")) {
-        // Stage classpath if `spark.jars` not set and not in local mode.
-        PipelineResources.prepareFilesForStaging(options);
-        // Set `spark.jars`, exclude JRE libs and jars causing conflicts using `userClassPathFirst`.
-        sparkConf.setJars(filterFilesToStage(options, SPARK_JAR_EXCLUDES));
-        // Enable `userClassPathFirst` to prevent issues with guava, jackson and others.
-        sparkConf.setIfMissing("spark.executor.userClassPathFirst", "true");
+      if (!master.startsWith("local")) {
+        if (options.getFilesToStage() != null && !options.getFilesToStage().isEmpty()) {
+          // Append the files to stage provided by the user to `spark.jars`.
+          PipelineResources.prepareFilesForStaging(options);
+          String[] filesToStage = filterFilesToStage(options, SPARK_JAR_EXCLUDES);
+          String[] jars = getSparkJars(sparkConf);
+          sparkConf.setJars(jars.length > 0 ? ArrayUtils.addAll(jars, filesToStage) : filesToStage);
+        } else if (!sparkConf.contains("spark.jars")) {
+          // Stage classpath if `spark.jars` not set and not in local mode.
+          PipelineResources.prepareFilesForStaging(options);
+          // Set `spark.jars`, exclude JRE libs and jars causing conflicts using
+          // `userClassPathFirst`.
+          sparkConf.setJars(filterFilesToStage(options, SPARK_JAR_EXCLUDES));
+          // Enable `userClassPathFirst` to prevent issues with guava, jackson and others.
+          sparkConf.setIfMissing("spark.executor.userClassPathFirst", "true");
+        }
       }
     }
 

runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/StructuredStreamingPipelineStateTest.java

@@ -104,6 +104,7 @@ public void processElement(final ProcessContext c) {
 
   private SparkStructuredStreamingPipelineOptions getBatchOptions() {
     options.setRunner(SparkStructuredStreamingRunner.class);
+    options.setSparkMaster("local[2]");
     options.setStreaming(false); // explicit because options is reused throughout the test.
     return options;
   }

sdks/go/pkg/beam/runners/prism/internal/worker/worker.go

@@ -446,7 +446,15 @@ func (wk *W) Data(data fnpb.BeamFnData_DataServer) error {
 //
 // State requests come from SDKs, and the runner responds.
 func (wk *W) State(state fnpb.BeamFnState_StateServer) error {
-	responses := make(chan *fnpb.StateResponse)
+	responses := make(chan *fnpb.StateResponse, 64)
+	sendResp := func(resp *fnpb.StateResponse) bool {
+		select {
+		case responses <- resp:
+			return true
+		case <-state.Context().Done():
+			return false
+		}
+	}
 	go func() {
 		// This go routine creates all responses to state requests from the worker
 		// so we want to close the State handler when it's all done.
@@ -472,6 +480,12 @@ func (wk *W) State(state fnpb.BeamFnState_StateServer) error {
 			wk.mu.Unlock()
 			if !ok {
 				slog.Warn("state request after bundle inactive", "instruction", req.GetInstructionId(), "worker", wk)
+				if !sendResp(&fnpb.StateResponse{
+					Id:    req.GetId(),
+					Error: fmt.Sprintf("state request for inactive bundle %v", req.GetInstructionId()),
+				}) {
+					return
+				}
 				continue
 			}
 			switch req.GetRequest().(type) {
@@ -558,13 +572,15 @@ func (wk *W) State(state fnpb.BeamFnState_StateServer) error {
 
 				// Encode the runner iterable (no length, just consecutive elements), and send it out.
 				// This is also where we can handle things like State Backed Iterables.
-				responses <- &fnpb.StateResponse{
+				if !sendResp(&fnpb.StateResponse{
 					Id: req.GetId(),
 					Response: &fnpb.StateResponse_Get{
 						Get: &fnpb.StateGetResponse{
 							Data: bytes.Join(data, []byte{}),
 						},
 					},
+				}) {
+					return
 				}
 
 			case *fnpb.StateRequest_Append:
@@ -585,11 +601,13 @@ func (wk *W) State(state fnpb.BeamFnState_StateServer) error {
 					panic(fmt.Sprintf("unsupported StateKey Append type: %T: %v", key.GetType(), prototext.Format(key)))
 				}
 
-				responses <- &fnpb.StateResponse{
+				if !sendResp(&fnpb.StateResponse{
 					Id: req.GetId(),
 					Response: &fnpb.StateResponse_Append{
 						Append: &fnpb.StateAppendResponse{},
 					},
+				}) {
+					return
 				}
 
 			case *fnpb.StateRequest_Clear:
@@ -602,7 +620,7 @@ func (wk *W) State(state fnpb.BeamFnState_StateServer) error {
 					mmkey := key.GetMultimapUserState()
 					b.OutputData.ClearMultimapState(engine.LinkID{Transform: mmkey.GetTransformId(), Local: mmkey.GetUserStateId()}, mmkey.GetWindow(), mmkey.GetKey(), mmkey.GetMapKey())
 				case *fnpb.StateKey_MultimapKeysUserState_:
-					mmkey := key.GetMultimapUserState()
+					mmkey := key.GetMultimapKeysUserState()
 					b.OutputData.ClearMultimapKeysState(engine.LinkID{Transform: mmkey.GetTransformId(), Local: mmkey.GetUserStateId()}, mmkey.GetWindow(), mmkey.GetKey())
 				case *fnpb.StateKey_OrderedListUserState_:
 					olkey := key.GetOrderedListUserState()
@@ -611,11 +629,13 @@ func (wk *W) State(state fnpb.BeamFnState_StateServer) error {
 				default:
 					panic(fmt.Sprintf("unsupported StateKey Clear type: %T: %v", key.GetType(), prototext.Format(key)))
 				}
-				responses <- &fnpb.StateResponse{
+				if !sendResp(&fnpb.StateResponse{
 					Id: req.GetId(),
 					Response: &fnpb.StateResponse_Clear{
 						Clear: &fnpb.StateClearResponse{},
 					},
+				}) {
+					return
 				}
 
 			default: