[linkedin#555] Response-only Streaming

zizhong · Mar 15, 2021 · 38430cd · 38430cd
1 parent 4f5c93d
commit 38430cd
Show file tree

Hide file tree

Showing 29 changed files with 463 additions and 13 deletions.
diff --git a/d2-test-api/src/main/java/com/linkedin/d2/balancer/clients/TestClient.java b/d2-test-api/src/main/java/com/linkedin/d2/balancer/clients/TestClient.java
@@ -102,6 +102,24 @@ public void restRequest(RestRequest request,
     callback.onResponse(TransportResponseImpl.success(response));
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
+    restRequest = request;
+    restRequestContext = requestContext;
+    restWireAttrs = wireAttrs;
+    streamCallback = callback;
+    StreamResponseBuilder builder = new StreamResponseBuilder();
+    StreamResponse response = _emptyResponse ? builder.build(EntityStreams.emptyStream())
+        : builder.build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy("This is not empty".getBytes()))));
+    if (_deferCallback)
+    {
+      scheduleTimeout(requestContext, callback);
+      return;
+    }
+    callback.onResponse(TransportResponseImpl.success(response, wireAttrs));
+  }
+
   @Override
   public void streamRequest(StreamRequest request,
       RequestContext requestContext,

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/D2ClientDelegator.java b/d2/src/main/java/com/linkedin/d2/balancer/D2ClientDelegator.java
@@ -93,6 +93,17 @@ public void streamRequest(StreamRequest request, RequestContext requestContext,
     _d2Client.streamRequest(request, requestContext, callback);
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, Callback<StreamResponse> callback) {
+    _d2Client.restRequestStreamResponse(request, callback);
+  }
+
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Callback<StreamResponse> callback) {
+    _d2Client.restRequestStreamResponse(request, requestContext, callback);
+  }
+
   @Override
   public void getMetadata(URI uri, Callback<Map<String, Object>> callback)
   {

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/BackupRequestsClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/BackupRequestsClient.java
@@ -390,6 +390,24 @@ public void streamRequest(StreamRequest request, RequestContext requestContext,
         decorateCallbackSync(request, requestContext, _d2Client::streamRequest, callback));
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, Callback<StreamResponse> callback) {
+    restRequestStreamResponse(request, new RequestContext(), callback);
+  }
+
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Callback<StreamResponse> callback) {
+    if (_isD2Async)
+    {
+      requestAsync(request, requestContext, _d2Client::restRequestStreamResponse, callback);
+      return;
+    }
+
+    _d2Client.restRequestStreamResponse(request, requestContext,
+        decorateCallbackSync(request, requestContext, _d2Client::restRequestStreamResponse, callback));
+  }
+
   private <R extends Request, T> Callback<T> decorateCallbackSync(R request, RequestContext requestContext,
       DecoratorClient<R, T> client, Callback<T> callback)
   {

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/DynamicClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/DynamicClient.java
@@ -107,7 +107,16 @@ public void streamRequest(StreamRequest request,
     _balancer.getClient(request, requestContext,
       getClientCallback(request, requestContext, true, callback, client -> client.streamRequest(request, requestContext, loggerCallback))
     );
+  }
+
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Callback<StreamResponse> callback) {
+    Callback<StreamResponse> loggerCallback = decorateLoggingCallback(callback, request, "stream");
 
+    _balancer.getClient(request, requestContext,
+        getClientCallback(request, requestContext, true, callback, client -> client.restRequestStreamResponse(request, requestContext, loggerCallback))
+    );
   }
 
   private Callback<TransportClient> getClientCallback(Request request, RequestContext requestContext, final boolean restOverStream, Callback<? extends Response> callback, SuccessCallback<Client> clientSuccessCallback)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/LazyClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/LazyClient.java
