Java Code Examples for io.grpc.LoadBalancer.PickResult#withSubchannel()

/**
 * For PICK_FIRST: creates a BackendEntry that includes all addresses.
 */
BackendEntry(Subchannel subchannel, TokenAttachingTracerFactory tracerFactory) {
  this.subchannel = checkNotNull(subchannel, "subchannel");
  this.result =
      PickResult.withSubchannel(subchannel, checkNotNull(tracerFactory, "tracerFactory"));
  this.token = null;
}
 

private static SubchannelPicker pickerOf(final Subchannel subchannel) {
  return new SubchannelPicker() {
    @Override
    public PickResult pickSubchannel(PickSubchannelArgs args) {
      return PickResult.withSubchannel(subchannel);
    }
  };
}
 

private static void deliverSubchannelState(
    final NoopSubchannel subchannel, ConnectivityState state) {
  SubchannelPicker picker = new SubchannelPicker() {
    @Override
    public PickResult pickSubchannel(PickSubchannelArgs args) {
      return PickResult.withSubchannel(subchannel);
    }
  };
  subchannel.helper.updateBalancingState(state, picker);
}
 

@Override
public PickResult pickSubchannel(PickSubchannelArgs args) {
  PickResult result = delegate().pickSubchannel(args);
  if (!result.getStatus().isOk()) {
    return result;
  }
  if (result.getSubchannel() == null) {
    return result;
  }
  ClientStreamTracer.Factory originFactory = result.getStreamTracerFactory();
  if (originFactory == null) {
    originFactory = NOOP_CLIENT_STREAM_TRACER_FACTORY;
  }
  return PickResult.withSubchannel(result.getSubchannel(), wrapTracerFactory(originFactory));
}
 

@Test
public void pickResult_equals() {
  PickResult sc1 = PickResult.withSubchannel(subchannel);
  PickResult sc2 = PickResult.withSubchannel(subchannel);
  PickResult sc3 = PickResult.withSubchannel(subchannel, tracerFactory);
  PickResult sc4 = PickResult.withSubchannel(subchannel2);
  PickResult nr = PickResult.withNoResult();
  PickResult error1 = PickResult.withError(status);
  PickResult error2 = PickResult.withError(status2);
  PickResult error3 = PickResult.withError(status2);
  PickResult drop1 = PickResult.withDrop(status);
  PickResult drop2 = PickResult.withDrop(status);
  PickResult drop3 = PickResult.withDrop(status2);

  assertThat(sc1).isNotEqualTo(nr);
  assertThat(sc1).isNotEqualTo(error1);
  assertThat(sc1).isNotEqualTo(drop1);
  assertThat(sc1).isEqualTo(sc2);
  assertThat(sc1).isNotEqualTo(sc3);
  assertThat(sc1).isNotEqualTo(sc4);

  assertThat(error1).isNotEqualTo(error2);
  assertThat(error2).isEqualTo(error3);

  assertThat(drop1).isEqualTo(drop2);
  assertThat(drop1).isNotEqualTo(drop3);

  assertThat(error1.getStatus()).isEqualTo(drop1.getStatus());
  assertThat(error1).isNotEqualTo(drop1);
}
 

@Test
public void pickResult_withSubchannelAndTracer() {
  PickResult result = PickResult.withSubchannel(subchannel, tracerFactory);
  assertThat(result.getSubchannel()).isSameInstanceAs(subchannel);
  assertThat(result.getStatus()).isSameInstanceAs(Status.OK);
  assertThat(result.getStreamTracerFactory()).isSameInstanceAs(tracerFactory);
  assertThat(result.isDrop()).isFalse();
}
 

@Test
public void pickResult_withSubchannel() {
  PickResult result = PickResult.withSubchannel(subchannel);
  assertThat(result.getSubchannel()).isSameInstanceAs(subchannel);
  assertThat(result.getStatus()).isSameInstanceAs(Status.OK);
  assertThat(result.getStreamTracerFactory()).isNull();
  assertThat(result.isDrop()).isFalse();
}
 

