Log Fused Stage Name Instead of computation id in more places

.github/workflows/beam_PostCommit_Java_Delta_IO_Dataflow.yml

@@ -64,8 +64,6 @@ jobs:
         job_phrase: ["Run PostCommit Java Delta IO Dataflow"]
     steps:
       - uses: actions/checkout@v7
-        with:
-          persist-credentials: false
       - name: Setup repository
         uses: ./.github/actions/setup-action
         with:

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkUnitClient.java

@@ -135,13 +135,13 @@ public Optional<WorkItem> getWorkItem() throws IOException {
 
     final String stage;
     if (work.getMapTask() != null) {
-      stage = work.getMapTask().getStageName();
+      stage = work.getMapTask().getSystemName();
       logger.info("Starting MapTask stage {}", stage);
     } else if (work.getSeqMapTask() != null) {
-      stage = work.getSeqMapTask().getStageName();
+      stage = work.getSeqMapTask().getSystemName();
       logger.info("Starting SeqMapTask stage {}", stage);
     } else if (work.getSourceOperationTask() != null) {
-      stage = work.getSourceOperationTask().getStageName();
+      stage = work.getSourceOperationTask().getSystemName();
       logger.info("Starting SourceOperationTask stage {}", stage);
     } else {
       stage = null;

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

@@ -924,7 +924,7 @@ static StreamingDataflowWorker forTesting(
                 mapTask,
                 workExecutor,
                 stateNameMap,
-                stateCache.forComputation(mapTask.getStageName())));
+                stateCache.forComputation(mapTask.getStageName(), mapTask.getSystemName())));
     MemoryMonitor memoryMonitor = MemoryMonitor.fromOptions(options);
     FailureTracker failureTracker =
         options.isEnableStreamingEngine()
@@ -1197,26 +1197,23 @@ void stop() {
   }
 
   private void onCompleteCommit(CompleteCommit completeCommit) {
+    Optional<ComputationState> computationState =
+        computationStateCache.getIfPresent(completeCommit.computationId());
     if (completeCommit.status() != Windmill.CommitStatus.OK) {
       readerCache.invalidateReader(
           WindmillComputationKey.create(
               completeCommit.computationId(), completeCommit.shardedKey()));
-      stateCache
-          .forComputation(completeCommit.computationId())
-          .invalidate(completeCommit.shardedKey());
+      computationState.ifPresent(
+          state ->
+              stateCache
+                  .forComputation(completeCommit.computationId(), state.getSystemName())
+                  .invalidate(completeCommit.shardedKey()));
     }
 
-    computationStateCache
-        .getIfPresent(completeCommit.computationId())
-        .ifPresent(
-            state -> {
-              if (completeCommit.retryableFailure()) {
-                state.reexecuteActiveWork(completeCommit.shardedKey(), completeCommit.workId());
-              } else {
-                state.completeWorkAndScheduleNextWorkForKey(
-                    completeCommit.shardedKey(), completeCommit.workId());
-              }
-            });
+    computationState.ifPresent(
+        state ->
+            state.completeWorkAndScheduleNextWorkForKey(
+                completeCommit.shardedKey(), completeCommit.workId()));
   }
 
   @AutoValue

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java

@@ -52,7 +52,6 @@
 import org.apache.beam.runners.dataflow.worker.counters.NameContext;
 import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
 import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
-import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
 import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
 import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
@@ -268,6 +267,10 @@
     return backlogBytes;
   }
 
+  public String getSystemName() {
+    return systemName;
+  }
+
   public long getMaxOutputKeyBytes() {
     return operationalLimits.getMaxOutputKeyBytes();
   }
@@ -585,7 +588,7 @@
       } catch (IOException e) {
         Windmill.WorkItem workItem = getWorkItem();
         long shardingKey = workItem != null ? workItem.getShardingKey() : -1L;
-        LOG.warn("Failed to close reader for {}-{}", computationId, shardingKey, e);
+        LOG.warn("Failed to close reader for {}-{}", systemName, shardingKey, e);
       }
     }
     activeReader = null;
