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

Example 11 with PickSubchannelArgs

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

the class GrpclbLoadBalancerTest method roundRobinPickerNoDrop.

@Test
public void roundRobinPickerNoDrop() {
    GrpclbClientLoadRecorder loadRecorder = new GrpclbClientLoadRecorder(fakeClock.getTimeProvider());
    Subchannel subchannel = mock(Subchannel.class);
    BackendEntry b1 = new BackendEntry(subchannel, loadRecorder, "LBTOKEN0001");
    BackendEntry b2 = new BackendEntry(subchannel, loadRecorder, "LBTOKEN0002");
    List<BackendEntry> pickList = Arrays.asList(b1, b2);
    RoundRobinPicker picker = new RoundRobinPicker(Collections.<DropEntry>emptyList(), pickList);
    PickSubchannelArgs args1 = mock(PickSubchannelArgs.class);
    Metadata headers1 = new Metadata();
    // The existing token on the headers will be replaced
    headers1.put(GrpclbConstants.TOKEN_METADATA_KEY, "LBTOKEN__OLD");
    when(args1.getHeaders()).thenReturn(headers1);
    assertSame(b1.result, picker.pickSubchannel(args1));
    verify(args1).getHeaders();
    assertThat(headers1.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0001");
    PickSubchannelArgs args2 = mock(PickSubchannelArgs.class);
    Metadata headers2 = new Metadata();
    when(args2.getHeaders()).thenReturn(headers2);
    assertSame(b2.result, picker.pickSubchannel(args2));
    verify(args2).getHeaders();
    assertThat(headers2.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0002");
    PickSubchannelArgs args3 = mock(PickSubchannelArgs.class);
    Metadata headers3 = new Metadata();
    when(args3.getHeaders()).thenReturn(headers3);
    assertSame(b1.result, picker.pickSubchannel(args3));
    verify(args3).getHeaders();
    assertThat(headers3.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0001");
    verify(subchannel, never()).getAttributes();
}

Example 12 with PickSubchannelArgs

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

the class GrpclbLoadBalancerTest method abundantInitialResponse.

@Test
public void abundantInitialResponse() {
    Metadata headers = new Metadata();
    PickSubchannelArgs args = mock(PickSubchannelArgs.class);
    when(args.getHeaders()).thenReturn(headers);
    List<EquivalentAddressGroup> grpclbBalancerList = createResolvedBalancerAddresses(1);
    deliverResolvedAddresses(Collections.<EquivalentAddressGroup>emptyList(), grpclbBalancerList);
    assertEquals(1, fakeOobChannels.size());
    verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture());
    StreamObserver<LoadBalanceResponse> lbResponseObserver = lbResponseObserverCaptor.getValue();
    // Simulate LB initial response
    assertEquals(0, fakeClock.numPendingTasks(LOAD_REPORTING_TASK_FILTER));
    lbResponseObserver.onNext(buildInitialResponse(1983));
    // Load reporting task is scheduled
    assertEquals(1, fakeClock.numPendingTasks(LOAD_REPORTING_TASK_FILTER));
    FakeClock.ScheduledTask scheduledTask = Iterables.getOnlyElement(fakeClock.getPendingTasks(LOAD_REPORTING_TASK_FILTER));
    assertEquals(1983, scheduledTask.getDelay(TimeUnit.MILLISECONDS));
    logs.clear();
    // Simulate an abundant LB initial response, with a different report interval
    lbResponseObserver.onNext(buildInitialResponse(9097));
    // This incident is logged
    assertThat(logs).containsExactly("DEBUG: [grpclb-<api.google.com>] Got an LB response: " + buildInitialResponse(9097), "WARNING: [grpclb-<api.google.com>] " + "Ignoring unexpected response type: INITIAL_RESPONSE").inOrder();
    // It doesn't affect load-reporting at all
    assertThat(fakeClock.getPendingTasks(LOAD_REPORTING_TASK_FILTER)).containsExactly(scheduledTask);
    assertEquals(1983, scheduledTask.getDelay(TimeUnit.MILLISECONDS));
}

Example 13 with PickSubchannelArgs

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

the class GrpclbLoadBalancerTest method loadReporting.

@Test
public void loadReporting() {
    Metadata headers = new Metadata();
    PickSubchannelArgs args = mock(PickSubchannelArgs.class);
    when(args.getHeaders()).thenReturn(headers);
    long loadReportIntervalMillis = 1983;
    List<EquivalentAddressGroup> grpclbBalancerList = createResolvedBalancerAddresses(1);
    deliverResolvedAddresses(Collections.<EquivalentAddressGroup>emptyList(), grpclbBalancerList);
    // Fallback timer is started as soon as address is resolved.
    assertEquals(1, fakeClock.numPendingTasks(FALLBACK_MODE_TASK_FILTER));
    assertEquals(1, fakeOobChannels.size());
    verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture());
    StreamObserver<LoadBalanceResponse> lbResponseObserver = lbResponseObserverCaptor.getValue();
    assertEquals(1, lbRequestObservers.size());
    StreamObserver<LoadBalanceRequest> lbRequestObserver = lbRequestObservers.poll();
    InOrder inOrder = inOrder(lbRequestObserver);
    InOrder helperInOrder = inOrder(helper, subchannelPool);
    inOrder.verify(lbRequestObserver).onNext(eq(LoadBalanceRequest.newBuilder().setInitialRequest(InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()).build()));
    // Simulate receiving LB response
    assertEquals(0, fakeClock.numPendingTasks(LOAD_REPORTING_TASK_FILTER));
    lbResponseObserver.onNext(buildInitialResponse(loadReportIntervalMillis));
    // Load reporting task is scheduled
    assertEquals(1, fakeClock.numPendingTasks(LOAD_REPORTING_TASK_FILTER));
    assertEquals(0, fakeClock.runDueTasks());
    List<ServerEntry> backends = Arrays.asList(new ServerEntry("127.0.0.1", 2000, "token0001"), // drop
    new ServerEntry("token0001"), new ServerEntry("127.0.0.1", 2010, "token0002"), // drop
    new ServerEntry("token0003"));
    lbResponseObserver.onNext(buildLbResponse(backends));
    assertEquals(2, mockSubchannels.size());
    Subchannel subchannel1 = mockSubchannels.poll();
    Subchannel subchannel2 = mockSubchannels.poll();
    deliverSubchannelState(subchannel1, ConnectivityStateInfo.forNonError(CONNECTING));
    deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(CONNECTING));
    deliverSubchannelState(subchannel1, ConnectivityStateInfo.forNonError(READY));
    deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(READY));
    helperInOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
    RoundRobinPicker picker = (RoundRobinPicker) pickerCaptor.getValue();
    assertThat(picker.dropList).containsExactly(null, new DropEntry(getLoadRecorder(), "token0001"), null, new DropEntry(getLoadRecorder(), "token0003")).inOrder();
    assertThat(picker.pickList).containsExactly(new BackendEntry(subchannel1, getLoadRecorder(), "token0001"), new BackendEntry(subchannel2, getLoadRecorder(), "token0002")).inOrder();
    // Report, no data
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().build());
    PickResult pick1 = picker.pickSubchannel(args);
    assertSame(subchannel1, pick1.getSubchannel());
    assertSame(getLoadRecorder(), pick1.getStreamTracerFactory());
    // Merely the pick will not be recorded as upstart.
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().build());
    ClientStreamTracer tracer1 = pick1.getStreamTracerFactory().newClientStreamTracer(STREAM_INFO, new Metadata());
    tracer1.streamCreated(Attributes.EMPTY, new Metadata());
    PickResult pick2 = picker.pickSubchannel(args);
    assertNull(pick2.getSubchannel());
    assertSame(DROP_PICK_RESULT, pick2);
    // Report includes upstart of pick1 and the drop of pick2
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().setNumCallsStarted(2).setNumCallsFinished(// pick2
    1).addCallsFinishedWithDrop(ClientStatsPerToken.newBuilder().setLoadBalanceToken("token0001").setNumCalls(// pick2
    1).build()).build());
    PickResult pick3 = picker.pickSubchannel(args);
    assertSame(subchannel2, pick3.getSubchannel());
    assertSame(getLoadRecorder(), pick3.getStreamTracerFactory());
    ClientStreamTracer tracer3 = pick3.getStreamTracerFactory().newClientStreamTracer(STREAM_INFO, new Metadata());
    tracer3.streamCreated(Attributes.EMPTY, new Metadata());
    // pick3 has sent out headers
    tracer3.outboundHeaders();
    // 3rd report includes pick3's upstart
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().setNumCallsStarted(1).build());
    PickResult pick4 = picker.pickSubchannel(args);
    assertNull(pick4.getSubchannel());
    assertSame(DROP_PICK_RESULT, pick4);
    // pick1 ended without sending anything
    tracer1.streamClosed(Status.CANCELLED);
    // 4th report includes end of pick1 and drop of pick4
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().setNumCallsStarted(// pick4
    1).setNumCallsFinished(2).setNumCallsFinishedWithClientFailedToSend(// pick1
    1).addCallsFinishedWithDrop(ClientStatsPerToken.newBuilder().setLoadBalanceToken("token0003").setNumCalls(// pick4
    1).build()).build());
    PickResult pick5 = picker.pickSubchannel(args);
    assertSame(subchannel1, pick1.getSubchannel());
    assertSame(getLoadRecorder(), pick5.getStreamTracerFactory());
    ClientStreamTracer tracer5 = pick5.getStreamTracerFactory().newClientStreamTracer(STREAM_INFO, new Metadata());
    tracer5.streamCreated(Attributes.EMPTY, new Metadata());
    // pick3 ended without receiving response headers
    tracer3.streamClosed(Status.DEADLINE_EXCEEDED);
    // pick5 sent and received headers
    tracer5.outboundHeaders();
    tracer5.inboundHeaders();
    // 5th report includes pick3's end and pick5's upstart
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().setNumCallsStarted(// pick5
    1).setNumCallsFinished(// pick3
    1).build());
    // pick5 ends
    tracer5.streamClosed(Status.OK);
    // 6th report includes pick5's end
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().setNumCallsFinished(1).setNumCallsFinishedKnownReceived(1).build());
    assertEquals(1, fakeClock.numPendingTasks());
    // Balancer closes the stream, scheduled reporting task cancelled
    lbResponseObserver.onError(Status.UNAVAILABLE.asException());
    assertEquals(0, fakeClock.numPendingTasks());
    // New stream created
    verify(mockLbService, times(2)).balanceLoad(lbResponseObserverCaptor.capture());
    lbResponseObserver = lbResponseObserverCaptor.getValue();
    assertEquals(1, lbRequestObservers.size());
    lbRequestObserver = lbRequestObservers.poll();
    inOrder = inOrder(lbRequestObserver);
    inOrder.verify(lbRequestObserver).onNext(eq(LoadBalanceRequest.newBuilder().setInitialRequest(InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()).build()));
    // Load reporting is also requested
    lbResponseObserver.onNext(buildInitialResponse(loadReportIntervalMillis));
    // No picker created because balancer is still using the results from the last stream
    helperInOrder.verify(helper, never()).updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
    // Make a new pick on that picker.  It will not show up on the report of the new stream, because
    // that picker is associated with the previous stream.
    PickResult pick6 = picker.pickSubchannel(args);
    assertNull(pick6.getSubchannel());
    assertSame(DROP_PICK_RESULT, pick6);
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().build());
    // New stream got the list update
    lbResponseObserver.onNext(buildLbResponse(backends));
    // Same backends, thus no new subchannels
    helperInOrder.verify(subchannelPool, never()).takeOrCreateSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class));
    // But the new RoundRobinEntries have a new loadRecorder, thus considered different from
    // the previous list, thus a new picker is created
    helperInOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
    picker = (RoundRobinPicker) pickerCaptor.getValue();
    PickResult pick1p = picker.pickSubchannel(args);
    assertSame(subchannel1, pick1p.getSubchannel());
    assertSame(getLoadRecorder(), pick1p.getStreamTracerFactory());
    pick1p.getStreamTracerFactory().newClientStreamTracer(STREAM_INFO, new Metadata());
    // The pick from the new stream will be included in the report
    assertNextReport(inOrder, lbRequestObserver, loadReportIntervalMillis, ClientStats.newBuilder().setNumCallsStarted(1).build());
    verify(args, atLeast(0)).getHeaders();
    verifyNoMoreInteractions(args);
}