void deliverSubchannelState(final Subchannel subchannel, ConnectivityState state) {
  SubchannelPicker picker = new SubchannelPicker() {
    @Override
    public PickResult pickSubchannel(PickSubchannelArgs args) {
      return PickResult.withSubchannel(subchannel);
    }
  };
  helper.updateBalancingState(state, picker);
}
 

/**
 * For ROUND_ROBIN: creates a BackendEntry whose usage will be reported to load recorder.
 */
BackendEntry(Subchannel subchannel, GrpclbClientLoadRecorder loadRecorder, String token) {
  this.subchannel = checkNotNull(subchannel, "subchannel");
  this.result =
      PickResult.withSubchannel(subchannel, checkNotNull(loadRecorder, "loadRecorder"));
  this.token = checkNotNull(token, "token");
}
 

@Override
public PickResult pickSubchannel(PickSubchannelArgs args) {
  Map<String, Object> affinity =
      args.getCallOptions().getOption(GrpcCallOptions.CALLOPTIONS_CUSTOME_KEY);
  GrpcURL refUrl = (GrpcURL) affinity.get(GrpcCallOptions.GRPC_REF_URL);
  if (size > 0) {
    Subchannel subchannel = nextSubchannel(refUrl);
    affinity.put(GrpcCallOptions.GRPC_NAMERESOVER_ATTRIBUTES, nameResovleCache);
    return PickResult.withSubchannel(subchannel);
  }
  if (status != null) {
    return PickResult.withError(status);
  }
  return PickResult.withNoResult();
}
 

@Test
public void pickResult_equals() {
  PickResult sc1 = PickResult.withSubchannel(subchannel);
  PickResult sc2 = PickResult.withSubchannel(subchannel);
  PickResult sc3 = PickResult.withSubchannel(subchannel, tracerFactory);
  PickResult sc4 = PickResult.withSubchannel(subchannel2);
  PickResult nr = PickResult.withNoResult();
  PickResult error1 = PickResult.withError(status);
  PickResult error2 = PickResult.withError(status2);
  PickResult error3 = PickResult.withError(status2);
  PickResult drop1 = PickResult.withDrop(status);
  PickResult drop2 = PickResult.withDrop(status);
  PickResult drop3 = PickResult.withDrop(status2);

  assertThat(sc1).isNotEqualTo(nr);
  assertThat(sc1).isNotEqualTo(error1);
  assertThat(sc1).isNotEqualTo(drop1);
  assertThat(sc1).isEqualTo(sc2);
  assertThat(sc1).isNotEqualTo(sc3);
  assertThat(sc1).isNotEqualTo(sc4);

  assertThat(error1).isNotEqualTo(error2);
  assertThat(error2).isEqualTo(error3);

  assertThat(drop1).isEqualTo(drop2);
  assertThat(drop1).isNotEqualTo(drop3);

  assertThat(error1.getStatus()).isEqualTo(drop1.getStatus());
  assertThat(error1).isNotEqualTo(drop1);
}
 

@Test
public void pickResult_withSubchannelAndTracer() {
  PickResult result = PickResult.withSubchannel(subchannel, tracerFactory);
  assertThat(result.getSubchannel()).isSameAs(subchannel);
  assertThat(result.getStatus()).isSameAs(Status.OK);
  assertThat(result.getStreamTracerFactory()).isSameAs(tracerFactory);
  assertThat(result.isDrop()).isFalse();
}
 

@Test
public void pickResult_withSubchannel() {
  PickResult result = PickResult.withSubchannel(subchannel);
  assertThat(result.getSubchannel()).isSameAs(subchannel);
  assertThat(result.getStatus()).isSameAs(Status.OK);
  assertThat(result.getStreamTracerFactory()).isNull();
  assertThat(result.isDrop()).isFalse();
}
 

