Spring Boot GRPC: ServerIntereceptor to read data in the request, and set it in the response - multithreading

Question

There is a field called "metadata" (not to be confused with GRPC metadata) that is present in every request proto that comes to the GRPC service:
message MyRequest {
RequestResponseMetadata metadata = 1;
...
}
And the same field is also present in all responses:
message MyResponse {
RequestResponseMetadata metadata = 1;
...
}
I am trying to write a ServerInterceptor (or something else, if it works) to read the "metadata" field from the request, keep it somewhere, and then set it in the response once done processing the request.
Attempt 1: ThreadLocal
public class ServerInterceptor implements io.grpc.ServerInterceptor {
private ThreadLocal<RequestResponseMetadata> metadataThreadLocal = new ThreadLocal<>();
#Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
ServerCall<ReqT, RespT> call,
final Metadata requestHeaders,
ServerCallHandler<ReqT, RespT> next) {
return new SimpleForwardingServerCallListener<ReqT>(
next.startCall(
new SimpleForwardingServerCall<ReqT, RespT>(call) {
#Override
public void sendMessage(RespT message) {
super.sendMessage(
(RespT)
MetadataUtils.setMetadata(
(GeneratedMessageV3) message, metadataThreadLocal.get()));
metadataThreadLocal.remove();
}
},
requestHeaders)) {
#Override
public void onMessage(ReqT request) {
// todo nava see if ReqT can extend GenericV3Message
var metadata = MetadataUtils.getMetadata((GeneratedMessageV3) request);
metadataThreadLocal.set(metadata);
super.onMessage(request);
}
};
}
}
I tried to use ThreadLocal, to later realise that sendMessage and onMessage need not necessary to be on the same thread.
Attempt 2: GRPC Context
public class ServerInterceptor implements io.grpc.ServerInterceptor {
public static final Context.Key<RequestResponseMetadata> METADATA_KEY = Context.key("metadata");
#Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
ServerCall<ReqT, RespT> call,
final Metadata requestHeaders,
ServerCallHandler<ReqT, RespT> next) {
return new SimpleForwardingServerCallListener<ReqT>(
next.startCall(
new SimpleForwardingServerCall<ReqT, RespT>(call) {
#Override
public void sendMessage(RespT message) {
super.sendMessage(
(RespT)
MetadataUtils.setMetadata(
(GeneratedMessageV3) message, METADATA_KEY.get()));
}
},
requestHeaders)) {
#Override
public void onMessage(ReqT request) {
var metadata = MetadataUtils.getMetadata((GeneratedMessageV3) request);
var newContext = Context.current().withValue(METADATA_KEY, metadata);
oldContext = newContext.attach();
super.onMessage(request);
}
};
}
}
I am planning to detach the context in a onComplete(), but before it reaches there itself, METADATA_KEY.get() in sendMessage returns null, while I was expecting it to return the data.
Even before hitting the sendMessage() function, I get this in the console, indicating that I am doing something wrong:
3289640 [grpc-default-executor-0] ERROR i.g.ThreadLocalContextStorage - Context was not attached when detaching
java.lang.Throwable: null
at io.grpc.ThreadLocalContextStorage.detach(ThreadLocalContextStorage.java:48)
at io.grpc.Context.detach(Context.java:421)
at io.grpc.Context$CancellableContext.detach(Context.java:761)
at io.grpc.internal.ContextRunnable.run(ContextRunnable.java:39)
at io.grpc.internal.SerializingExecutor.run(SerializingExecutor.java:123)
at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1128)
at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:628)
at java.base/java.lang.Thread.run(Thread.java:834)
How do I read data when a request is received, store it somewhere and use it when the response is send back?

Answer 1

You can use Metadata to pass values from the request to the response:
public class MetadataServerInterceptor implements ServerInterceptor {
public static final Metadata.Key<byte[]> METADATA_KEY = Metadata.Key.of("metadata-bin", Metadata.BINARY_BYTE_MARSHALLER);
#Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
var serverCall = new ForwardingServerCall.SimpleForwardingServerCall<>(call) {
#Override
public void sendMessage(RespT message) {
byte[] metadata = headers.get(METADATA_KEY);
message = (RespT) MetadataUtils.setMetadata((GeneratedMessageV3) message, metadata);
super.sendMessage(message);
}
};
ServerCall.Listener<ReqT> listenerWithContext = Contexts.interceptCall(Context.current(), serverCall, headers, next);
return new ForwardingServerCallListener.SimpleForwardingServerCallListener<>(listenerWithContext) {
#Override
public void onMessage(ReqT message) {
byte[] metadata = MetadataUtils.getMetadata((GeneratedMessageV3) message);
headers.put(METADATA_KEY, metadata);
super.onMessage(message);
}
};
}
}
Note: Since it is not possible to put the instance of RequestResponseMetadata in the metadata (at least without implementing a custom marshaller), you can save it there as a byte array. You can use toByteArray() on your RequestResponseMetadata object to get byte[] and RequestResponseMetadata.#parseFrom(byte[]) to get the object from byte[].