Log slow workflow task completion diagnostics

temporal-sdk/src/main/java/io/temporal/common/converter/CodecDataConverter.java

@@ -9,6 +9,8 @@
 import io.temporal.api.failure.v1.Failure;
 import io.temporal.api.failure.v1.ResetWorkflowFailureInfo;
 import io.temporal.api.failure.v1.TimeoutFailureInfo;
+import io.temporal.internal.worker.WorkflowTaskPayloadStats;
+import io.temporal.internal.worker.WorkflowTaskPayloadStatsContext;
 import io.temporal.payload.codec.ChainCodec;
 import io.temporal.payload.codec.PayloadCodec;
 import io.temporal.payload.context.SerializationContext;
@@ -17,6 +19,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -193,13 +196,43 @@ public CodecDataConverter withContext(@Nonnull SerializationContext context) {
   @Nonnull
   @Override
   public List<Payload> encode(@Nonnull List<Payload> payloads) {
-    return ConverterUtils.withContext(chainCodec, serializationContext).encode(payloads);
+
+    long start = System.nanoTime();
+
+    List<Payload> result =
+        ConverterUtils.withContext(chainCodec, serializationContext).encode(payloads);
+
+    WorkflowTaskPayloadStats stats = WorkflowTaskPayloadStatsContext.get();
+
+    if (stats != null) {
+      stats.recordUpload(
+          TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start),
+          payloads.size(),
+          payloads.stream().mapToLong(Payload::getSerializedSize).sum());
+    }
+
+    return result;
   }
 
   @Nonnull
   @Override
   public List<Payload> decode(@Nonnull List<Payload> payloads) {
-    return ConverterUtils.withContext(chainCodec, serializationContext).decode(payloads);
+
+    long start = System.nanoTime();
+
+    List<Payload> result =
+        ConverterUtils.withContext(chainCodec, serializationContext).decode(payloads);
+
+    WorkflowTaskPayloadStats stats = WorkflowTaskPayloadStatsContext.get();
+
+    if (stats != null) {
+      stats.recordDownload(
+          TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start),
+          payloads.size(),
+          payloads.stream().mapToLong(Payload::getSerializedSize).sum());
+    }
+
+    return result;
   }
 
   private Failure.Builder encodeFailure(Failure.Builder failure) {

temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowTaskPayloadStats.java

@@ -0,0 +1,43 @@
+package io.temporal.internal.worker;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+public final class WorkflowTaskPayloadStats {
+
+  private final AtomicLong uploadTimeMillis = new AtomicLong();
+  private final AtomicLong downloadTimeMillis = new AtomicLong();
+
+  private final AtomicLong uploadedPayloads = new AtomicLong();
+  private final AtomicLong downloadedPayloads = new AtomicLong();
+
+  private final AtomicLong uploadedBytes = new AtomicLong();
+  private final AtomicLong downloadedBytes = new AtomicLong();
+
+  public void recordUpload(long millis, int count, long bytes) {
+    uploadTimeMillis.addAndGet(millis);
+    uploadedPayloads.addAndGet(count);
+    uploadedBytes.addAndGet(bytes);
+  }
+
+  public void recordDownload(long millis, int count, long bytes) {
+    downloadTimeMillis.addAndGet(millis);
+    downloadedPayloads.addAndGet(count);
+    downloadedBytes.addAndGet(bytes);
+  }
+
+  @Override
+  public String toString() {
+    return "externalStorageUploadTime="
+        + uploadTimeMillis.get()
+        + "ms uploadedPayloads="
+        + uploadedPayloads.get()
+        + " uploadedBytes="
+        + uploadedBytes.get()
+        + " externalStorageDownloadTime="
+        + downloadTimeMillis.get()
+        + "ms downloadedPayloads="
+        + downloadedPayloads.get()
+        + " downloadedBytes="
+        + downloadedBytes.get();
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowTaskPayloadStatsContext.java

@@ -0,0 +1,18 @@
+package io.temporal.internal.worker;
+
+public final class WorkflowTaskPayloadStatsContext {
+
+  private static final ThreadLocal<WorkflowTaskPayloadStats> CURRENT = new ThreadLocal<>();
+
+  public static void set(WorkflowTaskPayloadStats stats) {
+    CURRENT.set(stats);
+  }
+
+  public static WorkflowTaskPayloadStats get() {
+    return CURRENT.get();
+  }
+
+  public static void clear() {
+    CURRENT.remove();
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java

@@ -38,7 +38,8 @@
 
 final class WorkflowWorker implements SuspendableWorker {
   private static final Logger log = LoggerFactory.getLogger(WorkflowWorker.class);
-
+  private static final long SLOW_WFT_INFO_THRESHOLD_MS = 5000;
+  private static final long SLOW_WFT_WARN_THRESHOLD_MS = 10000;
   private final WorkflowRunLockManager runLocks;
 
   private final WorkflowServiceStubs service;
@@ -60,7 +61,6 @@ final class WorkflowWorker implements SuspendableWorker {
   private final NamespaceCapabilities namespaceCapabilities;
 
   private PollTaskExecutor<WorkflowTask> pollTaskExecutor;
-
   // TODO this ideally should be volatile or final (and NoopWorker should go away)
   //  Currently the implementation looks safe without volatile, but it's brittle.
   @Nonnull private SuspendableWorker poller = new NoopWorker();
@@ -435,7 +435,13 @@ public void handle(WorkflowTask task) throws Exception {
           nextWFTResponse = Optional.empty();
           boolean iterationFailed = false;
           try {
+            long workflowTaskStartTime = System.nanoTime();
+            WorkflowTaskPayloadStats payloadStats = new WorkflowTaskPayloadStats();
+
+            WorkflowTaskPayloadStatsContext.set(payloadStats);
             WorkflowTaskHandler.Result result = handleTask(currentTask, workflowTypeScope);
+            long workflowTaskCompletionMillis =
+                TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - workflowTaskStartTime);
             WorkflowTaskFailedCause taskFailedCause = null;
             try {
               RespondWorkflowTaskCompletedRequest taskCompleted = result.getTaskCompleted();
@@ -507,6 +513,7 @@ public void handle(WorkflowTask task) throws Exception {
                   if (result.getApplyPostCompletionMetrics() != null) {
                     result.getApplyPostCompletionMetrics().run();
                   }
+                  logSlowWorkflowTaskCompletion(workflowTaskCompletionMillis, payloadStats);
                 } catch (GrpcMessageTooLargeException e) {
                   // Only fail workflow task on the first attempt, subsequent failures of the same
                   // workflow task should timeout.
@@ -587,13 +594,22 @@ public void handle(WorkflowTask task) throws Exception {
         MDC.remove(LoggerTag.WORKFLOW_ID);
         MDC.remove(LoggerTag.WORKFLOW_TYPE);
         MDC.remove(LoggerTag.RUN_ID);
-
+        WorkflowTaskPayloadStatsContext.clear();
         if (locked) {
           runLocks.unlock(runId);
         }
       }
     }
 
+    private void logSlowWorkflowTaskCompletion(
+        long durationMillis, WorkflowTaskPayloadStats payloadStats) {
+      if (durationMillis > SLOW_WFT_WARN_THRESHOLD_MS) {
+        log.warn("TMPRL1104 Workflow Task completion took {} ms. {}", durationMillis, payloadStats);
+      } else if (durationMillis > SLOW_WFT_INFO_THRESHOLD_MS) {
+        log.info("TMPRL1104 Workflow Task completion took {} ms. {}", durationMillis, payloadStats);
+      }
+    }
+
     @Override
     public Throwable wrapFailure(WorkflowTask task, Throwable failure) {
       WorkflowExecution execution = task.getResponse().getWorkflowExecution();

temporal-sdk/src/test/java/io/temporal/common/converter/CodecDataConverterTest.java

@@ -1,15 +1,15 @@
 package io.temporal.common.converter;
 
-import static org.junit.Assert.assertArrayEquals;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.*;
 
 import com.google.protobuf.ByteString;
 import io.temporal.api.common.v1.Payload;
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.failure.v1.Failure;
 import io.temporal.failure.ApplicationFailure;
 import io.temporal.failure.TemporalFailure;
+import io.temporal.internal.worker.WorkflowTaskPayloadStats;
+import io.temporal.internal.worker.WorkflowTaskPayloadStatsContext;
 import io.temporal.payload.codec.PayloadCodec;
 import io.temporal.payload.codec.PayloadCodecException;
 import java.util.Collections;
@@ -119,6 +119,36 @@ public void testRawValuePassThrough() {
     assertEquals(p, converted.getPayload());
   }
 
+  @Test
+  public void testPayloadStatsAreRecordedForEncodeAndDecode() {
+    WorkflowTaskPayloadStats stats = new WorkflowTaskPayloadStats();
+    WorkflowTaskPayloadStatsContext.set(stats);
+
+    try {
+      Payload payload = Payload.newBuilder().setData(ByteString.copyFromUtf8("test-data")).build();
+
+      List<Payload> encoded = dataConverter.encode(Collections.singletonList(payload));
+
+      dataConverter.decode(encoded);
+
+      String statsOutput = stats.toString();
+
+      assertTrue(statsOutput.contains("uploadedPayloads=1"));
+      assertTrue(statsOutput.contains("downloadedPayloads=1"));
+
+      assertTrue(statsOutput.contains("uploadedBytes="));
+      assertTrue(statsOutput.contains("downloadedBytes="));
+
+      assertFalse(statsOutput.contains("uploadedBytes=0"));
+      assertFalse(statsOutput.contains("downloadedBytes=0"));
+
+      assertTrue(statsOutput.contains("externalStorageUploadTime="));
+      assertTrue(statsOutput.contains("externalStorageDownloadTime="));
+    } finally {
+      WorkflowTaskPayloadStatsContext.clear();
+    }
+  }
+
   static boolean isEncoded(Payload payload) {
     return payload.getData().startsWith(PrefixPayloadCodec.PREFIX);
   }

temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowTaskPayloadStatsTest.java

@@ -0,0 +1,34 @@
+package io.temporal.internal.worker;
+
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+
+public class WorkflowTaskPayloadStatsTest {
+
+  @Test
+  public void recordsUploadStats() {
+    WorkflowTaskPayloadStats stats = new WorkflowTaskPayloadStats();
+
+    stats.recordUpload(100, 3, 1024);
+
+    String output = stats.toString();
+
+    assertTrue(output.contains("externalStorageUploadTime=100ms"));
+    assertTrue(output.contains("uploadedPayloads=3"));
+    assertTrue(output.contains("uploadedBytes=1024"));
+  }
+
+  @Test
+  public void recordsDownloadStats() {
+    WorkflowTaskPayloadStats stats = new WorkflowTaskPayloadStats();
+
+    stats.recordDownload(200, 5, 2048);
+
+    String output = stats.toString();
+
+    assertTrue(output.contains("externalStorageDownloadTime=200ms"));
+    assertTrue(output.contains("downloadedPayloads=5"));
+    assertTrue(output.contains("downloadedBytes=2048"));
+  }
+}

temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java

@@ -6,6 +6,9 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.*;
 
+import ch.qos.logback.classic.Level;
+import ch.qos.logback.classic.spi.ILoggingEvent;
+import ch.qos.logback.core.read.ListAppender;
 import com.google.common.util.concurrent.Futures;
 import com.google.protobuf.ByteString;
 import com.uber.m3.tally.NoopScope;
@@ -33,11 +36,9 @@
 import java.util.concurrent.*;
 import org.junit.Test;
 import org.mockito.stubbing.Answer;
-import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class WorkflowWorkerTest {
-  private static final Logger log = LoggerFactory.getLogger(WorkflowWorkerTest.class);
   private final TestStatsReporter reporter = new TestStatsReporter();
   private static final String WORKFLOW_ID = "test-workflow-id";
   private static final String RUN_ID = "test-run-id";
@@ -47,6 +48,12 @@ public class WorkflowWorkerTest {
   public void concurrentPollRequestLockTest() throws Exception {
     // Test that if the server sends multiple concurrent workflow tasks for the same workflow the
     // SDK holds the lock during all processing.
+    ch.qos.logback.classic.Logger workflowWorkerLogger =
+        (ch.qos.logback.classic.Logger) LoggerFactory.getLogger(WorkflowWorker.class);
+
+    ListAppender<ILoggingEvent> appender = new ListAppender<>();
+    appender.start();
+    workflowWorkerLogger.addAppender(appender);
     WorkflowServiceStubs client = mock(WorkflowServiceStubs.class);
     when(client.getServerCapabilities())
         .thenReturn(() -> GetSystemInfoResponse.Capabilities.newBuilder().build());
@@ -197,6 +204,13 @@ public void concurrentPollRequestLockTest() throws Exception {
         });
     // Wait for the worker to respond, by this time the other blocked tasks should have timed out
     respondTaskLatch.await();
+    assertTrue(
+        appender.list.stream()
+            .anyMatch(
+                event ->
+                    event.getLevel() == Level.INFO
+                        && event.getFormattedMessage().contains("TMPRL1104")
+                        && event.getFormattedMessage().contains("Workflow Task completion took")));
     // All slots should be available
     Eventually.assertEventually(
         Duration.ofSeconds(10),
@@ -209,6 +223,7 @@ public void concurrentPollRequestLockTest() throws Exception {
               100.0);
         });
     // Cleanup
+    workflowWorkerLogger.detachAppender(appender);
     worker.shutdown(new ShutdownManager(), false).get();
     // Verify we only handled two tasks
     verify(taskHandler, times(2)).handleWorkflowTask(any());