Java Code Examples for io.grpc.Attributes#Key

String getRemoteAddress() {
    Attributes attributes = serverStream.getAttributes();
    if (attributes == null) {
        return null;
    }

    try {
        // keys method is being considered for removal,
        Set<Attributes.Key<?>> keys = attributes.keys();
        if (keys == null) {
            if (isDebug) {
                logger.debug("can't attributes keys");
            }
            return null;
        }

        for (Attributes.Key<?> key : keys) {
            if (key != null && key.toString().equals("remote-addr")) {
                Object remoteAddress = attributes.get(key);
                if (remoteAddress instanceof SocketAddress) {
                    return getSocketAddressAsString((SocketAddress) remoteAddress);
                } else if (remoteAddress instanceof String) {
                    return (String) remoteAddress;
                }
            }
        }
    } catch (Exception e) {
        if (isDebug) {
            logger.debug("can't find keys method");
        }
    }

    return GrpcConstants.UNKNOWN_ADDRESS;
}
 

@Test
public void filterAddresses() {
  Attributes.Key<String> key1 = Attributes.Key.create("key1");
  Attributes attributes1 = Attributes.newBuilder().set(key1, "value1").build();
  EquivalentAddressGroup eag0 = new EquivalentAddressGroup(new InetSocketAddress(8000));
  EquivalentAddressGroup eag1 =
      new EquivalentAddressGroup(new InetSocketAddress(8001), attributes1);
  EquivalentAddressGroup eag2 = new EquivalentAddressGroup(new InetSocketAddress(8002));
  EquivalentAddressGroup eag3 =
      new EquivalentAddressGroup(
          Arrays.<SocketAddress>asList(new InetSocketAddress(8003), new InetSocketAddress(8083)));
  eag0 = AddressFilter.setPathFilter(eag0, Arrays.asList("A", "C"));
  eag1 = AddressFilter.setPathFilter(eag1, Arrays.asList("A", "B"));
  eag2 = AddressFilter.setPathFilter(eag2, Arrays.asList("D", "C"));
  eag3 = AddressFilter.setPathFilter(eag3, Arrays.asList("A", "B"));

  List<EquivalentAddressGroup> addresses =
      AddressFilter.filter(Arrays.asList(eag0, eag1, eag2, eag3), "A");
  assertThat(addresses).hasSize(3);
  addresses = AddressFilter.filter(addresses, "B");
  assertThat(addresses).hasSize(2);
  EquivalentAddressGroup filteredAddress0 = addresses.get(0);
  EquivalentAddressGroup filteredAddress1 = addresses.get(1);
  assertThat(filteredAddress0.getAddresses()).containsExactlyElementsIn(eag1.getAddresses());
  assertThat(filteredAddress0.getAttributes().get(key1)).isEqualTo("value1");
  assertThat(filteredAddress1.getAddresses()).containsExactlyElementsIn(eag3.getAddresses());
}
 

@Test public void index_looping() {
  Attributes.Key<String> key = Attributes.Key.create("some-key");
  Attributes attr1 = Attributes.newBuilder().set(key, "1").build();
  Attributes attr2 = Attributes.newBuilder().set(key, "2").build();
  Attributes attr3 = Attributes.newBuilder().set(key, "3").build();
  SocketAddress addr1 = new FakeSocketAddress();
  SocketAddress addr2 = new FakeSocketAddress();
  SocketAddress addr3 = new FakeSocketAddress();
  SocketAddress addr4 = new FakeSocketAddress();
  SocketAddress addr5 = new FakeSocketAddress();
  Index index = new Index(Arrays.asList(
      new EquivalentAddressGroup(Arrays.asList(addr1, addr2), attr1),
      new EquivalentAddressGroup(Arrays.asList(addr3), attr2),
      new EquivalentAddressGroup(Arrays.asList(addr4, addr5), attr3)));
  assertThat(index.getCurrentAddress()).isSameAs(addr1);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr1);
  assertThat(index.isAtBeginning()).isTrue();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameAs(addr2);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr1);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameAs(addr3);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr2);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameAs(addr4);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr3);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameAs(addr5);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr3);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isFalse();

  index.reset();
  assertThat(index.getCurrentAddress()).isSameAs(addr1);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr1);
  assertThat(index.isAtBeginning()).isTrue();
  assertThat(index.isValid()).isTrue();

  // We want to make sure both groupIndex and addressIndex are reset
  index.increment();
  index.increment();
  index.increment();
  index.increment();
  assertThat(index.getCurrentAddress()).isSameAs(addr5);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr3);
  index.reset();
  assertThat(index.getCurrentAddress()).isSameAs(addr1);
  assertThat(index.getCurrentEagAttributes()).isSameAs(attr1);
}
 

