Java Code Examples for io.grpc.Context#current()

@Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(final ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
    Header headerObject;
    try {
        headerObject = headerReader.extract(headers);
    } catch (Exception e) {
        if (logger.isInfoEnabled()) {
            logger.info("Header extract fail cause={}, method={} headers={}, attr={}",
                    e.getMessage(), call.getMethodDescriptor().getFullMethodName(), headers, call.getAttributes(), e);
        }
        call.close(Status.INVALID_ARGUMENT.withDescription(e.getMessage()), new Metadata());
        return new ServerCall.Listener<ReqT>() {
        };
    }

    final Context currentContext = Context.current();
    final Context newContext = currentContext.withValue(contextKey, headerObject);
    if (logger.isDebugEnabled()) {
        logger.debug("headerPropagation method={}, headers={}, attr={}", call.getMethodDescriptor().getFullMethodName(), headers, call.getAttributes());
    }

    ServerCall.Listener<ReqT> contextPropagateInterceptor = Contexts.interceptCall(newContext, call, headers, next);
    return contextPropagateInterceptor;
}
 

private ListenableFuture<Void> downloadBlob(
    Digest digest,
    OutputStream out,
    @Nullable Supplier<HashCode> hashSupplier) {
  Context ctx = Context.current();
  AtomicLong offset = new AtomicLong(0);
  ProgressiveBackoff progressiveBackoff = new ProgressiveBackoff(retrier::newBackoff);
  return Futures.catchingAsync(
      retrier.executeAsync(
          () ->
              ctx.call(() -> requestRead(offset, progressiveBackoff, digest, out, hashSupplier)),
          progressiveBackoff),
      StatusRuntimeException.class,
      (e) -> Futures.immediateFailedFuture(new IOException(e)),
      MoreExecutors.directExecutor());
}
 

@Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(final ServerCall<ReqT, RespT> serverCall, Metadata headers, ServerCallHandler<ReqT, RespT> serverCallHandler) {
    final Attributes attributes = serverCall.getAttributes();

    final TransportMetadata transportMetadata = attributes.get(MetadataServerTransportFilter.TRANSPORT_METADATA_KEY);
    if (transportMetadata == null) {
        if (logger.isInfoEnabled()) {
            logger.info("Close call. cause=transportMetadata is null, headers={}, attributes={}", headers, serverCall.getAttributes());
        }
        serverCall.close(Status.INTERNAL.withDescription("transportMetadata is null"), new Metadata());
        return new ServerCall.Listener<ReqT>() {
        };
    }

    final Context currentContext = Context.current();
    final Context newContext = currentContext.withValue(ServerContext.getTransportMetadataKey(), transportMetadata);
    if (logger.isDebugEnabled()) {
        logger.debug("bind metadata method={}, headers={}, attr={}", serverCall.getMethodDescriptor().getFullMethodName(), headers, serverCall.getAttributes());
    }
    ServerCall.Listener<ReqT> listener = Contexts.interceptCall(newContext, serverCall, headers, serverCallHandler);
    return listener;
}
 

@Test
public void testInject() {
  CustomHttpTextFormat propagator1 = new CustomHttpTextFormat("prop1");
  CustomHttpTextFormat propagator2 = new CustomHttpTextFormat("prop2");
  ContextPropagators propagators =
      DefaultContextPropagators.builder()
          .addHttpTextFormat(propagator1)
          .addHttpTextFormat(propagator2)
          .build();

  Context context = Context.current();
  context = context.withValue(propagator1.getKey(), "value1");
  context = context.withValue(propagator2.getKey(), "value2");

  Map<String, String> map = new HashMap<>();
  propagators.getHttpTextFormat().inject(context, map, MapSetter.INSTANCE);
  assertThat(map.get(propagator1.getKeyName())).isEqualTo("value1");
  assertThat(map.get(propagator2.getKeyName())).isEqualTo("value2");
}
 

@Test
public void withScopedContext() {
  Context prevCtx = Context.current();
  Context ctx = Context.current().withValue(SIMPLE_KEY, "value1");
  try (Scope scope = ContextUtils.withScopedContext(ctx)) {
    assertThat(scope).isNotNull();
    assertThat(Context.current()).isEqualTo(ctx);

    Context ctx2 = Context.current().withValue(SIMPLE_KEY, "value2");
    try (Scope scope2 = ContextUtils.withScopedContext(ctx2)) {
      assertThat(scope2).isNotNull();
      assertThat(Context.current()).isEqualTo(ctx2);
    }
    assertThat(Context.current()).isEqualTo(ctx);
  }
  assertThat(Context.current()).isEqualTo(prevCtx);
}
 

