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

Example 16 with PickSubchannelArgs

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

the class GrpclbLoadBalancerTest method roundRobinPicker.

@Test
public void roundRobinPicker() {
    Subchannel subchannel = mock(Subchannel.class);
    RoundRobinEntry r1 = new RoundRobinEntry(Status.UNAVAILABLE.withDescription("Just error"));
    RoundRobinEntry r2 = new RoundRobinEntry(subchannel, "LBTOKEN0001");
    RoundRobinEntry r3 = new RoundRobinEntry(subchannel, "LBTOKEN0002");
    List<RoundRobinEntry> list = Arrays.asList(r1, r2, r3);
    RoundRobinPicker picker = new RoundRobinPicker(list);
    PickSubchannelArgs args1 = mock(PickSubchannelArgs.class);
    Metadata headers1 = new Metadata();
    when(args1.getHeaders()).thenReturn(headers1);
    assertSame(r1.result, picker.pickSubchannel(args1));
    verify(args1).getHeaders();
    assertFalse(headers1.containsKey(GrpclbLoadBalancer.TOKEN_KEY));
    PickSubchannelArgs args2 = mock(PickSubchannelArgs.class);
    Metadata headers2 = new Metadata();
    // The existing token on the headers will be replaced
    headers2.put(GrpclbLoadBalancer.TOKEN_KEY, "LBTOKEN__OLD");
    when(args2.getHeaders()).thenReturn(headers2);
    assertSame(r2.result, picker.pickSubchannel(args2));
    verify(args2).getHeaders();
    assertThat(headers2.getAll(GrpclbLoadBalancer.TOKEN_KEY)).containsExactly("LBTOKEN0001");
    PickSubchannelArgs args3 = mock(PickSubchannelArgs.class);
    Metadata headers3 = new Metadata();
    when(args3.getHeaders()).thenReturn(headers3);
    assertSame(r3.result, picker.pickSubchannel(args3));
    verify(args3).getHeaders();
    assertThat(headers3.getAll(GrpclbLoadBalancer.TOKEN_KEY)).containsExactly("LBTOKEN0002");
    PickSubchannelArgs args4 = mock(PickSubchannelArgs.class);
    Metadata headers4 = new Metadata();
    when(args4.getHeaders()).thenReturn(headers4);
    assertSame(r1.result, picker.pickSubchannel(args4));
    verify(args4).getHeaders();
    assertFalse(headers4.containsKey(GrpclbLoadBalancer.TOKEN_KEY));
    verify(subchannel, never()).getAttributes();
}

Example 17 with PickSubchannelArgs

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

the class GrpclbLoadBalancerTest method errorPicker.

@Test
public void errorPicker() {
    PickSubchannelArgs mockArgs = mock(PickSubchannelArgs.class);
    Status error = Status.UNAVAILABLE.withDescription("Just don't know why");
    ErrorPicker picker = new ErrorPicker(error);
    assertSame(error, picker.pickSubchannel(mockArgs).getStatus());
    verifyNoMoreInteractions(mockArgs);
}

Example 18 with PickSubchannelArgs

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

the class RingHashLoadBalancerTest method subchannelLazyConnectUntilPicked.

@Test
public void subchannelLazyConnectUntilPicked() {
    RingHashConfig config = new RingHashConfig(10, 100);
    // one server
    List<EquivalentAddressGroup> servers = createWeightedServerAddrs(1);
    loadBalancer.handleResolvedAddresses(ResolvedAddresses.newBuilder().setAddresses(servers).setLoadBalancingPolicyConfig(config).build());
    verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
    Subchannel subchannel = Iterables.getOnlyElement(subchannels.values());
    verify(subchannel, never()).requestConnection();
    verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
    // Picking subchannel triggers connection.
    PickSubchannelArgs args = new PickSubchannelArgsImpl(TestMethodDescriptors.voidMethod(), new Metadata(), CallOptions.DEFAULT.withOption(XdsNameResolver.RPC_HASH_KEY, hashFunc.hashVoid()));
    PickResult result = pickerCaptor.getValue().pickSubchannel(args);
    assertThat(result.getStatus().isOk()).isTrue();
    assertThat(result.getSubchannel()).isNull();
    verify(subchannel).requestConnection();
    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(CONNECTING));
    verify(helper).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class));
    // Subchannel becomes ready, triggers pick again.
    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
    verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
    result = pickerCaptor.getValue().pickSubchannel(args);
    assertThat(result.getSubchannel()).isSameInstanceAs(subchannel);
    verifyNoMoreInteractions(helper);
}