@Override
public PickResult pickSubchannel(PickSubchannelArgs args) {
  Subchannel subchannel = null;
  if (stickinessState != null) {
    String stickinessValue = args.getHeaders().get(stickinessState.key);
    if (stickinessValue != null) {
      subchannel = stickinessState.getSubchannel(stickinessValue);
      if (subchannel == null || !RoundRobinLoadBalancer.isReady(subchannel)) {
        subchannel = stickinessState.maybeRegister(stickinessValue, nextSubchannel());
      }
    }
  }

  return PickResult.withSubchannel(subchannel != null ? subchannel : nextSubchannel());
}
 

/**
 * Creates a BackendEntry whose usage will be reported to load recorder.
 */
BackendEntry(Subchannel subchannel, GrpclbClientLoadRecorder loadRecorder, String token) {
  this.result = PickResult.withSubchannel(subchannel, loadRecorder);
  this.loadRecorder = checkNotNull(loadRecorder, "loadRecorder");
  this.token = checkNotNull(token, "token");
}
 

void setSubchannel(final InternalSubchannel subchannel) {
  log.log(Level.FINE, "[{0}] Created with [{1}]", new Object[] {this, subchannel});
  this.subchannel = subchannel;
  subchannelImpl = new AbstractSubchannel() {
      @Override
      public void shutdown() {
        subchannel.shutdown(Status.UNAVAILABLE.withDescription("OobChannel is shutdown"));
      }

      @Override
      ClientTransport obtainActiveTransport() {
        return subchannel.obtainActiveTransport();
      }

      @Override
      InternalInstrumented<ChannelStats> getInternalSubchannel() {
        return subchannel;
      }

      @Override
      public void requestConnection() {
        subchannel.obtainActiveTransport();
      }

      @Override
      public List<EquivalentAddressGroup> getAllAddresses() {
        return subchannel.getAddressGroups();
      }

      @Override
      public Attributes getAttributes() {
        return Attributes.EMPTY;
      }
  };

  subchannelPicker = new SubchannelPicker() {
      final PickResult result = PickResult.withSubchannel(subchannelImpl);

      @Override
      public PickResult pickSubchannel(PickSubchannelArgs args) {
        return result;
      }
    };
  delayedTransport.reprocess(subchannelPicker);
}
 

@Test
public void edsResourceUpdate_localityAssignmentChange() {
  deliverResolvedAddresses(null, null, fakeEndpointPickingPolicy);

  LbEndpoint endpoint11 = buildLbEndpoint("addr11.example.com", 8011, HEALTHY, 11);
  LbEndpoint endpoint12 = buildLbEndpoint("addr12.example.com", 8012, HEALTHY, 12);
  LocalityLbEndpoints localityLbEndpoints1 = buildLocalityLbEndpoints(
      "region1", "zone1", "subzone1",
      ImmutableList.of(endpoint11, endpoint12),
      1,
      0);

  LbEndpoint endpoint21 = buildLbEndpoint("addr21.example.com", 8021, HEALTHY, 21);
  LbEndpoint endpoint22 = buildLbEndpoint("addr22.example.com", 8022, HEALTHY, 22);
  LocalityLbEndpoints localityLbEndpoints2 = buildLocalityLbEndpoints(
      "region2", "zone2", "subzone2",
      ImmutableList.of(endpoint21, endpoint22),
      2,
      0);

  LbEndpoint endpoint31 = buildLbEndpoint("addr31.example.com", 8031, HEALTHY, 31);
  LocalityLbEndpoints localityLbEndpoints3 = buildLocalityLbEndpoints(
      "region3", "zone3", "subzone3",
      ImmutableList.of(endpoint31),
      3,
      0);

  ClusterLoadAssignment clusterLoadAssignment = buildClusterLoadAssignment(
      CLUSTER_NAME,
      ImmutableList.of(localityLbEndpoints1, localityLbEndpoints2, localityLbEndpoints3),
      ImmutableList.<DropOverload>of());
  deliverClusterLoadAssignments(clusterLoadAssignment);

  assertThat(childBalancers).hasSize(3);
  verify(childBalancers.get("subzone1")).handleResolvedAddresses(
      argThat(RoundRobinBackendsMatcher.builder()
          .addHostAndPort("addr11.example.com", 8011)
          .addHostAndPort("addr12.example.com", 8012)
          .build()));
  verify(childBalancers.get("subzone2")).handleResolvedAddresses(
      argThat(RoundRobinBackendsMatcher.builder()
          .addHostAndPort("addr21.example.com", 8021)
          .addHostAndPort("addr22.example.com", 8022)
          .build()));
  verify(childBalancers.get("subzone3")).handleResolvedAddresses(
      argThat(RoundRobinBackendsMatcher.builder()
          .addHostAndPort("addr31.example.com", 8031)
          .build()));
  assertThat(childHelpers).hasSize(3);
  Helper childHelper2 = childHelpers.get("subzone2");
  final Subchannel subchannel = mock(Subchannel.class);
  SubchannelPicker picker = new SubchannelPicker() {
    @Override
    public PickResult pickSubchannel(PickSubchannelArgs args) {
      return PickResult.withSubchannel(subchannel);
    }
  };
  verify(helper, never()).updateBalancingState(eq(READY), any(SubchannelPicker.class));
  childHelper2.updateBalancingState(READY, picker);
  assertLatestSubchannelPicker(subchannel);
}
 

