Examples with ClientStreamTracer io.grpc.ClientStreamTracer used on opensource projects

Example 56 with ClientStreamTracer

use of io.grpc.ClientStreamTracer in project grpc-java by grpc.

the class RetriableStreamTest method hedging_channelBufferLimitExceeded.

@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());
    verify(mockStream1).isReady();
    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());
    verify(mockStream1, times(2)).isReady();
    verify(mockStream2).isReady();
    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));
}

Example 57 with ClientStreamTracer

use of io.grpc.ClientStreamTracer in project grpc-java by grpc.

the class LeastRequestLoadBalancerTest method pickerLeastRequest.

@Test
public void pickerLeastRequest() throws Exception {
    int choiceCount = 2;
    // This should add inFlight counters to all subchannels.
    loadBalancer.handleResolvedAddresses(ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(Attributes.EMPTY).setLoadBalancingPolicyConfig(new LeastRequestConfig(choiceCount)).build());
    assertEquals(3, loadBalancer.getSubchannels().size());
    List<Subchannel> subchannels = Lists.newArrayList(loadBalancer.getSubchannels());
    // Make sure all inFlight counters have started at 0
    assertEquals(0, subchannels.get(0).getAttributes().get(IN_FLIGHTS).get());
    assertEquals(0, subchannels.get(1).getAttributes().get(IN_FLIGHTS).get());
    assertEquals(0, subchannels.get(2).getAttributes().get(IN_FLIGHTS).get());
    for (Subchannel sc : subchannels) {
        deliverSubchannelState(sc, ConnectivityStateInfo.forNonError(READY));
    }
    // Capture the active ReadyPicker once all subchannels are READY
    verify(mockHelper, times(4)).updateBalancingState(any(ConnectivityState.class), pickerCaptor.capture());
    assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class);
    ReadyPicker picker = (ReadyPicker) pickerCaptor.getValue();
    assertThat(picker.getList()).containsExactlyElementsIn(subchannels);
    // Make random return 0, then 2 for the sample indexes.
    when(mockRandom.nextInt(subchannels.size())).thenReturn(0, 2);
    PickResult pickResult1 = picker.pickSubchannel(mockArgs);
    verify(mockRandom, times(choiceCount)).nextInt(subchannels.size());
    assertEquals(subchannels.get(0), pickResult1.getSubchannel());
    // This simulates sending the actual RPC on the picked channel
    ClientStreamTracer streamTracer1 = pickResult1.getStreamTracerFactory().newClientStreamTracer(StreamInfo.newBuilder().build(), new Metadata());
    streamTracer1.streamCreated(Attributes.EMPTY, new Metadata());
    assertEquals(1, pickResult1.getSubchannel().getAttributes().get(IN_FLIGHTS).get());
    // For the second pick it should pick the one with lower inFlight.
    when(mockRandom.nextInt(subchannels.size())).thenReturn(0, 2);
    PickResult pickResult2 = picker.pickSubchannel(mockArgs);
    // Since this is the second pick we expect the total random samples to be choiceCount * 2
    verify(mockRandom, times(choiceCount * 2)).nextInt(subchannels.size());
    assertEquals(subchannels.get(2), pickResult2.getSubchannel());
    // For the third pick we unavoidably pick subchannel with index 1.
    when(mockRandom.nextInt(subchannels.size())).thenReturn(1, 1);
    PickResult pickResult3 = picker.pickSubchannel(mockArgs);
    verify(mockRandom, times(choiceCount * 3)).nextInt(subchannels.size());
    assertEquals(subchannels.get(1), pickResult3.getSubchannel());
    // Finally ensure a finished RPC decreases inFlight
    streamTracer1.streamClosed(Status.OK);
    assertEquals(0, pickResult1.getSubchannel().getAttributes().get(IN_FLIGHTS).get());
}