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

@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;
}
 

@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");
}
 

@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;
}
 

private <ReqT, RespT> Context copyDirectCallerContextIntoContext(Context context, ServerCall<ReqT, RespT> call) {
    Map<String, String> callerContext = new HashMap<>();

    String fullName = call.getMethodDescriptor().getFullMethodName();
    int methodBegin = fullName.indexOf('/');
    String serviceName;
    String methodName;
    if (methodBegin <= 0) {
        serviceName = fullName;
        methodName = fullName;
    } else {
        serviceName = fullName.substring(0, methodBegin);
        methodName = Character.toLowerCase(fullName.charAt(methodBegin + 1)) + fullName.substring(methodBegin + 2);
    }

    callerContext.put(CallMetadataHeaders.DIRECT_CALLER_CONTEXT_SERVICE_NAME, serviceName);
    callerContext.put(CallMetadataHeaders.DIRECT_CALLER_CONTEXT_SERVICE_METHOD, methodName);
    callerContext.put(CallMetadataHeaders.DIRECT_CALLER_CONTEXT_TRANSPORT_TYPE, "GRPC");
    callerContext.put(CallMetadataHeaders.DIRECT_CALLER_CONTEXT_TRANSPORT_SECURE, "?");

    return context.withValue(CALLER_CONTEXT_CONTEXT_KEY, callerContext);
}
 

@Override
public Context filterContext(Context context) {
  // Access directly the unsafe trace API to create the new Context. This is a safe usage
  // because gRPC always creates a new Context for each of the server calls and does not
  // inherit from the parent Context.
  return context.withValue(CONTEXT_SPAN_KEY, span);
}
 

@Override
public Context filterContext(Context context) {
  if (!tagger.empty().equals(parentCtx)) {
    return context.withValue(TAG_CONTEXT_KEY, parentCtx);
  }
  return 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);
}
 

@Override
public <ReqT, RespT> Listener<ReqT> interceptCall(
    ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
  Context ctx = Context.current();
  CallMetricRecorder callMetricRecorder = InternalCallMetricRecorder.CONTEXT_KEY.get(ctx);
  if (callMetricRecorder == null) {
    callMetricRecorder = InternalCallMetricRecorder.newCallMetricRecorder();
    ctx = ctx.withValue(InternalCallMetricRecorder.CONTEXT_KEY, callMetricRecorder);
  }
  final CallMetricRecorder finalCallMetricRecorder = callMetricRecorder;
  ServerCall<ReqT, RespT> trailerAttachingCall =
      new SimpleForwardingServerCall<ReqT, RespT>(call) {
        @Override
        public void close(Status status, Metadata trailers) {
          Map<String, Double> metricValues =
              InternalCallMetricRecorder.finalizeAndDump(finalCallMetricRecorder);
          // Only attach a metric report if there are some metric values to be reported.
          if (!metricValues.isEmpty()) {
            OrcaLoadReport report =
                OrcaLoadReport.newBuilder().putAllRequestCost(metricValues).build();
            trailers.put(ORCA_ENDPOINT_LOAD_METRICS_KEY, report);
          }
          super.close(status, trailers);
        }
      };
  return Contexts.interceptCall(
      ctx,
      trailerAttachingCall,
      headers,
      next);
}
 

@Override
public <C> Context extract(Context context, C carrier, Getter<C> getter) {
  String payload = getter.get(carrier, name);
  if (payload != null) {
    context = context.withValue(key, payload);
  }

  return context;
}
 

/**
 * Attaches an empty ambient context to the provided gRPC {@code Context}.
 *
 * @throws IllegalStateException  if an ambient context has already been attached to the
 * provided gRPC {@code Context}.
 */
public static Context initialize(Context context) {
    checkNotNull(context, "context");
    checkState(DATA_KEY.get(context) == null,
            "AmbientContext has already been created in the scope of the current context");
    return context.withValue(DATA_KEY, new AmbientContext());
}
 

