use of com.linkedin.d2.balancer.StaticLoadBalancerState in project rest.li by linkedin.
the class TestBackupRequestsClient method createAlwaysBackupClientWithHosts.
private BackupRequestsClient createAlwaysBackupClientWithHosts(List<String> uris, Deque<URI> hostsReceivingRequestList, int responseDelayNano, int backupDelayNano, boolean isD2Async) throws IOException {
Map<URI, Map<Integer, PartitionData>> partitionDescriptions = new HashMap<>();
uris.forEach(uri -> partitionDescriptions.put(URI.create(uri), Collections.singletonMap(0, new PartitionData(1))));
StaticLoadBalancerState LbState = new StaticLoadBalancerState() {
@Override
public TrackerClient getClient(String serviceName, URI uri) {
return new DegraderTrackerClientImpl(uri, partitionDescriptions.get(uri), null, SystemClock.instance(), null) {
@Override
public void restRequest(RestRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<RestResponse> callback) {
// whenever a trackerClient is used to make request, record down it's hostname
hostsReceivingRequestList.add(uri);
// delay response to allow backup request to happen
_executor.schedule(() -> callback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build())), responseDelayNano, TimeUnit.NANOSECONDS);
}
@Override
public void streamRequest(StreamRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
// whenever a trackerClient is used to make request, record down it's hostname
hostsReceivingRequestList.add(uri);
if (null != requestContext.getLocalAttr(R2Constants.BACKUP_REQUEST_BUFFERED_BODY)) {
callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().setHeader(BUFFERED_HEADER, String.valueOf(requestContext.getLocalAttr(R2Constants.BACKUP_REQUEST_BUFFERED_BODY) != null)).build(EntityStreams.emptyStream())));
return;
}
request.getEntityStream().setReader(new DrainReader() {
public void onDone() {
// delay response to allow backup request to happen
_executor.schedule(() -> callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().setHeader(BUFFERED_HEADER, String.valueOf(requestContext.getLocalAttr(R2Constants.BACKUP_REQUEST_BUFFERED_BODY) != null)).build(EntityStreams.emptyStream()))), responseDelayNano, TimeUnit.NANOSECONDS);
}
});
}
};
}
};
LbState.TEST_URIS_PARTITIONDESCRIPTIONS.putAll(partitionDescriptions);
LbState.TEST_SERVICE_BACKUP_REQUEST_PROPERTIES.add(createBackupRequestsConfiguration(5, "get"));
LbState.refreshDefaultProperties();
LoadBalancer loadBalancer = new SimpleLoadBalancer(LbState, _executor);
DynamicClient dynamicClient = new DynamicClient(loadBalancer, null);
return new BackupRequestsClient(dynamicClient, loadBalancer, _executor, null, 10, TimeUnit.SECONDS, isD2Async) {
@Override
Optional<TrackingBackupRequestsStrategy> getStrategyAfterUpdate(final String serviceName, final String operation) {
// constantly enable backup request after backupDelayNano time.
BackupRequestsStrategy alwaysBackup = new TestTrackingBackupRequestsStrategy.MockBackupRequestsStrategy(() -> Optional.of((long) backupDelayNano), () -> true);
return Optional.of(new TrackingBackupRequestsStrategy(alwaysBackup));
}
};
}
Aggregations