ListenableFuture<Void> start() {
  Context ctx = Context.current();
  ProgressiveBackoff progressiveBackoff = new ProgressiveBackoff(retrier::newBackoff);
  AtomicLong committedOffset = new AtomicLong(0);
  return Futures.transformAsync(
      retrier.executeAsync(
          () -> {
            if (committedOffset.get() < chunker.getSize()) {
              return ctx.call(() -> callAndQueryOnFailure(committedOffset, progressiveBackoff));
            }
            return Futures.immediateFuture(null);
          },
          progressiveBackoff),
      (result) -> {
        long committedSize = committedOffset.get();
        long expected = chunker.getSize();
        if (committedSize != expected) {
          String message =
              format(
                  "write incomplete: committed_size %d for %d total", committedSize, expected);
          return Futures.immediateFailedFuture(new IOException(message));
        }
        return Futures.immediateFuture(null);
      },
      MoreExecutors.directExecutor());
}
 

ClientCallImpl(
    MethodDescriptor<ReqT, RespT> method, Executor executor, CallOptions callOptions,
    ClientTransportProvider clientTransportProvider,
    ScheduledExecutorService deadlineCancellationExecutor,
    CallTracer channelCallsTracer,
    boolean retryEnabled) {
  this.method = method;
  // If we know that the executor is a direct executor, we don't need to wrap it with a
  // SerializingExecutor. This is purely for performance reasons.
  // See https://github.com/grpc/grpc-java/issues/368
  this.callExecutor = executor == directExecutor()
      ? new SerializeReentrantCallsDirectExecutor()
      : new SerializingExecutor(executor);
  this.channelCallsTracer = channelCallsTracer;
  // Propagate the context from the thread which initiated the call to all callbacks.
  this.context = Context.current();
  this.unaryRequest = method.getType() == MethodType.UNARY
      || method.getType() == MethodType.SERVER_STREAMING;
  this.callOptions = callOptions;
  this.clientTransportProvider = clientTransportProvider;
  this.deadlineCancellationExecutor = deadlineCancellationExecutor;
  this.retryEnabled = retryEnabled;
}
 

@Override
public ListenableFuture<ActionResult> downloadActionResult(
    ActionKey actionKey, boolean inlineOutErr) {
  GetActionResultRequest request =
      GetActionResultRequest.newBuilder()
          .setInstanceName(options.remoteInstanceName)
          .setActionDigest(actionKey.getDigest())
          .setInlineStderr(inlineOutErr)
          .setInlineStdout(inlineOutErr)
          .build();
  Context ctx = Context.current();
  return retrier.executeAsync(
      () -> ctx.call(() -> handleStatus(acFutureStub().getActionResult(request))));
}
 

MetadataApplierImpl(
    ClientTransport transport, MethodDescriptor<?, ?> method, Metadata origHeaders,
    CallOptions callOptions) {
  this.transport = transport;
  this.method = method;
  this.origHeaders = origHeaders;
  this.callOptions = callOptions;
  this.ctx = Context.current();
}
 

@Test
public void shouldAttachAndDetach() {
  final Context original = Context.current();
  assertThat(TEST_KEY.get(), is(nullValue()));
  final Context context = Context.current().withValue(TEST_KEY, "foobar");
  try (final TestScope ts = new TestScope(context)) {
    assertThat(Context.current(), is(context));
    assertThat(TEST_KEY.get(), is("foobar"));
  }
  assertThat(Context.current(), is(original));
  assertThat(TEST_KEY.get(), is(nullValue()));
}
 

public static SampleContext serverResolve() {
    if (Context.current() == Context.ROOT) {
        return CONTEXT_UNDEFINED;
    }

    SampleContext context = CALLER_ID_CONTEXT_KEY.get();
    return context == null ? CONTEXT_UNDEFINED : context;
}
 

@Before
public void setup() {
  MockitoAnnotations.initMocks(this);
  repoExecutor =
      new RemoteRepositoryRemoteExecutor(
          remoteCache,
          remoteExecutor,
          DIGEST_UTIL,
          Context.current(),
          /* remoteInstanceName= */ "foo",
          /* acceptCached= */ true);
}
 

@Test
public void noopPropagator() {
  ContextPropagators propagators = DefaultContextPropagators.builder().build();

  Context context = Context.current();
  Map<String, String> map = new HashMap<>();
  propagators.getHttpTextFormat().inject(context, map, MapSetter.INSTANCE);
  assertThat(map).isEmpty();

  assertThat(propagators.getHttpTextFormat().extract(context, map, MapGetter.INSTANCE))
      .isSameInstanceAs(context);
}
 

@Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(final ServerCall<ReqT, RespT> call,
                                                             final Metadata headers,
                                                             final ServerCallHandler<ReqT, RespT> next) {
    Context ctx = Context.current();
    final ServerStream stream = ServerStreamHelper.getServerStream(call);
    if (stream != null) {
        ctx = ctx.withValue(STREAM, stream);
    }
    return Contexts.interceptCall(ctx, call, headers, next);
}
 

