Examples with PickResult io.grpc.LoadBalancer.PickResult used on opensource projects

Example 6 with PickResult

use of io.grpc.LoadBalancer.PickResult in project grpc-java by grpc.

the class DelayedClientTransportTest method reprocess_newStreamRacesWithReprocess.

@Test
public void reprocess_newStreamRacesWithReprocess() throws Exception {
    final CyclicBarrier barrier = new CyclicBarrier(2);
    // In both phases, we only expect the first pickSubchannel() call to block on the barrier.
    final AtomicBoolean nextPickShouldWait = new AtomicBoolean(true);
    ///////// Phase 1: reprocess() twice with the same picker
    SubchannelPicker picker = mock(SubchannelPicker.class);
    doAnswer(new Answer<PickResult>() {

        @Override
        public PickResult answer(InvocationOnMock invocation) throws Throwable {
            if (nextPickShouldWait.compareAndSet(true, false)) {
                try {
                    barrier.await();
                    return PickResult.withNoResult();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
            return PickResult.withNoResult();
        }
    }).when(picker).pickSubchannel(any(PickSubchannelArgs.class));
    // Because there is no pending stream yet, it will do nothing but save the picker.
    delayedTransport.reprocess(picker);
    verify(picker, never()).pickSubchannel(any(PickSubchannelArgs.class));
    Thread sideThread = new Thread("sideThread") {

        @Override
        public void run() {
            // Will call pickSubchannel and wait on barrier
            delayedTransport.newStream(method, headers, callOptions, statsTraceCtx);
        }
    };
    sideThread.start();
    PickSubchannelArgsImpl args = new PickSubchannelArgsImpl(method, headers, callOptions);
    PickSubchannelArgsImpl args2 = new PickSubchannelArgsImpl(method, headers2, callOptions);
    // Is called from sideThread
    verify(picker, timeout(5000)).pickSubchannel(args);
    // Because stream has not been buffered (it's still stuck in newStream()), this will do nothing,
    // but incrementing the picker version.
    delayedTransport.reprocess(picker);
    verify(picker).pickSubchannel(args);
    // Now let the stuck newStream() through
    barrier.await(5, TimeUnit.SECONDS);
    sideThread.join(5000);
    assertFalse("sideThread should've exited", sideThread.isAlive());
    // newStream() detects that there has been a new picker while it's stuck, thus will pick again.
    verify(picker, times(2)).pickSubchannel(args);
    barrier.reset();
    nextPickShouldWait.set(true);
    ////////// Phase 2: reprocess() with a different picker
    // Create the second stream
    Thread sideThread2 = new Thread("sideThread2") {

        @Override
        public void run() {
            // Will call pickSubchannel and wait on barrier
            delayedTransport.newStream(method, headers2, callOptions, statsTraceCtx);
        }
    };
    sideThread2.start();
    // The second stream will see the first picker
    verify(picker, timeout(5000)).pickSubchannel(args2);
    // While the first stream won't use the first picker any more.
    verify(picker, times(2)).pickSubchannel(args);
    // Now use a different picker
    SubchannelPicker picker2 = mock(SubchannelPicker.class);
    when(picker2.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(PickResult.withNoResult());
    delayedTransport.reprocess(picker2);
    // The pending first stream uses the new picker
    verify(picker2).pickSubchannel(args);
    // The second stream is still pending in creation, doesn't use the new picker.
    verify(picker2, never()).pickSubchannel(args2);
    // Now let the second stream finish creation
    barrier.await(5, TimeUnit.SECONDS);
    sideThread2.join(5000);
    assertFalse("sideThread2 should've exited", sideThread2.isAlive());
    // The second stream should see the new picker
    verify(picker2, timeout(5000)).pickSubchannel(args2);
    // Wrapping up
    verify(picker, times(2)).pickSubchannel(args);
    verify(picker).pickSubchannel(args2);
    verify(picker2).pickSubchannel(args);
    verify(picker2).pickSubchannel(args);
}