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

Example 21 with SubchannelPicker

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

the class ManagedChannelImplTest method subtestCallsAndShutdown.

private void subtestCallsAndShutdown(boolean shutdownNow, boolean shutdownNowAfterShutdown) {
    FakeNameResolverFactory nameResolverFactory = new FakeNameResolverFactory.Builder(expectedUri).build();
    channelBuilder.nameResolverFactory(nameResolverFactory);
    createChannel();
    verify(executorPool).getObject();
    ClientStream mockStream = mock(ClientStream.class);
    ClientStream mockStream2 = mock(ClientStream.class);
    Metadata headers = new Metadata();
    Metadata headers2 = new Metadata();
    // Configure the picker so that first RPC goes to delayed transport, and second RPC goes to
    // real transport.
    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
    requestConnectionSafely(helper, subchannel);
    verify(mockTransportFactory).newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class), any(ChannelLogger.class));
    MockClientTransportInfo transportInfo = transports.poll();
    ConnectionClientTransport mockTransport = transportInfo.transport;
    verify(mockTransport).start(any(ManagedClientTransport.Listener.class));
    ManagedClientTransport.Listener transportListener = transportInfo.listener;
    when(mockTransport.newStream(same(method), same(headers), same(CallOptions.DEFAULT), ArgumentMatchers.<ClientStreamTracer[]>any())).thenReturn(mockStream);
    when(mockTransport.newStream(same(method), same(headers2), same(CallOptions.DEFAULT), ArgumentMatchers.<ClientStreamTracer[]>any())).thenReturn(mockStream2);
    transportListener.transportReady();
    when(mockPicker.pickSubchannel(new PickSubchannelArgsImpl(method, headers, CallOptions.DEFAULT))).thenReturn(PickResult.withNoResult());
    when(mockPicker.pickSubchannel(new PickSubchannelArgsImpl(method, headers2, CallOptions.DEFAULT))).thenReturn(PickResult.withSubchannel(subchannel));
    updateBalancingStateSafely(helper, READY, mockPicker);
    // First RPC, will be pending
    ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
    verify(mockTransportFactory).newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class), any(ChannelLogger.class));
    call.start(mockCallListener, headers);
    verify(mockTransport, never()).newStream(same(method), same(headers), same(CallOptions.DEFAULT), ArgumentMatchers.<ClientStreamTracer[]>any());
    // Second RPC, will be assigned to the real transport
    ClientCall<String, Integer> call2 = channel.newCall(method, CallOptions.DEFAULT);
    call2.start(mockCallListener2, headers2);
    verify(mockTransport).newStream(same(method), same(headers2), same(CallOptions.DEFAULT), ArgumentMatchers.<ClientStreamTracer[]>any());
    verify(mockTransport).newStream(same(method), same(headers2), same(CallOptions.DEFAULT), ArgumentMatchers.<ClientStreamTracer[]>any());
    verify(mockStream2).start(any(ClientStreamListener.class));
    // Shutdown
    if (shutdownNow) {
        channel.shutdownNow();
    } else {
        channel.shutdown();
        if (shutdownNowAfterShutdown) {
            channel.shutdownNow();
            shutdownNow = true;
        }
    }
    assertTrue(channel.isShutdown());
    assertFalse(channel.isTerminated());
    assertThat(nameResolverFactory.resolvers).hasSize(1);
    verify(mockLoadBalancerProvider).newLoadBalancer(any(Helper.class));
    // Further calls should fail without going to the transport
    ClientCall<String, Integer> call3 = channel.newCall(method, CallOptions.DEFAULT);
    call3.start(mockCallListener3, headers2);
    timer.runDueTasks();
    executor.runDueTasks();
    verify(mockCallListener3).onClose(statusCaptor.capture(), any(Metadata.class));
    assertSame(Status.Code.UNAVAILABLE, statusCaptor.getValue().getCode());
    if (shutdownNow) {
        // LoadBalancer and NameResolver are shut down as soon as delayed transport is terminated.
        verify(mockLoadBalancer).shutdown();
        assertTrue(nameResolverFactory.resolvers.get(0).shutdown);
        // call should have been aborted by delayed transport
        executor.runDueTasks();
        verify(mockCallListener).onClose(same(ManagedChannelImpl.SHUTDOWN_NOW_STATUS), any(Metadata.class));
    } else {
        // LoadBalancer and NameResolver are still running.
        verify(mockLoadBalancer, never()).shutdown();
        assertFalse(nameResolverFactory.resolvers.get(0).shutdown);
        // call and call2 are still alive, and can still be assigned to a real transport
        SubchannelPicker picker2 = mock(SubchannelPicker.class);
        when(picker2.pickSubchannel(new PickSubchannelArgsImpl(method, headers, CallOptions.DEFAULT))).thenReturn(PickResult.withSubchannel(subchannel));
        updateBalancingStateSafely(helper, READY, picker2);
        executor.runDueTasks();
        verify(mockTransport).newStream(same(method), same(headers), same(CallOptions.DEFAULT), ArgumentMatchers.<ClientStreamTracer[]>any());
        verify(mockStream).start(any(ClientStreamListener.class));
    }
    // After call is moved out of delayed transport, LoadBalancer, NameResolver and the transports
    // will be shutdown.
    verify(mockLoadBalancer).shutdown();
    assertTrue(nameResolverFactory.resolvers.get(0).shutdown);
    if (shutdownNow) {
        // Channel shutdownNow() all subchannels after shutting down LoadBalancer
        verify(mockTransport).shutdownNow(ManagedChannelImpl.SHUTDOWN_NOW_STATUS);
    } else {
        verify(mockTransport, never()).shutdownNow(any(Status.class));
    }
    // LoadBalancer should shutdown the subchannel
    shutdownSafely(helper, subchannel);
    if (shutdownNow) {
        verify(mockTransport).shutdown(same(ManagedChannelImpl.SHUTDOWN_NOW_STATUS));
    } else {
        verify(mockTransport).shutdown(same(ManagedChannelImpl.SHUTDOWN_STATUS));
    }
    // Killing the remaining real transport will terminate the channel
    transportListener.transportShutdown(Status.UNAVAILABLE);
    assertFalse(channel.isTerminated());
    verify(executorPool, never()).returnObject(any());
    transportListener.transportTerminated();
    assertTrue(channel.isTerminated());
    verify(executorPool).returnObject(executor.getScheduledExecutorService());
    verifyNoMoreInteractions(balancerRpcExecutorPool);
    verify(mockTransportFactory).newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class), any(ChannelLogger.class));
    verify(mockTransportFactory).close();
    verify(mockTransport, atLeast(0)).getLogId();
    verifyNoMoreInteractions(mockTransport);
}