@Test
public void handleResolvedAddresses() {
  ArgumentCaptor<ResolvedAddresses> resolvedAddressesCaptor = ArgumentCaptor.forClass(null);
  Attributes.Key<Object> fakeKey = Attributes.Key.create("fake_key");
  Object fakeValue = new Object();

  Map<String, WeightedPolicySelection> targets = ImmutableMap.of(
      // {foo, 10, config0}
      "target0", weightedLbConfig0,
      // {bar, 20, config1}
      "target1", weightedLbConfig1,
      // {bar, 30, config2}
      "target2", weightedLbConfig2,
      // {foo, 40, config3}
      "target3", weightedLbConfig3);
  EquivalentAddressGroup eag0 = new EquivalentAddressGroup(socketAddresses[0]);
  eag0 = AddressFilter.setPathFilter(eag0, ImmutableList.of("target0"));
  EquivalentAddressGroup eag1 = new EquivalentAddressGroup(socketAddresses[1]);
  eag1 = AddressFilter.setPathFilter(eag1, ImmutableList.of("target1"));
  EquivalentAddressGroup eag2 = new EquivalentAddressGroup(socketAddresses[2]);
  eag2 = AddressFilter.setPathFilter(eag2, ImmutableList.of("target2"));
  EquivalentAddressGroup eag3 = new EquivalentAddressGroup(socketAddresses[3]);
  eag3 = AddressFilter.setPathFilter(eag3, ImmutableList.of("target3"));
  weightedTargetLb.handleResolvedAddresses(
      ResolvedAddresses.newBuilder()
          .setAddresses(ImmutableList.of(eag0, eag1, eag2, eag3))
          .setAttributes(Attributes.newBuilder().set(fakeKey, fakeValue).build())
          .setLoadBalancingPolicyConfig(new WeightedTargetConfig(targets))
          .build());

  assertThat(childBalancers).hasSize(4);
  assertThat(childHelpers).hasSize(4);
  assertThat(fooLbCreated).isEqualTo(2);
  assertThat(barLbCreated).isEqualTo(2);

  for (int i = 0; i < childBalancers.size(); i++) {
    verify(childBalancers.get(i)).handleResolvedAddresses(resolvedAddressesCaptor.capture());
    ResolvedAddresses resolvedAddresses = resolvedAddressesCaptor.getValue();
    assertThat(resolvedAddresses.getLoadBalancingPolicyConfig()).isEqualTo(configs[i]);
    assertThat(resolvedAddresses.getAttributes().get(fakeKey)).isEqualTo(fakeValue);
    assertThat(Iterables.getOnlyElement(resolvedAddresses.getAddresses()).getAddresses())
        .containsExactly(socketAddresses[i]);
  }

  // Update new weighted target config for a typical workflow.
  // target0 removed. target1, target2, target3 changed weight and config. target4 added.
  int[] newWeights = new int[]{11, 22, 33, 44};
  Object[] newConfigs = new Object[]{"newConfig1", "newConfig2", "newConfig3", "newConfig4"};
  Map<String, WeightedPolicySelection> newTargets = ImmutableMap.of(
      "target1",
      new WeightedPolicySelection(
          newWeights[0], new PolicySelection(barLbProvider, null, newConfigs[0])),
      "target2",
      new WeightedPolicySelection(
          newWeights[1], new PolicySelection(barLbProvider, null, newConfigs[1])),
      "target3",
      new WeightedPolicySelection(
          newWeights[2], new PolicySelection(fooLbProvider, null, newConfigs[2])),
      "target4",
      new WeightedPolicySelection(
          newWeights[3], new PolicySelection(fooLbProvider, null, newConfigs[3])));
  weightedTargetLb.handleResolvedAddresses(
      ResolvedAddresses.newBuilder()
          .setAddresses(ImmutableList.<EquivalentAddressGroup>of())
          .setLoadBalancingPolicyConfig(new WeightedTargetConfig(newTargets))
          .build());

  assertThat(childBalancers).hasSize(5);
  assertThat(childHelpers).hasSize(5);
  assertThat(fooLbCreated).isEqualTo(3); // One more foo LB created for target4
  assertThat(barLbCreated).isEqualTo(2);

  verify(childBalancers.get(0)).shutdown();
  for (int i = 1; i < childBalancers.size(); i++) {
    verify(childBalancers.get(i), atLeastOnce())
        .handleResolvedAddresses(resolvedAddressesCaptor.capture());
    assertThat(resolvedAddressesCaptor.getValue().getLoadBalancingPolicyConfig())
        .isEqualTo(newConfigs[i - 1]);
  }
}
 

