Part 1: Log systemName in DataflowWorkUnitClient, Commit, and core worker states

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/StreamingModeExecutionContext.java

@@ -268,6 +268,10 @@ public final long getBacklogBytes() {
     return backlogBytes;
   }
 
+  public String getSystemName() {
+    return systemName;
+  }
+
   public long getMaxOutputKeyBytes() {
     return operationalLimits.getMaxOutputKeyBytes();
   }
@@ -585,7 +589,7 @@ public void invalidateCache() {
       } 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;

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 fused stage %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 fused stage {} 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 fused stage %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 fused stage {} 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/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/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

@@ -66,6 +66,10 @@ public final String computationId() {
     return computationState().getComputationId();
   }
 
+  public final String systemName() {
+    return computationState().getSystemName();
+  }
+
   public @Nullable WorkItemCommitRequest singleKeyRequest() {
     return singleKeyRequest;
   };
@@ -92,8 +96,8 @@ public final int getSerializedByteSize() {
   @Override
   public String toString() {
     Work work = workBatch.get(0);
-    return "[computationId="
-        + computationId()
+    return "[systemName="
+        + systemName()
         + ", shardingKey="
         + work.getShardedKey()
         + ", workId="

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkUnitClientTest.java

@@ -121,7 +121,7 @@ public void testCloudServiceCallMapTaskStagePropagation() throws Exception {
     // Publish and acquire a map task work item, and verify we're now processing that stage.
     final String stageName = "test_stage_name";
     MapTask mapTask = new MapTask();
-    mapTask.setStageName(stageName);
+    mapTask.setSystemName(stageName);
     WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
     workItem.setMapTask(mapTask);
 
@@ -141,7 +141,7 @@ public void testCloudServiceCallSeqMapTaskStagePropagation() throws Exception {
     // Publish and acquire a seq map task work item, and verify we're now processing that stage.
     final String stageName = "test_stage_name";
     SeqMapTask seqMapTask = new SeqMapTask();
-    seqMapTask.setStageName(stageName);
+    seqMapTask.setSystemName(stageName);
     WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
     workItem.setSeqMapTask(seqMapTask);