core: DelayedClientCall should propagate context

This was already being done for the Listener, but it was missed for the ClientCall draining itself. That's not too surprising, since very few things should be looking at the context in that path. We don't care too much about this client call case, but if the context _does_ end up mattering it could be painful to debug. Fixes grpc#9478

Author: ejona86

core/src/main/java/io/grpc/internal/DelayedClientCall.java

@@ -152,9 +152,9 @@ public final Runnable setCall(ClientCall<ReqT, RespT> call) {
  }
  setRealCall(checkNotNull(call, "call"));
  }
- return new Runnable() {
+ return new ContextRunnable(context) {
  @Override
- public void run() {
+ public void runInContext() {
  drainPendingCalls();
  }
  };

core/src/test/java/io/grpc/internal/DelayedClientCallTest.java

@@ -27,7 +27,9 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import io.grpc.ClientCall;
 import io.grpc.ClientCall.Listener;
+import io.grpc.Context;
 import io.grpc.Deadline;
+import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
 import io.grpc.ForwardingTestUtil;
 import io.grpc.Metadata;
 import io.grpc.Status;
@@ -36,6 +38,7 @@
 import java.lang.reflect.Modifier;
 import java.util.Arrays;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -191,6 +194,35 @@ public void setCallThenCancel() {
  verify(listener).onClose(Status.CANCELLED, null);
  }
 
+ @Test
+ public void delayedCallsRunUnderContext() throws Exception {
+ Context.Key<Object> contextKey = Context.key("foo");
+ Object goldenValue = new Object();
+ DelayedClientCall<String, Integer> delayedClientCall =
+ Context.current().withValue(contextKey, goldenValue).call(() ->
+ new DelayedClientCall<>(callExecutor, fakeClock.getScheduledExecutorService(), null));
+ AtomicReference<Context> readyContext = new AtomicReference<>();
+ delayedClientCall.start(new ClientCall.Listener<Integer>() {
+ @Override public void onReady() {
+ readyContext.set(Context.current());
+ }
+ }, new Metadata());
+ AtomicReference<Context> startContext = new AtomicReference<>();
+ Runnable r = delayedClientCall.setCall(new SimpleForwardingClientCall<String, Integer>(
+ mockRealCall) {
+ @Override public void start(Listener<Integer> listener, Metadata metadata) {
+ startContext.set(Context.current());
+ listener.onReady(); // Delayed until call finishes draining
+ assertThat(readyContext.get()).isNull();
+ super.start(listener, metadata);
+ }
+ });
+ assertThat(r).isNotNull();
+ r.run();
+ assertThat(contextKey.get(startContext.get())).isEqualTo(goldenValue);
+ assertThat(contextKey.get(readyContext.get())).isEqualTo(goldenValue);
+ }
+
  private void callMeMaybe(Runnable r) {
  if (r != null) {
  r.run();