@@ -726,11 +729,6 @@
         buildWorkItemTruncationRequestBuilder(currentWork, estimatedCommitSize);
     currentBuilder.clear();
     currentBuilder.mergeFrom(truncationBuilder.build());
-
-    // TODO: throw and retry when truncation is not on a single key bundle.
-    checkState(
-        !multiKeyBundleOptions.multiKeyBundleEnabled(),
-        "Commit truncation not implemented for multikey bundles");
   }
 
   private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBuilder(
@@ -774,7 +772,7 @@
     }
 
     @Nullable
     ExecutableWork additionalWork =
         executor.pollWork(computationId, activeWork.getKeyGroup(), handle);
     if (additionalWork != null) {
       flushStateInternal();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java

@@ -265,20 +265,26 @@ public long add(WindowedValue<T> data) throws IOException {
       }
       if (key.size() > context.getMaxOutputKeyBytes()) {
         if (context.throwExceptionsForLargeOutput()) {
-          throw new OutputTooLargeException("Key too large: " + key.size());
+          throw new OutputTooLargeException(
+              String.format(
+                  "Key for system %s too large: %s", context.getSystemName(), key.size()));
         } else {
           LOG.error(
-              "Trying to output too large key with size {}. Limit is {}. See https://cloud.google.com/dataflow/docs/guides/common-errors#key-commit-too-large-exception. Running with --experiments=throw_exceptions_on_large_output will instead throw an OutputTooLargeException which may be caught in user code.",
+              "Trying to output too large key for system {} with size {}. Limit is {}. See https://cloud.google.com/dataflow/docs/guides/common-errors#key-commit-too-large-exception. Running with --experiments=throw_exceptions_on_large_output will instead throw an OutputTooLargeException which may be caught in user code.",
+              context.getSystemName(),
               key.size(),
               context.getMaxOutputKeyBytes());
         }
       }
       if (value.size() > context.getMaxOutputValueBytes()) {
         if (context.throwExceptionsForLargeOutput()) {
-          throw new OutputTooLargeException("Value too large: " + value.size());
+          throw new OutputTooLargeException(
+              String.format(
+                  "Value for system %s too large: %s", context.getSystemName(), value.size()));
         } else {
           LOG.error(
-              "Trying to output too large value with size {}. Limit is {}. See https://cloud.google.com/dataflow/docs/guides/common-errors#key-commit-too-large-exception. Running with --experiments=throw_exceptions_on_large_output will instead throw an OutputTooLargeException which may be caught in user code.",
+              "Trying to output too large value for system {} with size {}. Limit is {}. See https://cloud.google.com/dataflow/docs/guides/common-errors#key-commit-too-large-exception. Running with --experiments=throw_exceptions_on_large_output will instead throw an OutputTooLargeException which may be caught in user code.",
+              context.getSystemName(),
               value.size(),
               context.getMaxOutputValueBytes());
         }

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

@@ -185,7 +185,7 @@ synchronized void failWorkForKey(ImmutableList<WorkIdWithShardingKey> failedWork
       executableWork.work().setFailed();
       LOG.debug(
           "Failing work {} {}. The work will be retried and is not lost.",
-          computationStateCache.getComputation(),
+          computationStateCache.getSystemName(),
           failedId);
     }
   }

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

@@ -69,6 +69,10 @@ public String getComputationId() {
     return computationId;
   }
 
+  public String getSystemName() {
+    return mapTask.getSystemName();
+  }
+
   public MapTask getMapTask() {
     return mapTask;
   }

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

@@ -28,6 +28,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.beam.runners.dataflow.worker.apiary.FixMultiOutputInfosOnParDoInstructions;
@@ -77,7 +78,8 @@ private ComputationStateCache(
   public static ComputationStateCache create(
       ComputationConfig.Fetcher computationConfigFetcher,
       BoundedQueueExecutor workUnitExecutor,
-      Function<String, WindmillStateCache.ForComputation> perComputationStateCacheViewFactory,
+      BiFunction<String, String, WindmillStateCache.ForComputation>
+          perComputationStateCacheViewFactory,
       IdGenerator idGenerator) {
     Function<MapTask, MapTask> fixMultiOutputInfosOnParDoInstructions =
         new FixMultiOutputInfosOnParDoInstructions(idGenerator);
@@ -105,7 +107,8 @@ public ComputationState load(String computationId) {
                         fixMultiOutputInfosOnParDoInstructions.apply(computationConfig.mapTask()),
                         workUnitExecutor,
                         transformUserNameToStateFamilyForComputation,
-                        perComputationStateCacheViewFactory.apply(computationId));
+                        perComputationStateCacheViewFactory.apply(
+                            computationId, computationConfig.mapTask().getSystemName()));
                   }
                 }),
         fixMultiOutputInfosOnParDoInstructions,
