io.grpc.ClientStreamTracer Java Examples

@Test
public void loadRecordingStreamTracerFactory_clientSideQueryCountsAggregation() {
  LoadRecordingStreamTracerFactory factory1 =
      new LoadRecordingStreamTracerFactory(counter, NOOP_CLIENT_STREAM_TRACER_FACTORY);
  ClientStreamTracer tracer = factory1.newClientStreamTracer(STREAM_INFO, new Metadata());
  ClientLoadSnapshot snapshot = counter.snapshot();
  assertQueryCounts(snapshot, 0, 1, 0, 1);
  tracer.streamClosed(Status.OK);
  snapshot = counter.snapshot();
  assertQueryCounts(snapshot, 1, 0, 0, 0);

  // Create a second LoadRecordingStreamTracerFactory with the same counter, stats are aggregated
  // together.
  LoadRecordingStreamTracerFactory factory2 =
      new LoadRecordingStreamTracerFactory(counter, NOOP_CLIENT_STREAM_TRACER_FACTORY);
  factory1.newClientStreamTracer(STREAM_INFO, new Metadata()).streamClosed(Status.ABORTED);
  factory2.newClientStreamTracer(STREAM_INFO, new Metadata()).streamClosed(Status.CANCELLED);
  snapshot = counter.snapshot();
  assertQueryCounts(snapshot, 0, 0, 2, 2);
}
 

@Override
public ClientStreamTracer newClientStreamTracer(StreamInfo info, Metadata headers) {
  counter.recordCallStarted();
  final ClientStreamTracer delegateTracer = delegate.newClientStreamTracer(info, headers);
  return new ForwardingClientStreamTracer() {
    @Override
    protected ClientStreamTracer delegate() {
      return delegateTracer;
    }

    @Override
    public void streamClosed(Status status) {
      counter.recordCallFinished(status);
      delegate().streamClosed(status);
    }
  };
}
 

@Test
public void noToken() {
  TokenAttachingTracerFactory factory = new TokenAttachingTracerFactory(delegate);
  ClientStreamTracer.StreamInfo info = ClientStreamTracer.StreamInfo.newBuilder()
      .setTransportAttrs(
          Attributes.newBuilder()
              .set(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS, Attributes.EMPTY).build())
      .build();

  Metadata headers = new Metadata();
  // Preexisting token should be removed
  headers.put(GrpclbConstants.TOKEN_METADATA_KEY, "preexisting-token");

  ClientStreamTracer tracer = factory.newClientStreamTracer(info, headers);
  verify(delegate).newClientStreamTracer(same(info), same(headers));
  assertThat(tracer).isSameInstanceAs(fakeTracer);
  assertThat(headers.get(GrpclbConstants.TOKEN_METADATA_KEY)).isNull();
}
 

