Switch logging to fused stage name in more places. This simplifies operations by removing backend details

Author: rwiggles

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

@@ -922,7 +922,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()
@@ -1194,21 +1194,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 ->
-                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/streaming/ActiveWorkState.java

@@ -180,7 +180,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/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/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;
@@ -81,7 +82,7 @@ final class ComputationWorkExecutorFactory {
   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;
@@ -110,7 +111,7 @@ final class ComputationWorkExecutorFactory {
       DataflowWorkerHarnessOptions options,
       DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
       ReaderCache readerCache,
-      Function<String, WindmillStateCache.ForComputation> stateCacheFactory,
+      BiFunction<String, String, WindmillStateCache.ForComputation> stateCacheFactory,
       DataflowExecutionStateSampler sampler,
       StreamingCounters streamingCounters,
       FailureTracker failureTracker,
@@ -283,7 +284,7 @@ private StreamingModeExecutionContext createExecutionContext(
         computationId,
         readerCache,
         computationState.getTransformUserNameToStateFamily(),
-        stateCacheFactory.apply(computationId),
+        stateCacheFactory.apply(computationId, stageInfo.systemName()),
         stageInfo.metricsContainerRegistry(),
         executionStateTracker,
         stageInfo.executionStateRegistry(),

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

@@ -26,7 +26,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 import java.util.function.Supplier;
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair;
@@ -113,7 +113,7 @@ public static StreamingWorkScheduler create(
       DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
       BoundedQueueExecutor workExecutor,
       ScheduledExecutorService commitFinalizerCleanupExecutor,
-      Function<String, WindmillStateCache.ForComputation> stateCacheFactory,
+      BiFunction<String, String, WindmillStateCache.ForComputation> stateCacheFactory,
       FailureTracker failureTracker,
       WorkFailureProcessor workFailureProcessor,
       StreamingCounters streamingCounters,

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

@@ -128,7 +128,7 @@ private StreamingModeExecutionContext createExecutionContext(
         WindmillStateCache.builder()
             .setSizeMb(options.getWorkerCacheMb())
             .build()
-            .forComputation("comp"),
+            .forComputation("comp", "systemName"),
         StreamingStepMetricsContainer.createRegistry(),
         new DataflowExecutionStateTracker(
             ExecutionStateSampler.newForTest(),

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

@@ -1003,7 +1003,7 @@ public void testFailedWorkItemsAbort() throws Exception {
             WindmillStateCache.builder()
                 .setSizeMb(options.getWorkerCacheMb())
                 .build()
-                .forComputation(COMPUTATION_ID),
+                .forComputation(COMPUTATION_ID, "systemName"),
             StreamingStepMetricsContainer.createRegistry(),
             new DataflowExecutionStateTracker(
                 ExecutionStateSampler.newForTest(),

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

@@ -84,7 +84,10 @@ private static ExecutableWork createWork(ShardedKey shardedKey, long workToken,
   public void setUp() {
     computationStateCache =
         ComputationStateCache.create(
-            configFetcher, workExecutor, ignored -> stateCache, IdGenerators.decrementingLongs());
+            configFetcher,
+            workExecutor,
+            (ignored1, ignored2) -> stateCache,
+            IdGenerators.decrementingLongs());
   }
 
   @Test