@Test
public void edsResourceRemoved() {
  deliverResolvedAddresses(null, null, fakeEndpointPickingPolicy);
  ClusterLoadAssignment clusterLoadAssignment =
      buildClusterLoadAssignment(CLUSTER_NAME,
          ImmutableList.of(
              buildLocalityLbEndpoints("region", "zone", "subzone",
                  ImmutableList.of(
                      buildLbEndpoint("192.168.0.1", 8080, HEALTHY, 2)),
                  1, 0)),
          ImmutableList.<DropOverload>of());
  deliverClusterLoadAssignments(clusterLoadAssignment);

  assertThat(childBalancers).hasSize(1);
  assertThat(childHelpers).hasSize(1);
  LoadBalancer localityBalancer = childBalancers.get("subzone");
  Helper localityBalancerHelper = childHelpers.get("subzone");
  final Subchannel subchannel = mock(Subchannel.class);
  SubchannelPicker picker = new SubchannelPicker() {
    @Override
    public PickResult pickSubchannel(PickSubchannelArgs args) {
      return PickResult.withSubchannel(subchannel);
    }
  };
  localityBalancerHelper.updateBalancingState(READY, picker);
  verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
  PickResult result = pickerCaptor.getValue().pickSubchannel(mock(PickSubchannelArgs.class));
  assertThat(result.getSubchannel()).isSameInstanceAs(subchannel);

  // The whole cluster is no longer accessible.
  // Note that EDS resource removal is achieved by CDS resource update.
  responseObserver.onNext(
      buildDiscoveryResponse(
          String.valueOf(versionIno++),
          Collections.<Any>emptyList(),
          XdsClientImpl.ADS_TYPE_URL_CDS,
          String.valueOf(nonce++)));

  verify(localityBalancer).shutdown();
  verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
  result = pickerCaptor.getValue().pickSubchannel(mock(PickSubchannelArgs.class));
  assertThat(result.getStatus().getCode()).isEqualTo(Code.UNAVAILABLE);
  assertThat(result.getStatus().getDescription())
      .isEqualTo("Resource " + CLUSTER_NAME + " is unavailable");
}
 

