From 73f728cb55285f46601463eb6038471dad55c970 Mon Sep 17 00:00:00 2001
From: "ayush.jaiswal" <ayush.jaiswal@harness.io>
Date: Wed, 17 Dec 2025 13:23:49 +0530
Subject: [PATCH] [AST-494] request-id distributed logging interceptor

---
 .../grpcutils/context/FastUUIDGenerator.java  |  4 +-
 .../grpcutils/context/RequestContext.java     |  6 ++
 ...equestContextLoggingServerInterceptor.java | 73 +++++++++++++++++++
 3 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/RequestContextLoggingServerInterceptor.java

diff --git a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/FastUUIDGenerator.java b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/FastUUIDGenerator.java
index 9cbfc5c..c2853cf 100644
--- a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/FastUUIDGenerator.java
+++ b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/FastUUIDGenerator.java
@@ -3,13 +3,13 @@
 import java.util.UUID;
 import java.util.concurrent.ThreadLocalRandom;
 
-class FastUUIDGenerator {
+public class FastUUIDGenerator {
   /**
    * This function generates UUIDs using ThreadLocalRandom, which is faster and doesn't block like
    * the default randomUUID method that relies on /dev/random. It's suitable for most random UUID
    * needs.
    */
-  static UUID randomUUID() {
+  public static UUID randomUUID() {
     long mostSigBits = ThreadLocalRandom.current().nextLong();
     long leastSigBits = ThreadLocalRandom.current().nextLong();
 
diff --git a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java
index 0aac33f..1e67f15 100644
--- a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java
+++ b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java
@@ -47,6 +47,12 @@ public static RequestContext forTenantId(String tenantId) {
             FastUUIDGenerator.randomUUID().toString());
   }
 
+  public static RequestContext withRequestId(String tenantId, String requestId) {
+    return new RequestContext()
+        .put(RequestContextConstants.TENANT_ID_HEADER_KEY, tenantId)
+        .put(RequestContextConstants.REQUEST_ID_HEADER_KEY, requestId);
+  }
+
   public static RequestContext fromMetadata(Metadata metadata) {
     RequestContext requestContext = new RequestContext();
 
diff --git a/grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/RequestContextLoggingServerInterceptor.java b/grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/RequestContextLoggingServerInterceptor.java
new file mode 100644
index 0000000..ba99f8e
--- /dev/null
+++ b/grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/RequestContextLoggingServerInterceptor.java
@@ -0,0 +1,73 @@
+package org.hypertrace.core.grpcutils.server;
+
+import static org.hypertrace.core.grpcutils.context.RequestContextConstants.REQUEST_ID_HEADER_KEY;
+
+import io.grpc.Context;
+import io.grpc.Contexts;
+import io.grpc.ForwardingServerCallListener;
+import io.grpc.Metadata;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import java.util.Optional;
+import lombok.extern.slf4j.Slf4j;
+import org.hypertrace.core.grpcutils.context.FastUUIDGenerator;
+import org.hypertrace.core.grpcutils.context.RequestContext;
+import org.slf4j.MDC;
+
+@Slf4j
+public final class RequestContextLoggingServerInterceptor implements ServerInterceptor {
+
+  @Override
+  public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
+      ServerCall<ReqT, RespT> serverCall,
+      Metadata metadata,
+      ServerCallHandler<ReqT, RespT> serverCallHandler) {
+    RequestContext currentContext =
+        Optional.ofNullable(RequestContext.CURRENT.get())
+            .orElseGet(() -> RequestContext.fromMetadata(metadata));
+    Optional<String> opRequestId = currentContext.getHeaderValue(REQUEST_ID_HEADER_KEY);
+    if (opRequestId.isEmpty()) {
+      opRequestId = Optional.of(FastUUIDGenerator.randomUUID().toString());
+    }
+    final String requestId = opRequestId.get();
+    ServerCall.Listener<ReqT> listener =
+        Contexts.interceptCall(
+            Context.current().withValue(RequestContext.CURRENT, currentContext),
+            serverCall,
+            metadata,
+            serverCallHandler);
+    return new ForwardingServerCallListener.SimpleForwardingServerCallListener<>(listener) {
+
+      @Override
+      public void onCancel() {
+        try {
+          MDC.clear();
+        } catch (Exception e) {
+          log.error("Error while clearing request context details from MDC params", e);
+        }
+        super.onCancel();
+      }
+
+      @Override
+      public void onComplete() {
+        try {
+          MDC.clear();
+        } catch (Exception e) {
+          log.error("Error while clearing request context details from MDC params", e);
+        }
+        super.onComplete();
+      }
+
+      @Override
+      public void onMessage(ReqT message) {
+        try {
+          MDC.put(REQUEST_ID_HEADER_KEY, requestId);
+        } catch (Exception e) {
+          log.error("Error while setting request context details in MDC params", e);
+        }
+        super.onMessage(message);
+      }
+    };
+  }
+}