add interceptor to add unique x-client-trace-id per RPC call

forcedotcom · Jun 25, 2024 · acca43e · acca43e
1 parent 3e6dbaa
commit acca43e
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/java/src/main/java/utility/CommonContext.java b/java/src/main/java/utility/CommonContext.java
@@ -65,8 +65,10 @@ public CommonContext(final ExampleConfigurations options) {
         callCredentials = setupCallCredentials(options);
         sessionToken = ((APISessionCredentials) callCredentials).getToken();
 
-        asyncStub = PubSubGrpc.newStub(channel).withCallCredentials(callCredentials);
-        blockingStub = PubSubGrpc.newBlockingStub(channel).withCallCredentials(callCredentials);
+        Channel interceptedChannel = ClientInterceptors.intercept(channel, new XClientTraceIdClientInterceptor());
+
+        asyncStub = PubSubGrpc.newStub(interceptedChannel).withCallCredentials(callCredentials);
+        blockingStub = PubSubGrpc.newBlockingStub(interceptedChannel).withCallCredentials(callCredentials);
     }
 
     /**

diff --git a/java/src/main/java/utility/XClientTraceIdClientInterceptor.java b/java/src/main/java/utility/XClientTraceIdClientInterceptor.java
@@ -0,0 +1,35 @@
+package utility;
+
+import java.util.UUID;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.grpc.*;
+
+public class XClientTraceIdClientInterceptor implements ClientInterceptor {
+    private static final Logger logger = LoggerFactory.getLogger(XClientTraceIdClientInterceptor.class.getClass());
+    private static final Metadata.Key<String> X_CLIENT_TRACE_ID = Metadata.Key.of("x-client-trace-id", Metadata.ASCII_STRING_MARSHALLER);
+
+    @Override
+    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
+                                                               CallOptions callOptions, Channel next) {
+        return new ForwardingClientCall.SimpleForwardingClientCall<>(next.newCall(method, callOptions)) {
+
+            @Override
+            public void start(Listener<RespT> responseListener, Metadata headers) {
+                String xClientTraceId = UUID.randomUUID().toString();
+                headers.put(X_CLIENT_TRACE_ID, xClientTraceId);
+                logger.info("sending request for xClientTraceId {}", xClientTraceId);
+
+                super.start(new ForwardingClientCallListener.SimpleForwardingClientCallListener<>(responseListener) {
+                    @Override
+                    public void onClose(Status status, Metadata trailers) {
+                        logger.info("request completed for xClientTraceId {} with status {}", xClientTraceId, status);
+                        super.onClose(status, trailers);
+                    }
+                }, headers);
+            }
+        };
+    }
+}