Example 22 with SubchannelPicker

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

the class ManagedChannelImplTest method allServersFailedToConnect.

/**
 * Verify that if all resolved addresses failed to connect, a fail-fast call will fail, while a
 * wait-for-ready call will still be buffered.
 */
@Test
public void allServersFailedToConnect() throws Exception {
    final SocketAddress addr1 = new SocketAddress() {

        @Override
        public String toString() {
            return "addr1";
        }
    };
    final SocketAddress addr2 = new SocketAddress() {

        @Override
        public String toString() {
            return "addr2";
        }
    };
    InOrder inOrder = inOrder(mockLoadBalancer, subchannelStateListener);
    List<SocketAddress> resolvedAddrs = Arrays.asList(addr1, addr2);
    FakeNameResolverFactory nameResolverFactory = new FakeNameResolverFactory.Builder(expectedUri).setServers(Collections.singletonList(new EquivalentAddressGroup(resolvedAddrs))).build();
    channelBuilder.nameResolverFactory(nameResolverFactory);
    createChannel();
    // Start a wait-for-ready call
    ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT.withWaitForReady());
    Metadata headers = new Metadata();
    call.start(mockCallListener, headers);
    // ... and a fail-fast call
    ClientCall<String, Integer> call2 = channel.newCall(method, CallOptions.DEFAULT.withoutWaitForReady());
    call2.start(mockCallListener2, headers);
    executor.runDueTasks();
    // Simulate name resolution results
    EquivalentAddressGroup addressGroup = new EquivalentAddressGroup(resolvedAddrs);
    inOrder.verify(mockLoadBalancer).handleResolvedAddresses(resolvedAddressCaptor.capture());
    assertThat(resolvedAddressCaptor.getValue().getAddresses()).containsExactly(addressGroup);
    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
    when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(PickResult.withSubchannel(subchannel));
    requestConnectionSafely(helper, subchannel);
    inOrder.verify(subchannelStateListener).onSubchannelState(stateInfoCaptor.capture());
    assertEquals(CONNECTING, stateInfoCaptor.getValue().getState());
    // Connecting to server1, which will fail
    verify(mockTransportFactory).newClientTransport(same(addr1), any(ClientTransportOptions.class), any(ChannelLogger.class));
    verify(mockTransportFactory, times(0)).newClientTransport(same(addr2), any(ClientTransportOptions.class), any(ChannelLogger.class));
    MockClientTransportInfo transportInfo1 = transports.poll();
    transportInfo1.listener.transportShutdown(Status.UNAVAILABLE);
    // Connecting to server2, which will fail too
    verify(mockTransportFactory).newClientTransport(same(addr2), any(ClientTransportOptions.class), any(ChannelLogger.class));
    MockClientTransportInfo transportInfo2 = transports.poll();
    Status server2Error = Status.UNAVAILABLE.withDescription("Server2 failed to connect");
    transportInfo2.listener.transportShutdown(server2Error);
    // ... which makes the subchannel enter TRANSIENT_FAILURE. The last error Status is propagated
    // to LoadBalancer.
    inOrder.verify(subchannelStateListener).onSubchannelState(stateInfoCaptor.capture());
    assertEquals(TRANSIENT_FAILURE, stateInfoCaptor.getValue().getState());
    assertSame(server2Error, stateInfoCaptor.getValue().getStatus());
    // A typical LoadBalancer would create a picker with error
    SubchannelPicker picker2 = mock(SubchannelPicker.class);
    when(picker2.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(PickResult.withError(server2Error));
    updateBalancingStateSafely(helper, TRANSIENT_FAILURE, picker2);
    executor.runDueTasks();
    // ... which fails the fail-fast call
    verify(mockCallListener2).onClose(same(server2Error), any(Metadata.class));
    // ... while the wait-for-ready call stays
    verifyNoMoreInteractions(mockCallListener);
    // No real stream was ever created
    verify(transportInfo1.transport, times(0)).newStream(any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class), ArgumentMatchers.<ClientStreamTracer[]>any());
    verify(transportInfo2.transport, times(0)).newStream(any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class), ArgumentMatchers.<ClientStreamTracer[]>any());
}