Example 14 with PickSubchannelArgs

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

the class ManagedChannelImplTest method newCallWithConfigSelector.

@Test
public void newCallWithConfigSelector() {
    FakeNameResolverFactory nameResolverFactory = new FakeNameResolverFactory.Builder(expectedUri).setServers(ImmutableList.of(addressGroup)).build();
    channelBuilder.nameResolverFactory(nameResolverFactory);
    channel = new ManagedChannelImpl(channelBuilder, mockTransportFactory, new FakeBackoffPolicyProvider(), balancerRpcExecutorPool, timer.getStopwatchSupplier(), Collections.<ClientInterceptor>emptyList(), timer.getTimeProvider());
    nameResolverFactory.nextConfigOrError.set(ConfigOrError.fromConfig(ManagedChannelServiceConfig.empty()));
    final Metadata.Key<String> metadataKey = Metadata.Key.of("test", Metadata.ASCII_STRING_MARSHALLER);
    final CallOptions.Key<String> callOptionsKey = CallOptions.Key.create("test");
    InternalConfigSelector configSelector = new InternalConfigSelector() {

        @Override
        public Result selectConfig(final PickSubchannelArgs args) {
            return Result.newBuilder().setConfig(ManagedChannelServiceConfig.empty()).setInterceptor(// An interceptor that mutates CallOptions based on headers value.
            new ClientInterceptor() {

                String value = args.getHeaders().get(metadataKey);

                @Override
                public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
                    callOptions = callOptions.withOption(callOptionsKey, value);
                    return next.newCall(method, callOptions);
                }
            }).build();
        }
    };
    nameResolverFactory.nextAttributes.set(Attributes.newBuilder().set(InternalConfigSelector.KEY, configSelector).build());
    channel.getState(true);
    Metadata headers = new Metadata();
    headers.put(metadataKey, "fooValue");
    ClientStream mockStream = mock(ClientStream.class);
    ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
    call.start(mockCallListener, headers);
    ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
    verify(mockLoadBalancerProvider).newLoadBalancer(helperCaptor.capture());
    helper = helperCaptor.getValue();
    // Make the transport available
    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;
    ManagedClientTransport.Listener transportListener = transportInfo.listener;
    when(mockTransport.newStream(same(method), same(headers), any(CallOptions.class), ArgumentMatchers.<ClientStreamTracer[]>any())).thenReturn(mockStream);
    transportListener.transportReady();
    when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(PickResult.withSubchannel(subchannel));
    updateBalancingStateSafely(helper, READY, mockPicker);
    executor.runDueTasks();
    ArgumentCaptor<CallOptions> callOptionsCaptor = ArgumentCaptor.forClass(null);
    verify(mockTransport).newStream(same(method), same(headers), callOptionsCaptor.capture(), ArgumentMatchers.<ClientStreamTracer[]>any());
    assertThat(callOptionsCaptor.getValue().getOption(callOptionsKey)).isEqualTo("fooValue");
    verify(mockStream).start(streamListenerCaptor.capture());
    // Clean up as much as possible to allow the channel to terminate.
    shutdownSafely(helper, subchannel);
    timer.forwardNanos(TimeUnit.SECONDS.toNanos(ManagedChannelImpl.SUBCHANNEL_SHUTDOWN_DELAY_SECONDS));
}

Example 15 with PickSubchannelArgs

use of io.grpc.LoadBalancer.PickSubchannelArgs 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();
    }
}