@@ -60,6 +60,12 @@ public void restRequest(RestRequest request,
     getWrappedClient().restRequest(request, requestContext, wireAttrs, callback);
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
+    getWrappedClient().restRequestStreamResponse(request, requestContext, wireAttrs, callback);
+  }
+
   @Override
   public void streamRequest(StreamRequest request,
                           RequestContext requestContext,

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/RequestTimeoutClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/RequestTimeoutClient.java
@@ -120,6 +120,20 @@ public void streamRequest(StreamRequest request, RequestContext requestContext,
     _d2Client.streamRequest(request, requestContext, transportCallback);
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, Callback<StreamResponse> callback) {
+    restRequestStreamResponse(request, new RequestContext(), callback);
+  }
+
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Callback<StreamResponse> callback) {
+    final Callback<StreamResponse> transportCallback =
+        decorateCallbackWithRequestTimeout(callback, request, requestContext);
+
+    _d2Client.restRequestStreamResponse(request, requestContext, transportCallback);
+  }
+
   /**
    * Enforces the user timeout to the layer below if necessary.
    *

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/RetryClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/RetryClient.java
@@ -186,6 +186,29 @@ public void streamRequest(StreamRequest request, RequestContext requestContext,
     }
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, Callback<StreamResponse> callback) {
+    restRequestStreamResponse(request, new RequestContext(), callback);
+  }
+
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Callback<StreamResponse> callback) {
+    if (_restRetryEnabled)
+    {
+      RestRequest newRequest = request.builder()
+          .setHeader(HttpConstants.HEADER_NUMBER_OF_RETRY_ATTEMPTS, "0")
+          .build();
+      ClientRetryTracker retryTracker = updateRetryTracker(newRequest.getURI(), false);
+      final Callback<StreamResponse> transportCallback = new ResponseOnlyStreamRetryRequestCallback(newRequest, requestContext, callback, retryTracker);
+      _d2Client.restRequestStreamResponse(newRequest, requestContext, transportCallback);
+    }
+    else
+    {
+      _d2Client.restRequestStreamResponse(request, requestContext, callback);
+    }
+  }
+
   private ClientRetryTracker updateRetryTracker(URI uri, boolean isRetry)
   {
     String serviceName = LoadBalancerUtil.getServiceNameFromUri(uri);
@@ -276,6 +299,28 @@ public boolean doRetryRequest(RestRequest request, RequestContext context, int n
     }
   }
 
+  /**
+   * Callback implementation for Retry {@link RestRequest} and {@link StreamResponse}
+   */
+  private class ResponseOnlyStreamRetryRequestCallback extends RetryRequestCallback<RestRequest, StreamResponse>
+  {
+    public ResponseOnlyStreamRetryRequestCallback(RestRequest request, RequestContext context, Callback<StreamResponse> callback, ClientRetryTracker retryTracker)
+    {
+      super(request, context, callback, retryTracker);
+    }
+
+    @Override
+    public boolean doRetryRequest(RestRequest request, RequestContext context, int numberOfRetryAttempts)
+    {
+      RestRequest newRequest = request.builder()
+          .setHeader(HttpConstants.HEADER_NUMBER_OF_RETRY_ATTEMPTS, Integer.toString(numberOfRetryAttempts))
+          .build();
+      updateRetryTracker(request.getURI(), true);
+      _d2Client.restRequestStreamResponse(newRequest, context, this);
+      return true;
+    }
+  }
+
   /**
    * Abstract callback implementation of retry requests.
    *

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/RewriteClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/RewriteClient.java
@@ -60,6 +60,14 @@ public void restRequest(RestRequest request, RequestContext requestContext, Map<
     _transportClient.restRequest(rewriteRequest(request), requestContext, wireAttrs, callback);
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request,
+      RequestContext requestContext,
+      Map<String, String> wireAttrs,
+      TransportCallback<StreamResponse> callback) {
+    _transportClient.restRequestStreamResponse(rewriteRequest(request), requestContext, wireAttrs, callback);
+  }
+
   /**
    * Asynchronously issues the given request. The given callback is invoked when the response is
    * received.

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/RewriteLoadBalancerClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/RewriteLoadBalancerClient.java
@@ -74,6 +74,13 @@ public void streamRequest(StreamRequest request,
     _client.streamRequest(request, requestContext, wireAttrs, callback);
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
+    assert _serviceName.equals(LoadBalancerUtil.getServiceNameFromUri(request.getURI()));
+    _client.restRequestStreamResponse(request, requestContext, wireAttrs, callback);
+  }
+
   @Override
   public void shutdown(Callback<None> callback)
   {

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClientImpl.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClientImpl.java
@@ -132,6 +132,12 @@ public void restRequest(RestRequest request,
     _transportClient.restRequest(request, requestContext, wireAttrs, new TrackerClientRestCallback(callback, _callTracker.startCall()));
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
+    _transportClient.restRequestStreamResponse(request, requestContext, wireAttrs, new TrackerClientStreamCallback(callback, _callTracker.startCall()));
+  }
+
   @Override
   public void streamRequest(StreamRequest request,
                             RequestContext requestContext,

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/simple/ClusterAwareTransportClient.java b/d2/src/main/java/com/linkedin/d2/balancer/simple/ClusterAwareTransportClient.java
@@ -72,6 +72,13 @@ public void restRequest(RestRequest request,
     getWrappedClient().restRequest(request, requestContext, wireAttrs, callback);
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
+    updateRequestContext(requestContext);
+    getWrappedClient().restRequestStreamResponse(request, requestContext, wireAttrs, callback);
+  }
+
   @Override
   public void streamRequest(StreamRequest request,
                           RequestContext requestContext,

diff --git a/d2/src/test/java/com/linkedin/d2/balancer/clients/LazyClientTest.java b/d2/src/test/java/com/linkedin/d2/balancer/clients/LazyClientTest.java
@@ -134,6 +134,12 @@ public void restRequest(RestRequest request,
           ++requestCount;
         }
 
+        @Override
+        public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+            Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
+          ++requestCount;
+        }
+
         @Override
         public void shutdown(Callback<None> callback)
         {

diff --git a/d2/src/test/java/com/linkedin/d2/balancer/clients/RetryTrackerClient.java b/d2/src/test/java/com/linkedin/d2/balancer/clients/RetryTrackerClient.java
@@ -74,6 +74,14 @@ public void streamRequest(StreamRequest request,
         () -> new StreamResponseBuilder().build(EntityStreams.emptyStream()));
   }
 
+  @Override
+  public void restRequestStreamResponse(RestRequest request, RequestContext requestContext,
+      Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
+    handleRequest(request, wireAttrs, callback,
+        r -> {},
+        () -> new StreamResponseBuilder().build(EntityStreams.emptyStream()));
+  }
+
   @Override
   public URI getUri()
   {