RATIS-2397. Add trace support for Log Appender

ratis-common/src/main/java/org/apache/ratis/trace/RatisAttributes.java

@@ -33,6 +33,8 @@ public final class RatisAttributes {
   public static final AttributeKey<String> OPERATION_NAME = AttributeKey.stringKey("raft.operation.name");
   public static final AttributeKey<String> OPERATION_TYPE = AttributeKey.stringKey("raft.operation.type");
 
+  /** Number of log entries in a single {@code AppendEntries} RPC (0 for heartbeat). */
+  public static final AttributeKey<Long> APPEND_ENTRIES_COUNT = AttributeKey.longKey("raft.append.entries.count");
 
   private RatisAttributes() {
   }

ratis-common/src/main/java/org/apache/ratis/trace/SpanNames.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.ratis.trace;
+
+public final class SpanNames {
+  private static final String SEPARATOR = ".";
+
+  public static final String ASYNC_SEND = "Async::send";
+
+  public static final String RAFT_SERVER_PREFIX = "raft" + SEPARATOR + "server";
+  public static final String SUBMIT_CLIENT_REQUEST_ASYNC = RAFT_SERVER_PREFIX + SEPARATOR + "submitClientRequestAsync";
+  public static final String APPEND_ENTRIES_ASYNC = RAFT_SERVER_PREFIX + SEPARATOR + "appendEntriesAsync";
+
+  private SpanNames() {
+  }
+}
+

ratis-common/src/main/java/org/apache/ratis/trace/TraceClient.java

@@ -43,7 +43,7 @@ public static <T, THROWABLE extends Throwable> CompletableFuture<T> asyncSend(
       return action.get();
     }
     return TraceUtils.traceAsyncMethod(action,
-        () -> createClientOperationSpan(type, server, "Async::send"));
+        () -> createClientOperationSpan(type, server, SpanNames.ASYNC_SEND));
   }
 
   private static Span createClientOperationSpan(RaftClientRequest.Type type, RaftPeerId server,

ratis-common/src/main/java/org/apache/ratis/trace/TraceServer.java

@@ -20,9 +20,14 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
+import org.apache.ratis.proto.RaftProtos.AppendEntriesRequestProto;
+import org.apache.ratis.proto.RaftProtos.RaftRpcRequestProto;
+import org.apache.ratis.proto.RaftProtos.SpanContextProto;
 import org.apache.ratis.protocol.RaftClientRequest;
+import org.apache.ratis.protocol.RaftPeerId;
 import org.apache.ratis.util.function.CheckedSupplier;
 
+import java.io.IOException;
 import java.util.concurrent.CompletableFuture;
 
 /** Server-side OpenTelemetry helpers. */
@@ -56,4 +61,34 @@ private static Span createServerSpanFromClientRequest(RaftClientRequest request,
     span.setAttribute(RatisAttributes.MEMBER_ID, memberId);
     return span;
   }
+
+  /**
+   * Traces follower handling of {@link AppendEntriesRequestProto} when the leader attached trace
+   * context (client-originated) for replication.
+   */
+  public static <T> CompletableFuture<T> traceAppendEntriesAsync(
+      CheckedSupplier<CompletableFuture<T>, IOException> action,
+      AppendEntriesRequestProto request, String memberId) throws IOException {
+    if (!TraceUtils.isEnabled()) {
+      return action.get();
+    }
+    final RaftRpcRequestProto rpc = request.getServerRequest();
+    final SpanContextProto spanContext = rpc.getSpanContext();
+    // If the leader sent no parent span context, still trace as a root span
+    // rather than skipping tracing entirely.
+    final Context remoteContext = (spanContext == null || spanContext.getContextMap().isEmpty())
+        ? Context.root()
+        : TraceUtils.extractContextFromProto(spanContext);
+    return TraceUtils.traceAsyncMethod(action, () -> {
+      final Span span = TraceUtils.getGlobalTracer()
+          .spanBuilder(SpanNames.APPEND_ENTRIES_ASYNC)
+          .setParent(remoteContext)
+          .setSpanKind(SpanKind.INTERNAL)
+          .startSpan();
+      span.setAttribute(RatisAttributes.MEMBER_ID, memberId);
+      span.setAttribute(RatisAttributes.PEER_ID, String.valueOf(RaftPeerId.valueOf(rpc.getRequestorId())));
+      span.setAttribute(RatisAttributes.APPEND_ENTRIES_COUNT, (long) request.getEntriesCount());
+      return span;
+    });
+  }
 }

ratis-common/src/main/java/org/apache/ratis/trace/TraceUtils.java

@@ -81,7 +81,7 @@ public static void setTracerWhenEnabled(boolean enabled) {
     }
   }
 