@@ -116,7 +119,8 @@ public ComputationState load(String computationId) {
   public static ComputationStateCache forTesting(
       ComputationConfig.Fetcher computationConfigFetcher,
       BoundedQueueExecutor workUnitExecutor,
-      Function<String, WindmillStateCache.ForComputation> perComputationStateCacheViewFactory,
+      BiFunction<String, String, WindmillStateCache.ForComputation>
+          perComputationStateCacheViewFactory,
       IdGenerator idGenerator,
       ConcurrentMap<String, String> pipelineUserNameToStateFamilyNameMap) {
     ComputationStateCache cache =
@@ -205,7 +209,7 @@ public void closeAndInvalidateAll() {
   public void appendSummaryHtml(PrintWriter writer) {
     writer.println("<h1>Specs</h1>");
     for (ComputationState computationState : getAllPresentComputations()) {
-      writer.println("<h3>" + computationState.getComputationId() + "</h3>");
+      writer.println("<h3>" + computationState.getSystemName() + "</h3>");
       writer.print("<script>document.write(JSON.stringify(");
       writer.print(computationState.getMapTask().toString());
       writer.println(", null, \"&nbsp&nbsp\").replace(/\\n/g, \"<br>\"))</script>");

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

@@ -59,7 +59,7 @@ public void appendSummaryHtml(PrintWriter writer) {
 
     writer.println("Active Keys: <br>");
     for (ComputationState computationState : allComputationStates.get()) {
-      writer.print(computationState.getComputationId());
+      writer.print(computationState.getSystemName());
       writer.print(":<br>");
       computationState.printActiveWork(writer);
       writer.println("<br>");

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java

@@ -30,7 +30,7 @@
 
 /** Value class for a queued commit. */
 @Internal
 public class Commit {
 
   private final ComputationState computationState;
   private final ImmutableList<Work> workBatch;
@@ -66,9 +66,11 @@
     return computationState().getComputationId();
   }
 
-  public @Nullable WorkItemCommitRequest singleKeyRequest() {
-    return singleKeyRequest;
-  };
+  public final String systemName() {
+    return computationState().getSystemName();
+  }
+
+  public abstract WorkItemCommitRequest request();
 
   public ComputationState computationState() {
     return computationState;
@@ -86,7 +88,7 @@
     if (multiKeyRequest() != null) {
       return checkStateNotNull(multiKeyRequest()).getSerializedSize();
     }
     return checkStateNotNull(singleKeyRequest()).getSerializedSize();
   }
 
   @Override

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java

@@ -112,7 +112,12 @@ public void commit(Commit commit) {
     // Do this check after adding to commitQueue, else commitQueue.put() can race with
     // drainCommitQueue() in stop() and leave commits orphaned in the queue.
     if (!this.isRunning.get()) {
-      LOG.debug("Trying to queue commit on shutdown, failing commit={}", commit);
+      LOG.debug(
+          "Trying to queue commit on shutdown, failing commit=[systemName={}, shardingKey={},"
+              + " workId={} ].",
+          commit.systemName(),
+          commit.work().getShardedKey(),
+          commit.work().id());
       drainCommitQueue();
     }
   }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateCache.java

@@ -170,8 +170,8 @@ public CacheStats getCacheStats() {
   }
 
   /** Returns a per-computation view of the state cache. */
-  public ForComputation forComputation(String computation) {
-    return new ForComputation(computation);
+  public ForComputation forComputation(String computation, String systemName) {
+    return new ForComputation(computation, systemName);
   }
 
   /** Print summary statistics of the cache to the given {@link PrintWriter}. */
@@ -353,16 +353,23 @@ private Optional<T> value() {
   public class ForComputation {
 
     private final String computation;
+    private final String systemName;
 
-    private ForComputation(String computation) {
+    private ForComputation(String computation, String systemName) {
       this.computation = computation;
+      this.systemName = systemName;
     }
 
     /** Returns the computation associated to this class. */
     public String getComputation() {
       return this.computation;
     }
 
+    /** Returns the system name associated to this class. */
+    public String getSystemName() {
+      return this.systemName;
+    }
+
     /** Invalidate all cache entries for this computation and {@code processingKey}. */
     public void invalidate(ByteString processingKey, long shardingKey) {
       WindmillComputationKey key =

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java

@@ -20,6 +20,7 @@
 import static org.apache.beam.runners.dataflow.DataflowRunner.hasExperiment;
 
 import com.google.api.services.dataflow.model.MapTask;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import org.apache.beam.runners.dataflow.internal.CustomSources;
 import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions;
@@ -82,12 +83,11 @@
   private final DataflowWorkerHarnessOptions options;
   private final DataflowMapTaskExecutorFactory mapTaskExecutorFactory;
   private final ReaderCache readerCache;
-  private final Function<String, WindmillStateCache.ForComputation> stateCacheFactory;
+  private final BiFunction<String, String, WindmillStateCache.ForComputation> stateCacheFactory;
   private final ReaderRegistry readerRegistry;
   private final SinkRegistry sinkRegistry;
   private final DataflowExecutionStateSampler sampler;
   private final CounterSet pendingDeltaCounters;
-  private final SideInputStateFetcherFactory sideInputStateFetcherFactory;
   private final StreamingCounters streamingCounters;
   private final FailureTracker failureTracker;
 
@@ -112,7 +112,7 @@
       DataflowWorkerHarnessOptions options,
       DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
       ReaderCache readerCache,
-      Function<String, WindmillStateCache.ForComputation> stateCacheFactory,
+      BiFunction<String, String, WindmillStateCache.ForComputation> stateCacheFactory,
       DataflowExecutionStateSampler sampler,
       StreamingCounters streamingCounters,
       FailureTracker failureTracker,
@@ -141,7 +141,7 @@
     this.throwExceptionOnLargeOutput =
         hasExperiment(options, THROW_EXCEPTIONS_ON_LARGE_OUTPUT_EXPERIMENT);
     this.hotKeyLogger = hotKeyLogger;
     this.sideInputStateFetcherFactory = sideInputStateFetcherFactory;
     this.multiKeyBundleOptions = multiKeyBundleOptions;
   }
 
@@ -287,7 +287,7 @@
         computationId,
         readerCache,
         computationState.getTransformUserNameToStateFamily(),
-        stateCacheFactory.apply(computationId),
+        stateCacheFactory.apply(computationId, stageInfo.systemName()),
         stageInfo.metricsContainerRegistry(),
         executionStateTracker,
         stageInfo.executionStateRegistry(),
@@ -302,7 +302,7 @@
         failureTracker,
         computationState.sourceBytesProcessCounterName(),
         multiKeyBundleOptions,
         sideInputStateFetcherFactory);
   }
 
   private DataflowMapTaskExecutor createMapTaskExecutor(