@Test
public void updateLocalityStore_pickResultInterceptedForLoadRecordingWhenSubchannelReady() {
  // Simulate receiving two localities.
  LocalityLbEndpoints localityInfo1 =
      new LocalityLbEndpoints(ImmutableList.of(lbEndpoint11, lbEndpoint12), 1, 0);
  LocalityLbEndpoints localityInfo2 =
      new LocalityLbEndpoints(ImmutableList.of(lbEndpoint21, lbEndpoint22), 2, 0);
  localityStore.updateLocalityStore(ImmutableMap.of(
      locality1, localityInfo1, locality2, localityInfo2));

  // Two child balancers are created.
  assertThat(loadBalancers).hasSize(2);

  ClientStreamTracer.Factory metricsTracingFactory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory metricsTracingFactory2 = mock(ClientStreamTracer.Factory.class);
  when(orcaPerRequestUtil.newOrcaClientStreamTracerFactory(any(ClientStreamTracer.Factory.class),
      any(OrcaPerRequestReportListener.class)))
      .thenReturn(metricsTracingFactory1, metricsTracingFactory2);

  Subchannel subchannel1 = mock(Subchannel.class);
  Subchannel subchannel2 = mock(Subchannel.class);
  final PickResult result1 = PickResult.withSubchannel(subchannel1);
  final PickResult result2 =
      PickResult.withSubchannel(subchannel2, mock(ClientStreamTracer.Factory.class));
  SubchannelPicker subchannelPicker1 = mock(SubchannelPicker.class);
  SubchannelPicker subchannelPicker2 = mock(SubchannelPicker.class);
  when(subchannelPicker1.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(result1);
  when(subchannelPicker2.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(result2);

  // Simulate picker updates for the two localities with dummy pickers.
  childHelpers.get("sz1").updateBalancingState(READY, subchannelPicker1);
  childHelpers.get("sz2").updateBalancingState(READY, subchannelPicker2);

  ArgumentCaptor<SubchannelPicker> interLocalityPickerCaptor = ArgumentCaptor.forClass(null);
  verify(helper, times(2)).updateBalancingState(eq(READY), interLocalityPickerCaptor.capture());
  WeightedRandomPicker interLocalityPicker =
      (WeightedRandomPicker) interLocalityPickerCaptor.getValue();
  assertThat(interLocalityPicker.weightedChildPickers).hasSize(2);

  // Verify each PickResult picked is intercepted with client stream tracer factory for
  // recording load and backend metrics.
  Map<Subchannel, Locality> localitiesBySubchannel
      = ImmutableMap.of(subchannel1, locality1, subchannel2, locality2);
  Map<Subchannel, ClientStreamTracer.Factory> metricsTracingFactoriesBySubchannel
      = ImmutableMap.of(subchannel1, metricsTracingFactory1, subchannel2, metricsTracingFactory2);
  for (int i = 0; i < interLocalityPicker.weightedChildPickers.size(); i++) {
    PickResult pickResult = interLocalityPicker.weightedChildPickers.get(i).getPicker()
        .pickSubchannel(pickSubchannelArgs);
    Subchannel expectedSubchannel = pickResult.getSubchannel();
    Locality expectedLocality = localitiesBySubchannel.get(expectedSubchannel);
    ArgumentCaptor<OrcaPerRequestReportListener> listenerCaptor = ArgumentCaptor.forClass(null);
    verify(orcaPerRequestUtil, times(i + 1))
        .newOrcaClientStreamTracerFactory(any(ClientStreamTracer.Factory.class),
            listenerCaptor.capture());
    assertThat(listenerCaptor.getValue()).isInstanceOf(MetricsRecordingListener.class);
    MetricsRecordingListener listener = (MetricsRecordingListener) listenerCaptor.getValue();
    assertThat(listener.getCounter())
        .isSameInstanceAs(fakeLoadStatsStore.localityCounters.get(expectedLocality));
    assertThat(pickResult.getStreamTracerFactory())
        .isInstanceOf(LoadRecordingStreamTracerFactory.class);
    LoadRecordingStreamTracerFactory loadRecordingFactory =
        (LoadRecordingStreamTracerFactory) pickResult.getStreamTracerFactory();
    assertThat(loadRecordingFactory.getCounter())
        .isSameInstanceAs(fakeLoadStatsStore.localityCounters.get(expectedLocality));
    assertThat(loadRecordingFactory.delegate())
        .isSameInstanceAs(metricsTracingFactoriesBySubchannel.get(expectedSubchannel));
  }
}
 

/**
 * Creates a BackendEntry whose usage will not be reported.
 */
BackendEntry(Subchannel subchannel) {
  this.result = PickResult.withSubchannel(subchannel);
  this.loadRecorder = null;
  this.token = null;
}