66import  io .grpc .ServerCall ;
77import  io .grpc .ServerCallHandler ;
88import  io .grpc .ServerInterceptor ;
9+ import  java .util .Optional ;
10+ import  java .util .UUID ;
911import  org .hypertrace .core .grpcutils .context .RequestContext ;
12+ import  org .hypertrace .core .grpcutils .context .RequestContextConstants ;
1013
1114/** 
1215 * Interceptor which intercepts the request headers to extract request context and sets it in the 
@@ -20,9 +23,18 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
2023      ServerCall <ReqT , RespT > serverCall ,
2124      Metadata  metadata ,
2225      ServerCallHandler <ReqT , RespT > serverCallHandler ) {
23-     RequestContext  requestContext  = RequestContext .fromMetadata (metadata );
24-     Context  ctx  = Context .current ().withValue (RequestContext .CURRENT , requestContext );
26+     RequestContext  currentContext  =
27+         Optional .ofNullable (RequestContext .CURRENT .get ())
28+             .orElseGet (() -> RequestContext .fromMetadata (metadata ));
29+     if  (currentContext .getHeaderValue (RequestContextConstants .REQUEST_ID_HEADER_KEY ).isEmpty ()) {
30+       currentContext .put (
31+           RequestContextConstants .REQUEST_ID_HEADER_KEY , UUID .randomUUID ().toString ());
32+     }
2533
26-     return  Contexts .interceptCall (ctx , serverCall , metadata , serverCallHandler );
34+     return  Contexts .interceptCall (
35+         Context .current ().withValue (RequestContext .CURRENT , currentContext ),
36+         serverCall ,
37+         metadata ,
38+         serverCallHandler );
2739  }
2840}
0 commit comments