@Test
public void hasToken() {
  TokenAttachingTracerFactory factory = new TokenAttachingTracerFactory(delegate);
  Attributes eagAttrs = Attributes.newBuilder()
      .set(GrpclbConstants.TOKEN_ATTRIBUTE_KEY, "token0001").build();
  ClientStreamTracer.StreamInfo info = ClientStreamTracer.StreamInfo.newBuilder()
      .setTransportAttrs(
          Attributes.newBuilder().set(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS, eagAttrs).build())
      .build();
  Metadata headers = new Metadata();
  // Preexisting token should be replaced
  headers.put(GrpclbConstants.TOKEN_METADATA_KEY, "preexisting-token");

  ClientStreamTracer tracer = factory.newClientStreamTracer(info, headers);
  verify(delegate).newClientStreamTracer(same(info), same(headers));
  assertThat(tracer).isSameInstanceAs(fakeTracer);
  assertThat(headers.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("token0001");
}
 

private Substream createSubstream(int previousAttemptCount) {
  Substream sub = new Substream(previousAttemptCount);
  // one tracer per substream
  final ClientStreamTracer bufferSizeTracer = new BufferSizeTracer(sub);
  ClientStreamTracer.Factory tracerFactory = new ClientStreamTracer.Factory() {
    @Override
    public ClientStreamTracer newClientStreamTracer(
        ClientStreamTracer.StreamInfo info, Metadata headers) {
      return bufferSizeTracer;
    }
  };

  Metadata newHeaders = updateHeaders(headers, previousAttemptCount);
  // NOTICE: This set _must_ be done before stream.start() and it actually is.
  sub.stream = newSubstream(tracerFactory, newHeaders);
  return sub;
}
 

@Override
public ClientStreamTracer newClientStreamTracer(
    ClientStreamTracer.StreamInfo info, Metadata headers) {
  Attributes transportAttrs = checkNotNull(info.getTransportAttrs(), "transportAttrs");
  Attributes eagAttrs =
      checkNotNull(transportAttrs.get(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS), "eagAttrs");
  String token = eagAttrs.get(GrpclbConstants.TOKEN_ATTRIBUTE_KEY);
  headers.discardAll(GrpclbConstants.TOKEN_METADATA_KEY);
  if (token != null) {
    headers.put(GrpclbConstants.TOKEN_METADATA_KEY, token);
  }
  if (delegate != null) {
    return delegate.newClientStreamTracer(info, headers);
  } else {
    return NOOP_TRACER;
  }
}
 

/**
 * Factory method for the client-side.
 */
public static StatsTraceContext newClientContext(
    final CallOptions callOptions, final Attributes transportAttrs, Metadata headers) {
  List<ClientStreamTracer.Factory> factories = callOptions.getStreamTracerFactories();
  if (factories.isEmpty()) {
    return NOOP;
  }
  ClientStreamTracer.StreamInfo info =
      ClientStreamTracer.StreamInfo.newBuilder()
          .setTransportAttrs(transportAttrs).setCallOptions(callOptions).build();
  // This array will be iterated multiple times per RPC. Use primitive array instead of Collection
  // so that for-each doesn't create an Iterator every time.
  StreamTracer[] tracers = new StreamTracer[factories.size()];
  for (int i = 0; i < tracers.length; i++) {
    tracers[i] = factories.get(i).newClientStreamTracer(info, headers);
  }
  return new StatsTraceContext(tracers);
}
 

@Test
public void subchannelPickerInterceptedWithLoadRecording() {
  List<EquivalentAddressGroup> backendAddrs = createResolvedBackendAddresses(2);
  deliverResolvedAddresses(backendAddrs, "round_robin");
  FakeLoadBalancer childBalancer = (FakeLoadBalancer) childBalancers.poll();
  NoopSubchannel subchannel = childBalancer.subchannels.values().iterator().next();
  deliverSubchannelState(subchannel, ConnectivityState.READY);
  assertThat(loadRecorder.recording).isTrue();
  ArgumentCaptor<SubchannelPicker> pickerCaptor = ArgumentCaptor.forClass(null);
  verify(helper).updateBalancingState(eq(ConnectivityState.READY), pickerCaptor.capture());
  SubchannelPicker picker = pickerCaptor.getValue();
  assertThat(picker).isInstanceOf(LoadRecordingSubchannelPicker.class);
  PickResult result = picker.pickSubchannel(mock(PickSubchannelArgs.class));
  ClientStreamTracer.Factory tracerFactory = result.getStreamTracerFactory();
  assertThat(((LoadRecordingStreamTracerFactory) tracerFactory).getCounter())
      .isSameInstanceAs(counter);
  loadBalancer.shutdown();
  assertThat(childBalancer.shutdown).isTrue();
  assertThat(loadRecorder.recording).isFalse();
}
 

private Substream createSubstream(int previousAttempts) {
  Substream sub = new Substream(previousAttempts);
  // one tracer per substream
  final ClientStreamTracer bufferSizeTracer = new BufferSizeTracer(sub);
  ClientStreamTracer.Factory tracerFactory = new ClientStreamTracer.Factory() {
    @Override
    public ClientStreamTracer newClientStreamTracer(CallOptions callOptions, Metadata headers) {
      return bufferSizeTracer;
    }
  };

  Metadata newHeaders = updateHeaders(headers, previousAttempts);
  // NOTICE: This set _must_ be done before stream.start() and it actually is.
  sub.stream = newSubstream(tracerFactory, newHeaders);
  return sub;
}
 

@Test
public void newStream_afterTermination() throws Exception {
  // We expect the same general behavior as duringShutdown, but for some transports (e.g., Netty)
  // dealing with afterTermination is harder than duringShutdown.
  server.start(serverListener);
  client = newClientTransport(server);
  startTransport(client, mockClientTransportListener);
  verify(mockClientTransportListener, timeout(TIMEOUT_MS)).transportReady();
  Status shutdownReason = Status.UNAVAILABLE.withDescription("shutdown called");
  client.shutdown(shutdownReason);
  verify(mockClientTransportListener, timeout(TIMEOUT_MS)).transportTerminated();
  Thread.sleep(100);
  ClientStream stream = client.newStream(methodDescriptor, new Metadata(), callOptions);
  ClientStreamListenerBase clientStreamListener = new ClientStreamListenerBase();
  stream.start(clientStreamListener);
  assertEquals(
      shutdownReason, clientStreamListener.status.get(TIMEOUT_MS, TimeUnit.MILLISECONDS));
  assertNotNull(clientStreamListener.trailers.get(TIMEOUT_MS, TimeUnit.MILLISECONDS));
  verify(mockClientTransportListener, never()).transportInUse(anyBoolean());
  verify(clientStreamTracerFactory).newClientStreamTracer(
      any(ClientStreamTracer.StreamInfo.class), any(Metadata.class));
  assertNull(clientStreamTracer1.getInboundTrailers());
  assertSame(shutdownReason, clientStreamTracer1.getStatus());
  // Assert no interactions
  assertNull(serverStreamTracer1.getServerCallInfo());
}
 

/**
 * Tests the case when parent policy creates its own {@link ClientStreamTracer.Factory}, ORCA
 * reports are only forwarded to the parent's listener.
 */
@Test
public void onlyParentPolicyReceivesReportsIfCreatesOwnTracer() {
  ClientStreamTracer.Factory parentFactory =
      OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(orcaListener1);
  ClientStreamTracer.Factory childFactory =
      mock(ClientStreamTracer.Factory.class,
          delegatesTo(OrcaPerRequestUtil.getInstance()
              .newOrcaClientStreamTracerFactory(parentFactory, orcaListener2)));
  ClientStreamTracer parentTracer =
      parentFactory.newClientStreamTracer(STREAM_INFO, new Metadata());
  Metadata trailer = new Metadata();
  trailer.put(
      OrcaReportingTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY,
      OrcaLoadReport.getDefaultInstance());
  parentTracer.inboundTrailers(trailer);
  verify(orcaListener1).onLoadReport(eq(OrcaLoadReport.getDefaultInstance()));
  verifyZeroInteractions(childFactory);
  verifyZeroInteractions(orcaListener2);
}
 

@Override
public ClientStreamTracer newClientStreamTracer(
    ClientStreamTracer.StreamInfo info, Metadata headers) {
  ClientTracer tracer = new ClientTracer(module, startCtx);
  // TODO(zhangkun83): Once retry or hedging is implemented, a ClientCall may start more than
  // one streams.  We will need to update this file to support them.
  if (streamTracerUpdater != null) {
    checkState(
        streamTracerUpdater.compareAndSet(this, null, tracer),
        "Are you creating multiple streams per call? This class doesn't yet support this case");
  } else {
    checkState(
        streamTracer == null,
        "Are you creating multiple streams per call? This class doesn't yet support this case");
    streamTracer = tracer;
  }
  if (module.propagateTags) {
    headers.discardAll(module.statsHeader);
    if (!module.tagger.empty().equals(parentCtx)) {
      headers.put(module.statsHeader, parentCtx);
    }
  }
  return tracer;
}
 

@Override
public ClientStreamTracer newClientStreamTracer(CallOptions callOptions, Metadata headers) {
  ClientTracer tracer = new ClientTracer();
  // TODO(zhangkun83): Once retry or hedging is implemented, a ClientCall may start more than
  // one streams.  We will need to update this file to support them.
  if (streamTracerUpdater != null) {
    checkState(
        streamTracerUpdater.compareAndSet(this, null, tracer),
        "Are you creating multiple streams per call? This class doesn't yet support this case");
  } else {
    checkState(
        streamTracer == null,
        "Are you creating multiple streams per call? This class doesn't yet support this case");
    streamTracer = tracer;
  }
  if (module.propagateTags) {
    headers.discardAll(module.statsHeader);
    if (!module.tagger.empty().equals(parentCtx)) {
      headers.put(module.statsHeader, parentCtx);
    }
  }
  return tracer;
}
 

@Test
public void allMethodsForwarded() throws Exception {
  ForwardingTestUtil.testMethodsForwarded(
      ClientStreamTracer.class,
      mockDelegate,
      new ForwardingClientStreamTracerTest.TestClientStreamTracer(),
      Collections.<Method>emptyList());
}
 

/**
 * Tests a single load balance policy's listener receive per-request ORCA reports upon call
 * trailer arrives.
 */
@Test
public void singlePolicyTypicalWorkflow() {
  // Use a mocked noop stream tracer factory as the original stream tracer factory.
  ClientStreamTracer.Factory fakeDelegateFactory = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer fakeTracer = mock(ClientStreamTracer.class);
  doNothing().when(fakeTracer).inboundTrailers(any(Metadata.class));
  when(fakeDelegateFactory.newClientStreamTracer(
          any(ClientStreamTracer.StreamInfo.class), any(Metadata.class)))
      .thenReturn(fakeTracer);

  // The OrcaReportingTracerFactory will augment the StreamInfo passed to its
  // newClientStreamTracer method. The augmented StreamInfo's CallOptions will contain
  // a OrcaReportBroker, in which has the registered listener.
  ClientStreamTracer.Factory factory =
      OrcaPerRequestUtil.getInstance()
          .newOrcaClientStreamTracerFactory(fakeDelegateFactory, orcaListener1);
  ClientStreamTracer tracer = factory.newClientStreamTracer(STREAM_INFO, new Metadata());
  ArgumentCaptor<ClientStreamTracer.StreamInfo> streamInfoCaptor = ArgumentCaptor.forClass(null);
  verify(fakeDelegateFactory)
      .newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class));
  ClientStreamTracer.StreamInfo capturedInfo = streamInfoCaptor.getValue();
  assertThat(capturedInfo).isNotEqualTo(STREAM_INFO);

  // When the trailer does not contain ORCA report, listener callback will not be invoked.
  Metadata trailer = new Metadata();
  tracer.inboundTrailers(trailer);
  verifyNoMoreInteractions(orcaListener1);

  // When the trailer contains an ORCA report, listener callback will be invoked.
  trailer.put(
      OrcaReportingTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY,
      OrcaLoadReport.getDefaultInstance());
  tracer.inboundTrailers(trailer);
  ArgumentCaptor<OrcaLoadReport> reportCaptor = ArgumentCaptor.forClass(null);
  verify(orcaListener1).onLoadReport(reportCaptor.capture());
  assertThat(reportCaptor.getValue()).isEqualTo(OrcaLoadReport.getDefaultInstance());
}
 