@Test public void index_looping() {
  Attributes.Key<String> key = Attributes.Key.create("some-key");
  Attributes attr1 = Attributes.newBuilder().set(key, "1").build();
  Attributes attr2 = Attributes.newBuilder().set(key, "2").build();
  Attributes attr3 = Attributes.newBuilder().set(key, "3").build();
  SocketAddress addr1 = new FakeSocketAddress();
  SocketAddress addr2 = new FakeSocketAddress();
  SocketAddress addr3 = new FakeSocketAddress();
  SocketAddress addr4 = new FakeSocketAddress();
  SocketAddress addr5 = new FakeSocketAddress();
  Index index = new Index(Arrays.asList(
      new EquivalentAddressGroup(Arrays.asList(addr1, addr2), attr1),
      new EquivalentAddressGroup(Arrays.asList(addr3), attr2),
      new EquivalentAddressGroup(Arrays.asList(addr4, addr5), attr3)));
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr1);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr1);
  assertThat(index.isAtBeginning()).isTrue();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr2);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr1);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr3);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr2);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr4);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr3);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr5);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr3);
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isTrue();

  index.increment();
  assertThat(index.isAtBeginning()).isFalse();
  assertThat(index.isValid()).isFalse();

  index.reset();
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr1);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr1);
  assertThat(index.isAtBeginning()).isTrue();
  assertThat(index.isValid()).isTrue();

  // We want to make sure both groupIndex and addressIndex are reset
  index.increment();
  index.increment();
  index.increment();
  index.increment();
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr5);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr3);
  index.reset();
  assertThat(index.getCurrentAddress()).isSameInstanceAs(addr1);
  assertThat(index.getCurrentEagAttributes()).isSameInstanceAs(attr1);
}
 

@Test
public void pickAfterResolvedUpdatedHosts() throws Exception {
  Subchannel removedSubchannel = mock(Subchannel.class);
  Subchannel oldSubchannel = mock(Subchannel.class);
  Subchannel newSubchannel = mock(Subchannel.class);

  Attributes.Key<String> key = Attributes.Key.create("check-that-it-is-propagated");
  FakeSocketAddress removedAddr = new FakeSocketAddress("removed");
  EquivalentAddressGroup removedEag = new EquivalentAddressGroup(removedAddr);
  FakeSocketAddress oldAddr = new FakeSocketAddress("old");
  EquivalentAddressGroup oldEag1 = new EquivalentAddressGroup(oldAddr);
  EquivalentAddressGroup oldEag2 = new EquivalentAddressGroup(
      oldAddr, Attributes.newBuilder().set(key, "oldattr").build());
  FakeSocketAddress newAddr = new FakeSocketAddress("new");
  EquivalentAddressGroup newEag = new EquivalentAddressGroup(
      newAddr, Attributes.newBuilder().set(key, "newattr").build());

  subchannels.put(Collections.singletonList(removedEag), removedSubchannel);
  subchannels.put(Collections.singletonList(oldEag1), oldSubchannel);
  subchannels.put(Collections.singletonList(newEag), newSubchannel);

  List<EquivalentAddressGroup> currentServers = Lists.newArrayList(removedEag, oldEag1);

  InOrder inOrder = inOrder(mockHelper);

  loadBalancer.handleResolvedAddresses(
      ResolvedAddresses.newBuilder().setAddresses(currentServers).setAttributes(affinity)
          .build());

  inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());

  deliverSubchannelState(removedSubchannel, ConnectivityStateInfo.forNonError(READY));
  deliverSubchannelState(oldSubchannel, ConnectivityStateInfo.forNonError(READY));

  inOrder.verify(mockHelper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());
  SubchannelPicker picker = pickerCaptor.getValue();
  assertThat(getList(picker)).containsExactly(removedSubchannel, oldSubchannel);

  verify(removedSubchannel, times(1)).requestConnection();
  verify(oldSubchannel, times(1)).requestConnection();

  assertThat(loadBalancer.getSubchannels()).containsExactly(removedSubchannel,
      oldSubchannel);

  // This time with Attributes
  List<EquivalentAddressGroup> latestServers = Lists.newArrayList(oldEag2, newEag);

  loadBalancer.handleResolvedAddresses(
      ResolvedAddresses.newBuilder().setAddresses(latestServers).setAttributes(affinity).build());

  verify(newSubchannel, times(1)).requestConnection();
  verify(oldSubchannel, times(1)).updateAddresses(Arrays.asList(oldEag2));
  verify(removedSubchannel, times(1)).shutdown();

  deliverSubchannelState(removedSubchannel, ConnectivityStateInfo.forNonError(SHUTDOWN));
  deliverSubchannelState(newSubchannel, ConnectivityStateInfo.forNonError(READY));

  assertThat(loadBalancer.getSubchannels()).containsExactly(oldSubchannel,
      newSubchannel);

  verify(mockHelper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
  inOrder.verify(mockHelper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());

  picker = pickerCaptor.getValue();
  assertThat(getList(picker)).containsExactly(oldSubchannel, newSubchannel);

  // test going from non-empty to empty
  loadBalancer.handleResolvedAddresses(
      ResolvedAddresses.newBuilder()
          .setAddresses(Collections.<EquivalentAddressGroup>emptyList())
          .setAttributes(affinity)
          .build());

  inOrder.verify(mockHelper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
  assertEquals(PickResult.withNoResult(), pickerCaptor.getValue().pickSubchannel(mockArgs));

  verifyNoMoreInteractions(mockHelper);
}
 

/**
 * Partition the limit by a request attribute.
 * @return Chainable builder
 */
public GrpcServerLimiterBuilder partitionByAttribute(Attributes.Key<String> attribute) {
    return partitionResolver(context -> context.getCall().getAttributes().get(attribute));
}