[AST-494] request-id distributed logging interceptor

Author: iamajais

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();
 

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();
 

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);
+      }
+    };
+  }
+}