/**
 * Factory method for the client-side.
 */
public static StatsTraceContext newClientContext(CallOptions callOptions, Metadata headers) {
  List<ClientStreamTracer.Factory> factories = callOptions.getStreamTracerFactories();
  if (factories.isEmpty()) {
    return NOOP;
  }
  // This array will be iterated multiple times per RPC. Use primitive array instead of Collection
  // so that for-each doesn't create an Iterator every time.
  StreamTracer[] tracers = new StreamTracer[factories.size()];
  for (int i = 0; i < tracers.length; i++) {
    tracers[i] = factories.get(i).newClientStreamTracer(callOptions, headers);
  }
  return new StatsTraceContext(tracers);
}
 

@Test
public void newStream_duringShutdown() throws Exception {
  InOrder inOrder = inOrder(clientStreamTracerFactory);
  server.start(serverListener);
  client = newClientTransport(server);
  startTransport(client, mockClientTransportListener);
  // Stream prevents termination
  ClientStream stream = client.newStream(methodDescriptor, new Metadata(), callOptions);
  inOrder.verify(clientStreamTracerFactory).newClientStreamTracer(
      any(ClientStreamTracer.StreamInfo.class), any(Metadata.class));
  ClientStreamListenerBase clientStreamListener = new ClientStreamListenerBase();
  stream.start(clientStreamListener);
  client.shutdown(Status.UNAVAILABLE);
  verify(mockClientTransportListener, timeout(TIMEOUT_MS)).transportShutdown(any(Status.class));

  ClientStream stream2 = client.newStream(methodDescriptor, new Metadata(), callOptions);
  inOrder.verify(clientStreamTracerFactory).newClientStreamTracer(
      any(ClientStreamTracer.StreamInfo.class), any(Metadata.class));
  ClientStreamListenerBase clientStreamListener2 = new ClientStreamListenerBase();
  stream2.start(clientStreamListener2);
  Status clientStreamStatus2 =
      clientStreamListener2.status.get(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  assertNotNull(clientStreamListener2.trailers.get(TIMEOUT_MS, TimeUnit.MILLISECONDS));
  assertCodeEquals(Status.UNAVAILABLE, clientStreamStatus2);
  assertNull(clientStreamTracer2.getInboundTrailers());
  assertSame(clientStreamStatus2, clientStreamTracer2.getStatus());

  // Make sure earlier stream works.
  MockServerTransportListener serverTransportListener
      = serverListener.takeListenerOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  serverTransport = serverTransportListener.transport;
  // TODO(zdapeng): Increased timeout to 20 seconds to see if flakiness of #2328 persists. Take
  // further action after sufficient observation.
  StreamCreation serverStreamCreation
      = serverTransportListener.takeStreamOrFail(20 * TIMEOUT_MS, TimeUnit.MILLISECONDS);
  serverStreamCreation.stream.close(Status.OK, new Metadata());
  assertCodeEquals(Status.OK, clientStreamListener.status.get(TIMEOUT_MS, TimeUnit.MILLISECONDS));
  assertNotNull(clientStreamListener.trailers.get(TIMEOUT_MS, TimeUnit.MILLISECONDS));
}
 

@Override
public ClientStreamTracer newClientStreamTracer(CallOptions callOptions, Metadata headers) {
  if (span != BlankSpan.INSTANCE) {
    headers.discardAll(tracingHeader);
    headers.put(tracingHeader, span.getContext());
  }
  return new ClientTracer(span);
}
 

@Test
public void clientStreamTracerTransfers() {
    ClientStreamTracer.Factory factory1 = new ClientStreamTracer.Factory() {};
    ClientStreamTracer.Factory factory2 = new ClientStreamTracer.Factory() {};

    CallOptions baseOptions = CallOptions.DEFAULT.withStreamTracerFactory(factory1);
    CallOptions defaultOptions = CallOptions.DEFAULT.withStreamTracerFactory(factory2);

    DefaultCallOptionsClientInterceptor interceptor = new DefaultCallOptionsClientInterceptor(defaultOptions);

    CallOptions patchedOptions = interceptor.patchOptions(baseOptions);

    assertThat(patchedOptions.getStreamTracerFactories()).containsExactly(factory1, factory2);
}
 

@Override
public ClientStreamTracer newClientStreamTracer(StreamInfo info, Metadata metadata) {
  metadata.put(tracerHeaderKey, tracerKeyValue);
  TestClientStreamTracer tracer = tracers.poll();
  if (tracer != null) {
    return tracer;
  }
  return new TestClientStreamTracer();
}
 

@Override
ClientStream newSubstream(ClientStreamTracer.Factory tracerFactory, Metadata metadata) {
  bufferSizeTracer =
      tracerFactory.newClientStreamTracer(CallOptions.DEFAULT, new Metadata());
  int actualPreviousRpcAttemptsInHeader = metadata.get(GRPC_PREVIOUS_RPC_ATTEMPTS) == null
      ? 0 : Integer.valueOf(metadata.get(GRPC_PREVIOUS_RPC_ATTEMPTS));
  return retriableStreamRecorder.newSubstream(actualPreviousRpcAttemptsInHeader);
}
 

@Test
public void hedging_channelBufferLimitExceeded() {
  ClientStream mockStream1 = mock(ClientStream.class);
  ClientStream mockStream2 = mock(ClientStream.class);
  doReturn(mockStream1).when(retriableStreamRecorder).newSubstream(0);
  doReturn(mockStream2).when(retriableStreamRecorder).newSubstream(1);

  hedgingStream.start(masterListener);
  ArgumentCaptor<ClientStreamListener> sublistenerCaptor1 =
      ArgumentCaptor.forClass(ClientStreamListener.class);
  verify(mockStream1).start(sublistenerCaptor1.capture());

  ClientStreamTracer bufferSizeTracer1 = bufferSizeTracer;
  bufferSizeTracer1.outboundWireSize(100);

  fakeClock.forwardTime(HEDGING_DELAY_IN_SECONDS, TimeUnit.SECONDS);
  ArgumentCaptor<ClientStreamListener> sublistenerCaptor2 =
      ArgumentCaptor.forClass(ClientStreamListener.class);
  verify(mockStream2).start(sublistenerCaptor2.capture());

  ClientStreamTracer bufferSizeTracer2 = bufferSizeTracer;
  bufferSizeTracer2.outboundWireSize(100);

  verify(retriableStreamRecorder, never()).postCommit();

  //  channel bufferLimitExceeded
  channelBufferUsed.addAndGet(CHANNEL_BUFFER_LIMIT - 200);
  bufferSizeTracer2.outboundWireSize(101);

  ArgumentCaptor<Status> statusCaptor = ArgumentCaptor.forClass(Status.class);
  verify(mockStream1).cancel(statusCaptor.capture());
  assertEquals(Status.CANCELLED.getCode(), statusCaptor.getValue().getCode());
  assertEquals(CANCELLED_BECAUSE_COMMITTED, statusCaptor.getValue().getDescription());
  verify(retriableStreamRecorder).postCommit();
  verifyNoMoreInteractions(mockStream1);
  verifyNoMoreInteractions(mockStream2);
  // verify channel buffer is adjusted
  assertEquals(CHANNEL_BUFFER_LIMIT - 200, channelBufferUsed.addAndGet(0));
}
 

@Test
public void nullDelegate() {
  TokenAttachingTracerFactory factory = new TokenAttachingTracerFactory(null);
  ClientStreamTracer.StreamInfo info = ClientStreamTracer.StreamInfo.newBuilder()
      .setTransportAttrs(
          Attributes.newBuilder()
              .set(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS, Attributes.EMPTY).build())
      .build();

  Metadata headers = new Metadata();

  ClientStreamTracer tracer = factory.newClientStreamTracer(info, headers);
  assertThat(tracer).isNotNull();
  assertThat(headers.get(GrpclbConstants.TOKEN_METADATA_KEY)).isNull();
}
 

@Test
public void pickerReturnsStreamTracer_delayed() {
  ClientStream mockStream = mock(ClientStream.class);
  ClientStreamTracer.Factory factory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory factory2 = mock(ClientStreamTracer.Factory.class);
  createChannel();

  CallOptions callOptions = CallOptions.DEFAULT.withStreamTracerFactory(factory1);
  ClientCall<String, Integer> call = channel.newCall(method, callOptions);
  call.start(mockCallListener, new Metadata());

  Subchannel subchannel =
      createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
  requestConnectionSafely(helper, subchannel);
  MockClientTransportInfo transportInfo = transports.poll();
  transportInfo.listener.transportReady();
  ClientTransport mockTransport = transportInfo.transport;
  when(mockTransport.newStream(
          any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class)))
      .thenReturn(mockStream);
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(
      PickResult.withSubchannel(subchannel, factory2));

  updateBalancingStateSafely(helper, READY, mockPicker);
  assertEquals(1, executor.runDueTasks());

  verify(mockPicker).pickSubchannel(any(PickSubchannelArgs.class));
  verify(mockTransport).newStream(same(method), any(Metadata.class), callOptionsCaptor.capture());
  assertEquals(
      Arrays.asList(factory1, factory2),
      callOptionsCaptor.getValue().getStreamTracerFactories());
  // The factories are safely not stubbed because we do not expect any usage of them.
  verifyZeroInteractions(factory1);
  verifyZeroInteractions(factory2);
}
 

