6
6
import io .grpc .ServerCall ;
7
7
import io .grpc .ServerCallHandler ;
8
8
import io .grpc .ServerInterceptor ;
9
+ import java .util .Optional ;
10
+ import java .util .UUID ;
9
11
import org .hypertrace .core .grpcutils .context .RequestContext ;
12
+ import org .hypertrace .core .grpcutils .context .RequestContextConstants ;
10
13
11
14
/**
12
15
* 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(
20
23
ServerCall <ReqT , RespT > serverCall ,
21
24
Metadata metadata ,
22
25
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
+ }
25
33
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 );
27
39
}
28
40
}
0 commit comments