@Override
public Context filterContext(Context context) {
  Context newCtx = super.filterContext(context);
  return newCtx.withValue(SERVER_TRACER_ADDED_KEY, "context added by tracer");
}
 

@Test
public void shareCallMetricRecorderInContext() throws IOException {
  final CallMetricRecorder callMetricRecorder =
      InternalCallMetricRecorder.newCallMetricRecorder();
  ServerStreamTracer.Factory callMetricRecorderSharingStreamTracerFactory =
      new ServerStreamTracer.Factory() {
    @Override
    public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata headers) {
      return new ServerStreamTracer() {
        @Override
        public Context filterContext(Context context) {
          return context.withValue(InternalCallMetricRecorder.CONTEXT_KEY, callMetricRecorder);
        }
      };
    }
  };

  final AtomicReference<CallMetricRecorder> callMetricRecorderCapture = new AtomicReference<>();
  SimpleServiceGrpc.SimpleServiceImplBase simpleServiceImpl =
      new SimpleServiceGrpc.SimpleServiceImplBase() {
        @Override
        public void unaryRpc(
            SimpleRequest request, StreamObserver<SimpleResponse> responseObserver) {
          callMetricRecorderCapture.set(CallMetricRecorder.getCurrent());
          SimpleResponse response =
              SimpleResponse.newBuilder().setResponseMessage("Simple response").build();
          responseObserver.onNext(response);
          responseObserver.onCompleted();
        }
      };

  ServerInterceptor metricReportingServerInterceptor = new OrcaMetricReportingServerInterceptor();
  String serverName = InProcessServerBuilder.generateName();
  grpcCleanupRule.register(
      InProcessServerBuilder
          .forName(serverName)
          .directExecutor()
          .addStreamTracerFactory(callMetricRecorderSharingStreamTracerFactory)
          .addService(
              ServerInterceptors.intercept(simpleServiceImpl, metricReportingServerInterceptor))
          .build().start());

  ManagedChannel channel =
      grpcCleanupRule.register(InProcessChannelBuilder.forName(serverName).build());
  ClientCalls.blockingUnaryCall(channel, SIMPLE_METHOD, CallOptions.DEFAULT, REQUEST);

  assertThat(callMetricRecorderCapture.get()).isSameInstanceAs(callMetricRecorder);
}
 

private void attachContext(Header header) {
    final Context currentContext = Context.current();
    Context newContext = currentContext.withValue(ServerContext.getAgentInfoKey(), header);
    newContext.attach();
}
 

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

private static Context copyIntoContext(Context context, Metadata headers, Metadata.Key<String> headerKey, Context.Key<String> contextKey) {
    Object value = headers.get(headerKey);
    return value == null ? context : context.withValue(contextKey, value.toString());
}
 

@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);
}
 

@Override
public Context filterContext(Context context) {
  Context newCtx = super.filterContext(context);
  return newCtx.withValue(SERVER_TRACER_ADDED_KEY, "context added by tracer");
}
 

/**
 * Similar to {@link #initialize(Context)}, {@code fork()} attaches a shallow clone of this {@code AmbientContext}
 * to a provided gRPC {@code Context}. Use {@code fork()} when you want create a temporary context scope.
 *
 * @param context
 * @return
 */
public Context fork(Context context) {
    return context.withValue(DATA_KEY, new AmbientContext(this));
}
 

/**
 * Creates a new {@code Context} with the given value set.
 *
 * @param corrContext the value to be set.
 * @param context the parent {@code Context}.
 * @return a new context with the given value set.
 * @since 0.1.0
 */
public static Context withCorrelationContext(CorrelationContext corrContext, Context context) {
  return context.withValue(CORR_CONTEXT_KEY, corrContext);
}
 

/**
 * Creates a new {@code Context} with the given {@link Span} set.
 *
 * @param span the value to be set.
 * @param context the parent {@code Context}.
 * @return a new context with the given value set.
 * @since 0.1.0
 */
public static Context withSpan(Span span, Context context) {
  return context.withValue(CONTEXT_SPAN_KEY, span);
}