@Override
public ClientStreamTracer newClientStreamTracer(
    ClientStreamTracer.StreamInfo info, Metadata headers) {
  TestClientStreamTracer tracer = new TestClientStreamTracer();
  clientStreamTracers.add(tracer);
  return tracer;
}
 

@Test
public void pickerReturnsStreamTracer_noDelay() {
  ClientStream mockStream = mock(ClientStream.class);
  ClientStreamTracer.Factory factory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory factory2 = mock(ClientStreamTracer.Factory.class);
  createChannel();
  Subchannel subchannel =
      createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
  requestConnectionSafely(helper, subchannel);
  MockClientTransportInfo transportInfo = transports.poll();
  transportInfo.listener.transportReady();
  ClientTransport mockTransport = transportInfo.transport;
  when(mockTransport.newStream(
          any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class)))
      .thenReturn(mockStream);

  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(
      PickResult.withSubchannel(subchannel, factory2));
  updateBalancingStateSafely(helper, READY, mockPicker);

  CallOptions callOptions = CallOptions.DEFAULT.withStreamTracerFactory(factory1);
  ClientCall<String, Integer> call = channel.newCall(method, callOptions);
  call.start(mockCallListener, new Metadata());

  verify(mockPicker).pickSubchannel(any(PickSubchannelArgs.class));
  verify(mockTransport).newStream(same(method), any(Metadata.class), callOptionsCaptor.capture());
  assertEquals(
      Arrays.asList(factory1, factory2),
      callOptionsCaptor.getValue().getStreamTracerFactories());
  // The factories are safely not stubbed because we do not expect any usage of them.
  verifyZeroInteractions(factory1);
  verifyZeroInteractions(factory2);
}
 