Example 23 with SubchannelPicker

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

the class ManagedChannelImplTest method oobChannelHasNoChannelCallCredentials.

@Test
public void oobChannelHasNoChannelCallCredentials() {
    Metadata.Key<String> metadataKey = Metadata.Key.of("token", Metadata.ASCII_STRING_MARSHALLER);
    String channelCredValue = "channel-provided call cred";
    channelBuilder = new ManagedChannelImplBuilder(TARGET, InsecureChannelCredentials.create(), new FakeCallCredentials(metadataKey, channelCredValue), new UnsupportedClientTransportFactoryBuilder(), new FixedPortProvider(DEFAULT_PORT));
    channelBuilder.disableRetry();
    configureBuilder(channelBuilder);
    createChannel();
    // Verify that the normal channel has call creds, to validate configuration
    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
    requestConnectionSafely(helper, subchannel);
    MockClientTransportInfo transportInfo = transports.poll();
    assertNotNull(transportInfo);
    transportInfo.listener.transportReady();
    when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(PickResult.withSubchannel(subchannel));
    updateBalancingStateSafely(helper, READY, mockPicker);
    String callCredValue = "per-RPC call cred";
    CallOptions callOptions = CallOptions.DEFAULT.withCallCredentials(new FakeCallCredentials(metadataKey, callCredValue));
    Metadata headers = new Metadata();
    ClientCall<String, Integer> call = channel.newCall(method, callOptions);
    call.start(mockCallListener, headers);
    verify(transportInfo.transport).newStream(same(method), same(headers), same(callOptions), ArgumentMatchers.<ClientStreamTracer[]>any());
    assertThat(headers.getAll(metadataKey)).containsExactly(channelCredValue, callCredValue).inOrder();
    // Verify that the oob channel does not
    ManagedChannel oob = helper.createOobChannel(Collections.singletonList(addressGroup), "oobauthority");
    headers = new Metadata();
    call = oob.newCall(method, callOptions);
    call.start(mockCallListener2, headers);
    transportInfo = transports.poll();
    assertNotNull(transportInfo);
    transportInfo.listener.transportReady();
    balancerRpcExecutor.runDueTasks();
    verify(transportInfo.transport).newStream(same(method), same(headers), same(callOptions), ArgumentMatchers.<ClientStreamTracer[]>any());
    assertThat(headers.getAll(metadataKey)).containsExactly(callCredValue);
    oob.shutdownNow();
    // Verify that resolving oob channel does not
    oob = helper.createResolvingOobChannelBuilder("oobauthority").nameResolverFactory(new FakeNameResolverFactory.Builder(URI.create("oobauthority")).build()).defaultLoadBalancingPolicy(MOCK_POLICY_NAME).idleTimeout(ManagedChannelImplBuilder.IDLE_MODE_MAX_TIMEOUT_DAYS, TimeUnit.DAYS).disableRetry().build();
    oob.getState(true);
    ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(Helper.class);
    verify(mockLoadBalancerProvider, times(2)).newLoadBalancer(helperCaptor.capture());
    Helper oobHelper = helperCaptor.getValue();
    subchannel = createSubchannelSafely(oobHelper, addressGroup, Attributes.EMPTY, subchannelStateListener);
    requestConnectionSafely(oobHelper, subchannel);
    transportInfo = transports.poll();
    assertNotNull(transportInfo);
    transportInfo.listener.transportReady();
    SubchannelPicker mockPicker2 = mock(SubchannelPicker.class);
    when(mockPicker2.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(PickResult.withSubchannel(subchannel));
    updateBalancingStateSafely(oobHelper, READY, mockPicker2);
    headers = new Metadata();
    call = oob.newCall(method, callOptions);
    call.start(mockCallListener2, headers);
    // CallOptions may contain StreamTracerFactory for census that is added by default.
    verify(transportInfo.transport).newStream(same(method), same(headers), any(CallOptions.class), ArgumentMatchers.<ClientStreamTracer[]>any());
    assertThat(headers.getAll(metadataKey)).containsExactly(callCredValue);
    oob.shutdownNow();
}

Example 24 with SubchannelPicker

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

the class WeightedRandomPicker method pickSubchannel.

@Override
public final PickResult pickSubchannel(PickSubchannelArgs args) {
    SubchannelPicker childPicker = null;
    if (totalWeight == 0) {
        childPicker = weightedChildPickers.get(random.nextInt(weightedChildPickers.size())).getPicker();
    } else {
        int rand = random.nextInt(totalWeight);
        // Find the first idx such that rand < accumulatedWeights[idx]
        // Not using Arrays.binarySearch for better readability.
        int accumulatedWeight = 0;
        for (int idx = 0; idx < weightedChildPickers.size(); idx++) {
            accumulatedWeight += weightedChildPickers.get(idx).getWeight();
            if (rand < accumulatedWeight) {
                childPicker = weightedChildPickers.get(idx).getPicker();
                break;
            }
        }
        checkNotNull(childPicker, "childPicker not found");
    }
    return childPicker.pickSubchannel(args);
}

Example 25 with SubchannelPicker

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

the class DelayedClientTransport method newStream.

/**
   * If a {@link SubchannelPicker} is being, or has been provided via {@link #reprocess}, the last
   * picker will be consulted.
   *
   * <p>Otherwise, if the delayed transport is not shutdown, then a {@link PendingStream} is
   * returned; if the transport is shutdown, then a {@link FailingClientStream} is returned.
   */
@Override
public final ClientStream newStream(MethodDescriptor<?, ?> method, Metadata headers, CallOptions callOptions, StatsTraceContext statsTraceCtx) {
    try {
        SubchannelPicker picker = null;
        PickSubchannelArgs args = new PickSubchannelArgsImpl(method, headers, callOptions);
        long pickerVersion = -1;
        synchronized (lock) {
            if (!shutdown) {
                if (lastPicker == null) {
                    return createPendingStream(args, statsTraceCtx);
                }
                picker = lastPicker;
                pickerVersion = lastPickerVersion;
            }
        }
        if (picker != null) {
            while (true) {
                PickResult pickResult = picker.pickSubchannel(args);
                ClientTransport transport = GrpcUtil.getTransportFromPickResult(pickResult, callOptions.isWaitForReady());
                if (transport != null) {
                    return transport.newStream(args.getMethodDescriptor(), args.getHeaders(), args.getCallOptions(), statsTraceCtx);
                }
                // picker.
                synchronized (lock) {
                    if (shutdown) {
                        break;
                    }
                    if (pickerVersion == lastPickerVersion) {
                        return createPendingStream(args, statsTraceCtx);
                    }
                    picker = lastPicker;
                    pickerVersion = lastPickerVersion;
                }
            }
        }
        return new FailingClientStream(Status.UNAVAILABLE.withDescription("Channel has shutdown (reported by delayed transport)"));
    } finally {
        channelExecutor.drain();
    }
}