MetadataApplierImpl(
    ClientTransport transport, MethodDescriptor<?, ?> method, Metadata origHeaders,
    CallOptions callOptions) {
  this.transport = transport;
  this.method = method;
  this.origHeaders = origHeaders;
  this.callOptions = callOptions;
  this.ctx = Context.current();
}
 

@Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {

    Context context = Context.current();

    // The authorization header has the credentials (e.g., username and password for Basic Authentication).
    // The form of the header is: <Method> <Token> (CustomMethod static-token, or Basic XYZ...., for example)
    String credentials = headers.get(Metadata.Key.of(AuthConstants.AUTHORIZATION, Metadata.ASCII_STRING_MARSHALLER));

    if (!Strings.isNullOrEmpty(credentials)) {
        String[] parts = credentials.split("\\s+", 2);
        if (parts.length == 2) {
            String method = parts[0];
            String token = parts[1];
            if (!Strings.isNullOrEmpty(method)) {
                if (method.equals(handler.getHandlerName())) {
                    log.debug("Handler [{}] successfully matched auth method [{}]", handler, method);
                    Principal principal;
                    try {
                        if ((principal = handler.authenticate(token)) == null) {
                            log.warn("Handler for method [{}] returned a null Principal upon authentication for the"
                                    + "given token", method);
                            call.close(Status.fromCode(Status.Code.UNAUTHENTICATED), headers);
                            return null;
                        }
                    } catch (AuthException e) {
                        log.warn("Authentication failed", e);
                        call.close(Status.fromCode(Status.Code.UNAUTHENTICATED), headers);
                        return null;
                    }
                    // Creates a new Context with the given key/value pairs.
                    context = context.withValues(PRINCIPAL_OBJECT_KEY, principal, AUTH_INTERCEPTOR_OBJECT_KEY, this);
                }
            } else {
                log.debug("Credentials are present, but method [{}] is null or empty", method);
            }
        }
    }

    // reaching this point means that the handler wasn't applicable to this request.
    return Contexts.interceptCall(context, call, headers, next);
}
 

@Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
    Context wrappedContext = Context.current();

    Object debugValue = headers.get(DEBUG_KEY);
    if (debugValue != null) {
        boolean debugEnabled = Boolean.parseBoolean(debugValue.toString());
        if (debugEnabled) {
            wrappedContext = wrappedContext.withValue(DEBUG_CONTEXT_KEY, "true");
        }
    }
    Object compressionValue = headers.get(COMPRESSION_KEY);
    if (compressionValue != null) {
        String compressionType = compressionValue.toString();
        if (ALLOWED_COMPRESSION_TYPES.contains(compressionType)) {
            call.setCompression(compressionType);
            wrappedContext = wrappedContext.withValue(COMPRESSION_CONTEXT_KEY, compressionType);
        }
    }

    wrappedContext = copyIntoContext(wrappedContext, headers, CALLER_ID_KEY, CALLER_ID_CONTEXT_KEY);
    wrappedContext = copyIntoContext(wrappedContext, headers, CALLER_TYPE_KEY, CALLER_TYPE_CONTEXT_KEY);
    wrappedContext = copyIntoContext(wrappedContext, headers, DIRECT_CALLER_ID_KEY, DIRECT_CALLER_ID_CONTEXT_KEY);
    wrappedContext = copyIntoContext(wrappedContext, headers, CALL_REASON_KEY, CALL_REASON_CONTEXT_KEY);
    wrappedContext = copyDirectCallerContextIntoContext(wrappedContext, call);

    Object callMetadataValue = headers.get(CALL_METADATA_KEY);
    if (callMetadataValue != null) {
        try {
            com.netflix.titus.grpc.protogen.CallMetadata grpcCallMetadata = com.netflix.titus.grpc.protogen.CallMetadata.parseFrom((byte[]) callMetadataValue);
            wrappedContext = wrappedContext.withValue(CALL_METADATA_CONTEXT_KEY, CommonRuntimeGrpcModelConverters.toCallMetadata(grpcCallMetadata));
        } catch (Exception e) {
            // Ignore bad header value.
            logger.info("Invalid CallMetadata in a request header", e);
        }
    }

    return wrappedContext == Context.current()
            ? next.startCall(call, headers)
            : Contexts.interceptCall(wrappedContext, call, headers, next);
}
 

private void attachContext(TransportMetadata transportMetadata) {
    final Context currentContext = Context.current();
    Context newContext = currentContext.withValue(ServerContext.getTransportMetadataKey(), transportMetadata);
    newContext.attach();
}
 

@Override
public TraceContext getCurrentContext() {
  return new OpenTelemetryContext(tracer, Context.current());
}
 

public static Header getAgentInfo() {
    final Context current = Context.current();
    return getAgentInfo(current);
}