-  static boolean isEnabled() {
+  public static boolean isEnabled() {
     return TRACER.get() != null;
   }
 

ratis-server/src/main/java/org/apache/ratis/server/impl/LeaderStateImpl.java

@@ -26,6 +26,7 @@
 import org.apache.ratis.proto.RaftProtos.RaftPeerRole;
 import org.apache.ratis.proto.RaftProtos.ReplicationLevel;
 import org.apache.ratis.proto.RaftProtos.RoleInfoProto;
+import org.apache.ratis.proto.RaftProtos.SpanContextProto;
 import org.apache.ratis.protocol.Message;
 import org.apache.ratis.protocol.RaftClientReply;
 import org.apache.ratis.protocol.RaftClientRequest;
@@ -353,6 +354,7 @@ boolean isApplied() {
   private final LogAppenderMetrics logAppenderMetrics;
   private final long followerMaxGapThreshold;
   private final PendingStepDown pendingStepDown;
+  private final LeaderTracer leaderTracer;
 
   private final ReadIndexHeartbeats readIndexHeartbeats;
   private final RaftServerConfigKeys.Read.ReadIndex.Type readIndexType;
@@ -387,6 +389,7 @@ boolean isApplied() {
     this.logMetadataEnabled = RaftServerConfigKeys.Log.logMetadataEnabled(properties);
     long maxPendingRequests = RaftServerConfigKeys.Write.elementLimit(properties);
     double followerGapRatioMax = RaftServerConfigKeys.Write.followerGapRatioMax(properties);
+    leaderTracer = new LeaderTracer();
 
     if (followerGapRatioMax == -1) {
       this.followerMaxGapThreshold = -1;
@@ -483,6 +486,7 @@ CompletableFuture<Void> stop() {
     raftServerMetrics.unregister();
     pendingRequests.close();
     watchRequests.close();
+    leaderTracer.close();
     return f;
   }
 
@@ -558,7 +562,9 @@ PendingRequest addPendingRequest(PendingRequests.Permit permit, RaftClientReques
       LOG.debug("{}: addPendingRequest at {}, entry={}", this, request,
           LogProtoUtils.toLogEntryString(entry.getLogEntry()));
     }
-    return pendingRequests.add(permit, request, entry);
+    final PendingRequest pending = pendingRequests.add(permit, request, entry);
+    leaderTracer.tracePendingRequest(pending);
+    return pending;
   }
 
   CompletableFuture<RaftClientReply> streamAsync(RaftClientRequest request) {
@@ -645,9 +651,10 @@ public AppendEntriesRequestProto newAppendEntriesRequestProto(FollowerInfo follo
       List<LogEntryProto> entries, TermIndex previous, long callId) {
     final boolean initializing = !isCaughtUp(follower);
     final RaftPeerId targetId = follower.getId();
+    final SpanContextProto tracing = leaderTracer.traceAppendEntries(entries);
     return ServerProtoUtils.toAppendEntriesRequestProto(server.getMemberId(), targetId, getCurrentTerm(), entries,
         ServerImplUtils.effectiveCommitIndex(readIndexSupplier.get(), previous, entries.size()),
-        initializing, previous, server.getCommitInfos(), callId);
+        initializing, previous, server.getCommitInfos(), callId, tracing);
   }
 
   /**
@@ -1243,6 +1250,7 @@ private boolean checkLeaderLease() {
 
   void replyPendingRequest(TermIndex termIndex, RaftClientReply reply, RetryCacheImpl.CacheEntry cacheEntry) {
     final PendingRequest pending = pendingRequests.remove(termIndex);
+    leaderTracer.removePendingRequest(pending);
 
     final LongSupplier replyMethod = () -> {
       cacheEntry.updateResult(reply);

ratis-server/src/main/java/org/apache/ratis/server/impl/LeaderTracer.java

@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.ratis.server.impl;
+
+import org.apache.ratis.proto.RaftProtos.LogEntryProto;
+import org.apache.ratis.proto.RaftProtos.SpanContextProto;
+import org.apache.ratis.trace.TraceUtils;
+import org.apache.ratis.util.AutoCloseableLock;
+
+import java.util.List;
+import java.util.NavigableMap;
+import java.util.TreeMap;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+class LeaderTracer {
+  private final AppendEntriesSpans appendEntriesSpans;
+  public LeaderTracer() {
+    appendEntriesSpans = new AppendEntriesSpans();
+  }
+  /**
+   * Client-originated trace context keyed by log index and propagated in appendEntries requests,
+   * so follower appendEntries spans can join the same trace as the client write.
+   */
+  static class AppendEntriesSpans {
+    private final NavigableMap<Long, SpanContextProto> sorted = new TreeMap<>();
+    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+
+    SpanContextProto get(long first, long last) {
+      try (AutoCloseableLock ignored = AutoCloseableLock.acquire(lock.readLock())) {
+        for (SpanContextProto sc : sorted.subMap(first, true, last, true).values()) {
+          if (sc != null && !sc.getContextMap().isEmpty()) {
+            return sc;
+          }
+        }
+      }
+      return null;
+    }
+
+    void put(long index, SpanContextProto spanContext) {
+      try (AutoCloseableLock ignored = AutoCloseableLock.acquire(lock.writeLock())) {
+        sorted.put(index, spanContext);
+      }
+    }
+
+    void remove(long index) {
+      try (AutoCloseableLock ignored = AutoCloseableLock.acquire(lock.writeLock())) {
+        sorted.remove(index);
+      }
+    }
+
+    void clear() {
+      try (AutoCloseableLock ignored = AutoCloseableLock.acquire(lock.writeLock())) {
+        sorted.clear();
+      }
+    }
+  }
+
+  void tracePendingRequest(PendingRequest pending) {
+    if (pending == null || !TraceUtils.isEnabled()) {
+      return;
+    }
+    final SpanContextProto spanContext = pending.getRequest().getSpanContext();
+    if (spanContext != null && !spanContext.getContextMap().isEmpty()) {
+      appendEntriesSpans.put(pending.getTermIndex().getIndex(), spanContext);
+    }
+  }
+
+  void removePendingRequest(PendingRequest pending) {
+    appendEntriesSpans.remove(pending.getTermIndex().getIndex());
+  }
+
+  SpanContextProto traceAppendEntries(List<LogEntryProto> entries) {
+    if (entries == null || entries.isEmpty()) {
+      return null;
+    }
+    final long first = entries.get(0).getIndex();
+    final long last = entries.get(entries.size() - 1).getIndex();
+    return appendEntriesSpans.get(first, last);
+  }
+
+  void close() {
+    appendEntriesSpans.clear();
+  }
+}

ratis-server/src/main/java/org/apache/ratis/server/impl/RaftServerImpl.java

@@ -101,6 +101,7 @@
 import org.apache.ratis.statemachine.impl.TransactionContextImpl;
 import org.apache.ratis.thirdparty.com.google.common.annotations.VisibleForTesting;
 import org.apache.ratis.thirdparty.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.ratis.trace.SpanNames;
 import org.apache.ratis.trace.TraceServer;
 import org.apache.ratis.trace.TraceUtils;
 import org.apache.ratis.util.CodeInjectionForTesting;
@@ -982,7 +983,7 @@ public CompletableFuture<RaftClientReply> submitClientRequestAsync(
       RaftClientRequest request) throws IOException {
     return TraceServer.traceAsyncMethod(
         () -> submitClientRequestAsyncInternal(request),
-        request, getMemberId().toString(), "raft.server.submitClientRequestAsync");
+        request, getMemberId().toString(), SpanNames.SUBMIT_CLIENT_REQUEST_ASYNC);
   }
 
   private CompletableFuture<RaftClientReply> submitClientRequestAsyncInternal(
@@ -1553,7 +1554,6 @@ public CompletableFuture<AppendEntriesReplyProto> appendEntriesAsync(AppendEntri
     try {
       final RaftPeerId leaderId = RaftPeerId.valueOf(request.getRequestorId());
       final RaftGroupId leaderGroupId = ProtoUtils.toRaftGroupId(request.getRaftGroupId());
-
       CodeInjectionForTesting.execute(APPEND_ENTRIES, getId(), leaderId, previous, r);
 
       assertLifeCycleState(LifeCycle.States.STARTING_OR_RUNNING);
@@ -1562,8 +1562,8 @@ public CompletableFuture<AppendEntriesReplyProto> appendEntriesAsync(AppendEntri
       }
       assertGroup(getMemberId(), leaderId, leaderGroupId);
       assertEntries(r, previous, state);
-
-      return appendEntriesAsync(leaderId, request.getCallId(), previous, r);
+      return TraceServer.traceAppendEntriesAsync(() -> appendEntriesAsync(leaderId, request.getCallId(), previous, r),
+          r, getMemberId().toString());
     } catch(Exception t) {
       LOG.error("{}: Failed appendEntries* {}", getMemberId(),
           toAppendEntriesRequestString(r, stateMachine::toStateMachineLogEntryString), t);

ratis-server/src/main/java/org/apache/ratis/server/impl/ServerProtoUtils.java

@@ -157,9 +157,13 @@ static AppendEntriesReplyProto toAppendEntriesReplyProto(
   static AppendEntriesRequestProto toAppendEntriesRequestProto(
       RaftGroupMemberId requestorId, RaftPeerId replyId, long leaderTerm,
       List<LogEntryProto> entries, long leaderCommit, boolean initializing,
-      TermIndex previous, Collection<CommitInfoProto> commitInfos, long callId) {
+      TermIndex previous, Collection<CommitInfoProto> commitInfos, long callId,
+      SpanContextProto tracingContext) {
     final RaftRpcRequestProto.Builder rpcRequest = ClientProtoUtils.toRaftRpcRequestProtoBuilder(requestorId, replyId)
         .setCallId(callId);
+    if (tracingContext != null && !tracingContext.getContextMap().isEmpty()) {
+      rpcRequest.setSpanContext(tracingContext);
+    }
     final AppendEntriesRequestProto.Builder b = AppendEntriesRequestProto
         .newBuilder()
         .setServerRequest(rpcRequest)