Java Code Examples for io.grpc.LoadBalancer.PickResult#getSubchannel()

The following examples show how to use io.grpc.LoadBalancer.PickResult#getSubchannel() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: ClientLoadCounter.java    From grpc-java with Apache License 2.0 5 votes vote down vote up 
@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));
}
Example 2
Source File: GrpcUtil.java    From grpc-nebula-java with Apache License 2.0 4 votes vote down vote up 
/**
 * Returns a transport out of a PickResult, or {@code null} if the result is "buffer".
 */
@Nullable
static ClientTransport getTransportFromPickResult(PickResult result, boolean isWaitForReady) {
  final ClientTransport transport;
  Subchannel subchannel = result.getSubchannel();
  if (subchannel != null) {
    transport = ((AbstractSubchannel) subchannel).obtainActiveTransport();
  } else {
    transport = null;
  }
  if (transport != null) {
    final ClientStreamTracer.Factory streamTracerFactory = result.getStreamTracerFactory();
    if (streamTracerFactory == null) {
      return transport;
    }
    return new ClientTransport() {
      @Override
      public ClientStream newStream(
          MethodDescriptor<?, ?> method, Metadata headers, CallOptions callOptions) {
        return transport.newStream(
            method, headers, callOptions.withStreamTracerFactory(streamTracerFactory));
      }

      @Override
      public void ping(PingCallback callback, Executor executor) {
        transport.ping(callback, executor);
      }

      @Override
      public InternalLogId getLogId() {
        return transport.getLogId();
      }

      @Override
      public ListenableFuture<SocketStats> getStats() {
        return transport.getStats();
      }
    };
  }
  if (!result.getStatus().isOk()) {
    if (result.isDrop()) {
      return new FailingClientTransport(result.getStatus(), RpcProgress.DROPPED);
    }
    if (!isWaitForReady) {
      return new FailingClientTransport(result.getStatus(), RpcProgress.PROCESSED);
    }
  }
  return null;
}
Example 3
Source File: LocalityStoreTest.java    From grpc-java with Apache License 2.0 4 votes vote down vote up 
@Test
public void updateLocalityStore_pickResultInterceptedForLoadRecordingWhenSubchannelReady() {
  // Simulate receiving two localities.
  LocalityLbEndpoints localityInfo1 =
      new LocalityLbEndpoints(ImmutableList.of(lbEndpoint11, lbEndpoint12), 1, 0);
  LocalityLbEndpoints localityInfo2 =
      new LocalityLbEndpoints(ImmutableList.of(lbEndpoint21, lbEndpoint22), 2, 0);
  localityStore.updateLocalityStore(ImmutableMap.of(
      locality1, localityInfo1, locality2, localityInfo2));

  // Two child balancers are created.
  assertThat(loadBalancers).hasSize(2);

  ClientStreamTracer.Factory metricsTracingFactory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory metricsTracingFactory2 = mock(ClientStreamTracer.Factory.class);
  when(orcaPerRequestUtil.newOrcaClientStreamTracerFactory(any(ClientStreamTracer.Factory.class),
      any(OrcaPerRequestReportListener.class)))
      .thenReturn(metricsTracingFactory1, metricsTracingFactory2);

  Subchannel subchannel1 = mock(Subchannel.class);
  Subchannel subchannel2 = mock(Subchannel.class);
  final PickResult result1 = PickResult.withSubchannel(subchannel1);
  final PickResult result2 =
      PickResult.withSubchannel(subchannel2, mock(ClientStreamTracer.Factory.class));
  SubchannelPicker subchannelPicker1 = mock(SubchannelPicker.class);
  SubchannelPicker subchannelPicker2 = mock(SubchannelPicker.class);
  when(subchannelPicker1.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(result1);
  when(subchannelPicker2.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(result2);

  // Simulate picker updates for the two localities with dummy pickers.
  childHelpers.get("sz1").updateBalancingState(READY, subchannelPicker1);
  childHelpers.get("sz2").updateBalancingState(READY, subchannelPicker2);

  ArgumentCaptor<SubchannelPicker> interLocalityPickerCaptor = ArgumentCaptor.forClass(null);
  verify(helper, times(2)).updateBalancingState(eq(READY), interLocalityPickerCaptor.capture());
  WeightedRandomPicker interLocalityPicker =
      (WeightedRandomPicker) interLocalityPickerCaptor.getValue();
  assertThat(interLocalityPicker.weightedChildPickers).hasSize(2);

  // Verify each PickResult picked is intercepted with client stream tracer factory for
  // recording load and backend metrics.
  Map<Subchannel, Locality> localitiesBySubchannel
      = ImmutableMap.of(subchannel1, locality1, subchannel2, locality2);
  Map<Subchannel, ClientStreamTracer.Factory> metricsTracingFactoriesBySubchannel
      = ImmutableMap.of(subchannel1, metricsTracingFactory1, subchannel2, metricsTracingFactory2);
  for (int i = 0; i < interLocalityPicker.weightedChildPickers.size(); i++) {
    PickResult pickResult = interLocalityPicker.weightedChildPickers.get(i).getPicker()
        .pickSubchannel(pickSubchannelArgs);
    Subchannel expectedSubchannel = pickResult.getSubchannel();
    Locality expectedLocality = localitiesBySubchannel.get(expectedSubchannel);
    ArgumentCaptor<OrcaPerRequestReportListener> listenerCaptor = ArgumentCaptor.forClass(null);
    verify(orcaPerRequestUtil, times(i + 1))
        .newOrcaClientStreamTracerFactory(any(ClientStreamTracer.Factory.class),
            listenerCaptor.capture());
    assertThat(listenerCaptor.getValue()).isInstanceOf(MetricsRecordingListener.class);
    MetricsRecordingListener listener = (MetricsRecordingListener) listenerCaptor.getValue();
    assertThat(listener.getCounter())
        .isSameInstanceAs(fakeLoadStatsStore.localityCounters.get(expectedLocality));
    assertThat(pickResult.getStreamTracerFactory())
        .isInstanceOf(LoadRecordingStreamTracerFactory.class);
    LoadRecordingStreamTracerFactory loadRecordingFactory =
        (LoadRecordingStreamTracerFactory) pickResult.getStreamTracerFactory();
    assertThat(loadRecordingFactory.getCounter())
        .isSameInstanceAs(fakeLoadStatsStore.localityCounters.get(expectedLocality));
    assertThat(loadRecordingFactory.delegate())
        .isSameInstanceAs(metricsTracingFactoriesBySubchannel.get(expectedSubchannel));
  }
}
Example 4
Source File: GrpcUtil.java    From grpc-java with Apache License 2.0 4 votes vote down vote up 
/**
 * Returns a transport out of a PickResult, or {@code null} if the result is "buffer".
 */
@Nullable
static ClientTransport getTransportFromPickResult(PickResult result, boolean isWaitForReady) {
  final ClientTransport transport;
  Subchannel subchannel = result.getSubchannel();
  if (subchannel != null) {
    transport = ((TransportProvider) subchannel.getInternalSubchannel()).obtainActiveTransport();
  } else {
    transport = null;
  }
  if (transport != null) {
    final ClientStreamTracer.Factory streamTracerFactory = result.getStreamTracerFactory();
    if (streamTracerFactory == null) {
      return transport;
    }
    return new ClientTransport() {
      @Override
      public ClientStream newStream(
          MethodDescriptor<?, ?> method, Metadata headers, CallOptions callOptions) {
        return transport.newStream(
            method, headers, callOptions.withStreamTracerFactory(streamTracerFactory));
      }

      @Override
      public void ping(PingCallback callback, Executor executor) {
        transport.ping(callback, executor);
      }

      @Override
      public InternalLogId getLogId() {
        return transport.getLogId();
      }

      @Override
      public ListenableFuture<SocketStats> getStats() {
        return transport.getStats();
      }
    };
  }
  if (!result.getStatus().isOk()) {
    if (result.isDrop()) {
      return new FailingClientTransport(result.getStatus(), RpcProgress.DROPPED);
    }
    if (!isWaitForReady) {
      return new FailingClientTransport(result.getStatus(), RpcProgress.PROCESSED);
    }
  }
  return null;
}