@Override
public PickResult pickSubchannel(PickSubchannelArgs args) {
  PickResult result = delegate().pickSubchannel(args);
  if (!result.getStatus().isOk()) {
    return result;
  }
  if (result.getSubchannel() == null) {
    return result;
  }
  ClientStreamTracer.Factory originFactory = result.getStreamTracerFactory();
  if (originFactory == null) {
    originFactory = NOOP_CLIENT_STREAM_TRACER_FACTORY;
  }
  return PickResult.withSubchannel(result.getSubchannel(), wrapTracerFactory(originFactory));
}
 

@Override
public ClientStreamTracer newClientStreamTracer(
    ClientStreamTracer.StreamInfo info, Metadata headers) {
  if (span != BlankSpan.INSTANCE) {
    headers.discardAll(tracingHeader);
    headers.put(tracingHeader, span.getContext());
  }
  return new ClientTracer(span);
}
 

@Override
public ClientStreamTracer newClientStreamTracer(StreamInfo info, Metadata headers) {
  OrcaReportBroker broker = info.getCallOptions().getOption(ORCA_REPORT_BROKER_KEY);
  boolean augmented = false;
  if (broker == null) {
    broker = new OrcaReportBroker();
    info =
        info.toBuilder()
            .setCallOptions(info.getCallOptions().withOption(ORCA_REPORT_BROKER_KEY, broker))
            .build();
    augmented = true;
  }
  broker.addListener(listener);
  ClientStreamTracer tracer = delegate.newClientStreamTracer(info, headers);
  if (augmented) {
    final ClientStreamTracer currTracer = tracer;
    final OrcaReportBroker currBroker = broker;
    // The actual tracer that performs ORCA report deserialization.
    tracer =
        new ForwardingClientStreamTracer() {
          @Override
          protected ClientStreamTracer delegate() {
            return currTracer;
          }

          @Override
          public void inboundTrailers(Metadata trailers) {
            OrcaLoadReport report = trailers.get(ORCA_ENDPOINT_LOAD_METRICS_KEY);
            if (report != null) {
              currBroker.onReport(report);
            }
            delegate().inboundTrailers(trailers);
          }
        };
  }
  return tracer;
}
 