Example 19 with PickSubchannelArgs

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

the class RingHashLoadBalancerTest method deterministicPickWithHostsPartiallyRemoved.

@Test
public void deterministicPickWithHostsPartiallyRemoved() {
    RingHashConfig config = new RingHashConfig(10, 100);
    List<EquivalentAddressGroup> servers = createWeightedServerAddrs(1, 1, 1, 1, 1);
    loadBalancer.handleResolvedAddresses(ResolvedAddresses.newBuilder().setAddresses(servers).setLoadBalancingPolicyConfig(config).build());
    InOrder inOrder = Mockito.inOrder(helper);
    inOrder.verify(helper, times(5)).createSubchannel(any(CreateSubchannelArgs.class));
    inOrder.verify(helper).updateBalancingState(eq(IDLE), any(SubchannelPicker.class));
    // Bring all subchannels to READY so that next pick always succeeds.
    for (Subchannel subchannel : subchannels.values()) {
        deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
        inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
    }
    // Simulate rpc hash hits one ring entry exactly for server1.
    long rpcHash = hashFunc.hashAsciiString("[FakeSocketAddress-server1]_0");
    PickSubchannelArgs args = new PickSubchannelArgsImpl(TestMethodDescriptors.voidMethod(), new Metadata(), CallOptions.DEFAULT.withOption(XdsNameResolver.RPC_HASH_KEY, rpcHash));
    pickerCaptor.getValue().pickSubchannel(args);
    PickResult result = pickerCaptor.getValue().pickSubchannel(args);
    Subchannel subchannel = result.getSubchannel();
    assertThat(subchannel.getAddresses()).isEqualTo(servers.get(1));
    List<EquivalentAddressGroup> updatedServers = new ArrayList<>();
    for (EquivalentAddressGroup addr : servers.subList(0, 2)) {
        // only server0 and server1 left
        Attributes attr = addr.getAttributes().toBuilder().set(CUSTOM_KEY, "custom value").build();
        updatedServers.add(new EquivalentAddressGroup(addr.getAddresses(), attr));
    }
    loadBalancer.handleResolvedAddresses(ResolvedAddresses.newBuilder().setAddresses(updatedServers).setLoadBalancingPolicyConfig(config).build());
    verify(subchannels.get(Collections.singletonList(servers.get(0)))).updateAddresses(Collections.singletonList(updatedServers.get(0)));
    verify(subchannels.get(Collections.singletonList(servers.get(1)))).updateAddresses(Collections.singletonList(updatedServers.get(1)));
    inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
    assertThat(pickerCaptor.getValue().pickSubchannel(args).getSubchannel()).isSameInstanceAs(subchannel);
    verifyNoMoreInteractions(helper);
}

Example 20 with PickSubchannelArgs

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

the class RingHashLoadBalancerTest method nameResolutionErrorWithActiveSubchannels.

@Test
public void nameResolutionErrorWithActiveSubchannels() {
    RingHashConfig config = new RingHashConfig(10, 100);
    List<EquivalentAddressGroup> servers = createWeightedServerAddrs(1);
    loadBalancer.handleResolvedAddresses(ResolvedAddresses.newBuilder().setAddresses(servers).setLoadBalancingPolicyConfig(config).build());
    verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
    verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
    // Picking subchannel triggers subchannel creation and connection.
    PickSubchannelArgs args = new PickSubchannelArgsImpl(TestMethodDescriptors.voidMethod(), new Metadata(), CallOptions.DEFAULT.withOption(XdsNameResolver.RPC_HASH_KEY, hashFunc.hashVoid()));
    pickerCaptor.getValue().pickSubchannel(args);
    deliverSubchannelState(Iterables.getOnlyElement(subchannels.values()), ConnectivityStateInfo.forNonError(READY));
    verify(helper).updateBalancingState(eq(READY), any(SubchannelPicker.class));
    loadBalancer.handleNameResolutionError(Status.NOT_FOUND.withDescription("target not found"));
    verifyNoMoreInteractions(helper);
}