From 113d1e5f47522d23ccd63a76ce85268ac6924a06 Mon Sep 17 00:00:00 2001
From: Wei-Chiu Chuang <weichiu@apache.org>
Date: Wed, 13 Sep 2023 15:45:46 -0700
Subject: [PATCH] Add starting point traces for StateMachine

Change-Id: I2cf55dcfa91fd2733880dc042a6dcec4c3dd6c81
(cherry picked from commit 4ca599671c843f0734974e03c0f972440e7df61c)
---
 .../server/ratis/ContainerStateMachine.java   | 32 +++++++++++++++++--
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/transport/server/ratis/ContainerStateMachine.java b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/transport/server/ratis/ContainerStateMachine.java
index 0758b245fbe2..455f0789d8c4 100644
--- a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/transport/server/ratis/ContainerStateMachine.java
+++ b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/transport/server/ratis/ContainerStateMachine.java
@@ -56,6 +56,7 @@
 import org.apache.hadoop.hdds.scm.ScmConfigKeys;
 import org.apache.hadoop.hdds.scm.container.common.helpers.ContainerNotOpenException;
 import org.apache.hadoop.hdds.scm.container.common.helpers.StorageContainerException;
+import org.apache.hadoop.hdds.tracing.TracingUtil;
 import org.apache.hadoop.hdds.utils.Cache;
 import org.apache.hadoop.hdds.utils.ResourceLimitCache;
 import org.apache.hadoop.ozone.OzoneConfigKeys;
@@ -94,6 +95,8 @@
 import org.apache.ratis.util.TaskQueue;
 import org.apache.ratis.util.function.CheckedSupplier;
 import org.apache.ratis.util.JavaUtils;
+
+import io.opentracing.Span;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -348,9 +351,21 @@ public long takeSnapshot() throws IOException {
   @Override
   public TransactionContext startTransaction(RaftClientRequest request)
       throws IOException {
-    long startTime = Time.monotonicNowNanos();
     final ContainerCommandRequestProto proto =
         message2ContainerCommandRequestProto(request.getMessage());
+    Span span = TracingUtil.importAndCreateSpan("startTransaction", proto.getTraceID());
+    try {
+      return startTransactionTraced(request, proto);
+    } finally {
+      span.finish();
+    }
+  }
+
+  private TransactionContext startTransactionTraced(RaftClientRequest request,
+      ContainerCommandRequestProto proto)
+      throws IOException {
+    long startTime = Time.monotonicNowNanos();
+
     Preconditions.checkArgument(request.getRaftGroupId().equals(gid));
     try {
       dispatcher.validateContainerCommand(proto);
@@ -491,6 +506,7 @@ private CompletableFuture<Message> handleWriteChunk(
     // thread.
     CompletableFuture<ContainerCommandResponseProto> writeChunkFuture =
         CompletableFuture.supplyAsync(() -> {
+          Span span = TracingUtil.importAndCreateSpan("writeStateMachineData", requestProto.getTraceID());
           try {
             return runCommand(requestProto, context);
           } catch (Exception e) {
@@ -503,6 +519,8 @@ private CompletableFuture<Message> handleWriteChunk(
             stateMachineHealthy.set(false);
             raftFuture.completeExceptionally(e);
             throw e;
+          } finally {
+            span.finish();
           }
         }, getChunkExecutor(requestProto.getWriteChunk()));
 
@@ -709,8 +727,13 @@ private ByteString readStateMachineData(
         new DispatcherContext.Builder().setTerm(term).setLogIndex(index)
             .setReadFromTmpFile(true).build();
     // read the chunk
-    ContainerCommandResponseProto response =
-        dispatchCommand(dataContainerCommandProto, context);
+    Span span = TracingUtil.importAndCreateSpan("readStateMachineData", requestProto.getTraceID());
+    ContainerCommandResponseProto response;
+    try {
+      response = dispatchCommand(dataContainerCommandProto, context);
+    } finally {
+      span.finish();
+    }
     if (response.getResult() != ContainerProtos.Result.SUCCESS) {
       StorageContainerException sce =
           new StorageContainerException(response.getMessage(),
@@ -853,11 +876,14 @@ private CompletableFuture<ContainerCommandResponseProto> submitTask(
     final long containerId = request.getContainerID();
     final CheckedSupplier<ContainerCommandResponseProto, Exception> task
         = () -> {
+          Span span = TracingUtil.importAndCreateSpan("applyTransaction", request.getTraceID());
           try {
             return runCommand(request, context.build());
           } catch (Exception e) {
             exceptionHandler.accept(e);
             throw e;
+          } finally {
+            span.finish();
           }
         };
     return containerTaskQueues.submit(containerId, task, executor);