/**
 * Tests parent-child load balance policies' listeners both receive per-request ORCA reports upon
 * call trailer arrives and ORCA report deserialization happens only once.
 */
@Test
public void twoLevelPoliciesTypicalWorkflow() {
  ClientStreamTracer.Factory parentFactory =
      mock(ClientStreamTracer.Factory.class,
          delegatesTo(
              OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(orcaListener1)));

  ClientStreamTracer.Factory childFactory =
      OrcaPerRequestUtil.getInstance()
          .newOrcaClientStreamTracerFactory(parentFactory, orcaListener2);
  // Child factory will augment the StreamInfo and pass it to the parent factory.
  ClientStreamTracer childTracer =
      childFactory.newClientStreamTracer(STREAM_INFO, new Metadata());
  ArgumentCaptor<ClientStreamTracer.StreamInfo> streamInfoCaptor = ArgumentCaptor.forClass(null);
  verify(parentFactory).newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class));
  ClientStreamTracer.StreamInfo parentStreamInfo = streamInfoCaptor.getValue();
  assertThat(parentStreamInfo).isNotEqualTo(STREAM_INFO);

  // When the trailer does not contain ORCA report, no listener callback will be invoked.
  Metadata trailer = new Metadata();
  childTracer.inboundTrailers(trailer);
  verifyNoMoreInteractions(orcaListener1);
  verifyNoMoreInteractions(orcaListener2);

  // When the trailer contains an ORCA report, callbacks for both listeners will be invoked.
  // Both listener will receive the same ORCA report instance, which means deserialization
  // happens only once.
  trailer.put(
      OrcaReportingTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY,
      OrcaLoadReport.getDefaultInstance());
  childTracer.inboundTrailers(trailer);
  ArgumentCaptor<OrcaLoadReport> parentReportCap = ArgumentCaptor.forClass(null);
  ArgumentCaptor<OrcaLoadReport> childReportCap = ArgumentCaptor.forClass(null);
  verify(orcaListener1).onLoadReport(parentReportCap.capture());
  verify(orcaListener2).onLoadReport(childReportCap.capture());
  assertThat(parentReportCap.getValue()).isEqualTo(OrcaLoadReport.getDefaultInstance());
  assertThat(childReportCap.getValue()).isSameInstanceAs(parentReportCap.getValue());
}