xDS: implement ADS stream flow control mechanism

xds/src/main/java/io/grpc/xds/CdsLoadBalancer2.java

@@ -320,7 +320,7 @@
 
       private void start() {
         shutdown = false;
-        xdsClient.watchXdsResource(XdsClusterResource.getInstance(), name, this);
+        xdsClient.watchXdsResource(XdsClusterResource.getInstance(), name, this, syncContext);
       }
 
       void shutdown() {
@@ -341,102 +341,85 @@
                 String.format("Unable to load CDS %s. xDS server returned: %s: %s",
                   name, error.getCode(), error.getDescription()))
             .withCause(error.getCause());
-        syncContext.execute(new Runnable() {
-          @Override
-          public void run() {
-            if (shutdown) {
-              return;
-            }
-            // All watchers should receive the same error, so we only propagate it once.
-            if (ClusterState.this == root) {
-              handleClusterDiscoveryError(status);
-            }
-          }
-        });
+        if (shutdown) {
+          return;
+        }
+        // All watchers should receive the same error, so we only propagate it once.
+        if (ClusterState.this == root) {
+          handleClusterDiscoveryError(status);
+        }
       }
 
       @Override
       public void onResourceDoesNotExist(String resourceName) {
-        syncContext.execute(new Runnable() {
-          @Override
-          public void run() {
-            if (shutdown) {
-              return;
-            }
-            discovered = true;
-            result = null;
-            if (childClusterStates != null) {
-              for (ClusterState state : childClusterStates.values()) {
-                state.shutdown();
-              }
-              childClusterStates = null;
-            }
-            handleClusterDiscovered();
+        if (shutdown) {
+          return;
+        }
+        discovered = true;
+        result = null;
+        if (childClusterStates != null) {
+          for (ClusterState state : childClusterStates.values()) {
+            state.shutdown();
           }
-        });
+          childClusterStates = null;
+        }
+        handleClusterDiscovered();
       }
 
       @Override
       public void onChanged(final CdsUpdate update) {
-        class ClusterDiscovered implements Runnable {
-          @Override
-          public void run() {
-            if (shutdown) {
-              return;
+        if (shutdown) {
+          return;
+        }
+        logger.log(XdsLogLevel.DEBUG, "Received cluster update {0}", update);
+        discovered = true;
+        result = update;
+        if (update.clusterType() == ClusterType.AGGREGATE) {
+          isLeaf = false;
+          logger.log(XdsLogLevel.INFO, "Aggregate cluster {0}, underlying clusters: {1}",
+              update.clusterName(), update.prioritizedClusterNames());
+          Map<String, ClusterState> newChildStates = new LinkedHashMap<>();
+          for (String cluster : update.prioritizedClusterNames()) {
+            if (newChildStates.containsKey(cluster)) {
+              logger.log(XdsLogLevel.WARNING,
+                  String.format("duplicate cluster name %s in aggregate %s is being ignored",
+                      cluster, update.clusterName()));
+              continue;
             }
-
-            logger.log(XdsLogLevel.DEBUG, "Received cluster update {0}", update);
-            discovered = true;
-            result = update;
-            if (update.clusterType() == ClusterType.AGGREGATE) {
-              isLeaf = false;
-              logger.log(XdsLogLevel.INFO, "Aggregate cluster {0}, underlying clusters: {1}",
-                  update.clusterName(), update.prioritizedClusterNames());
-              Map<String, ClusterState> newChildStates = new LinkedHashMap<>();
-              for (String cluster : update.prioritizedClusterNames()) {
-                if (newChildStates.containsKey(cluster)) {
-                  logger.log(XdsLogLevel.WARNING,
-                      String.format("duplicate cluster name %s in aggregate %s is being ignored",
-                          cluster, update.clusterName()));
-                  continue;
-                }
-                if (childClusterStates == null || !childClusterStates.containsKey(cluster)) {
-                  ClusterState childState;
-                  if (clusterStates.containsKey(cluster)) {
-                    childState = clusterStates.get(cluster);
-                    if (childState.shutdown) {
-                      childState.start();
-                    }
-                  } else {
-                    childState = new ClusterState(cluster);
-                    clusterStates.put(cluster, childState);
-                    childState.start();
-                  }
-                  newChildStates.put(cluster, childState);
-                } else {
-                  newChildStates.put(cluster, childClusterStates.remove(cluster));
-                }
-              }
-              if (childClusterStates != null) {  // stop subscribing to revoked child clusters
-                for (ClusterState watcher : childClusterStates.values()) {
-                  watcher.shutdown();
+            if (childClusterStates == null || !childClusterStates.containsKey(cluster)) {
+              ClusterState childState;
+              if (clusterStates.containsKey(cluster)) {
+                childState = clusterStates.get(cluster);
+                if (childState.shutdown) {
+                  childState.start();
                 }
+              } else {
+                childState = new ClusterState(cluster);
+                clusterStates.put(cluster, childState);
+                childState.start();
               }
-              childClusterStates = newChildStates;
-            } else if (update.clusterType() == ClusterType.EDS) {
-              isLeaf = true;
-              logger.log(XdsLogLevel.INFO, "EDS cluster {0}, edsServiceName: {1}",
-                  update.clusterName(), update.edsServiceName());
-            } else {  // logical DNS
-              isLeaf = true;
-              logger.log(XdsLogLevel.INFO, "Logical DNS cluster {0}", update.clusterName());
+              newChildStates.put(cluster, childState);
+            } else {
+              newChildStates.put(cluster, childClusterStates.remove(cluster));
+            }
+          }
+          if (childClusterStates != null) {  // stop subscribing to revoked child clusters
+            for (ClusterState watcher : childClusterStates.values()) {
+              watcher.shutdown();
             }
-            handleClusterDiscovered();
           }
+          childClusterStates = newChildStates;
+        } else if (update.clusterType() == ClusterType.EDS) {
+          isLeaf = true;
+          logger.log(XdsLogLevel.INFO, "EDS cluster {0}, edsServiceName: {1}",
+              update.clusterName(), update.edsServiceName());
+        } else {  // logical DNS
+          isLeaf = true;
+          logger.log(XdsLogLevel.INFO, "Logical DNS cluster {0}", update.clusterName());
         }
-
-        syncContext.execute(new ClusterDiscovered());
+        handleClusterDiscovered();
       }
+
     }
   }
 }

xds/src/main/java/io/grpc/xds/ClusterResolverLoadBalancer.java

@@ -366,7 +366,8 @@
       void start() {
         String resourceName = edsServiceName != null ? edsServiceName : name;
         logger.log(XdsLogLevel.INFO, "Start watching EDS resource {0}", resourceName);
-        xdsClient.watchXdsResource(XdsEndpointResource.getInstance(), resourceName, this);
+        xdsClient.watchXdsResource(XdsEndpointResource.getInstance(),
+            resourceName, this, syncContext);
       }
 
       @Override
@@ -452,7 +453,7 @@
           }
         }
 
-        syncContext.execute(new EndpointsUpdated());
+        new EndpointsUpdated().run();
       }
 
       private List<String> generatePriorityNames(String name,
@@ -491,38 +492,28 @@
 
       @Override
       public void onResourceDoesNotExist(final String resourceName) {
-        syncContext.execute(new Runnable() {
-          @Override
-          public void run() {
-            if (shutdown) {
-              return;
-            }
-            logger.log(XdsLogLevel.INFO, "Resource {0} unavailable", resourceName);
-            status = Status.OK;
-            resolved = true;
-            result = null;  // resource revoked
-            handleEndpointResourceUpdate();
-          }
-        });
+        if (shutdown) {
+          return;
+        }
+        logger.log(XdsLogLevel.INFO, "Resource {0} unavailable", resourceName);
+        status = Status.OK;
+        resolved = true;
+        result = null;  // resource revoked
+        handleEndpointResourceUpdate();
       }
 
       @Override
       public void onError(final Status error) {
-        syncContext.execute(new Runnable() {
-          @Override
-          public void run() {
-            if (shutdown) {
-              return;
-            }
-            String resourceName = edsServiceName != null ? edsServiceName : name;
-            status = Status.UNAVAILABLE
-                .withDescription(String.format("Unable to load EDS %s. xDS server returned: %s: %s",
-                      resourceName, error.getCode(), error.getDescription()))
-                .withCause(error.getCause());
-            logger.log(XdsLogLevel.WARNING, "Received EDS error: {0}", error);
-            handleEndpointResolutionError();
-          }
-        });
+        if (shutdown) {
+          return;
+        }
+        String resourceName = edsServiceName != null ? edsServiceName : name;
+        status = Status.UNAVAILABLE
+            .withDescription(String.format("Unable to load EDS %s. xDS server returned: %s: %s",
+                  resourceName, error.getCode(), error.getDescription()))
+            .withCause(error.getCause());
+        logger.log(XdsLogLevel.WARNING, "Received EDS error: {0}", error);
+        handleEndpointResolutionError();
       }
     }
 

xds/src/main/java/io/grpc/xds/ControlPlaneClient.java

@@ -38,9 +38,9 @@
 import io.grpc.internal.BackoffPolicy;
 import io.grpc.stub.ClientCallStreamObserver;
 import io.grpc.stub.ClientResponseObserver;
-import io.grpc.stub.StreamObserver;
 import io.grpc.xds.Bootstrapper.ServerInfo;
 import io.grpc.xds.EnvoyProtoData.Node;
+import io.grpc.xds.XdsClient.ProcessingTracker;
 import io.grpc.xds.XdsClient.ResourceStore;
 import io.grpc.xds.XdsClient.XdsResponseHandler;
 import io.grpc.xds.XdsClientImpl.XdsChannelFactory;
@@ -288,6 +288,8 @@ private abstract class AbstractAdsStream {
 
     abstract boolean isReady();
 
+    abstract void request(int count);
+
     /**
      * Sends a discovery request with the given {@code versionInfo}, {@code nonce} and
      * {@code errorDetail}. Used for reacting to a specific discovery response. For
@@ -314,7 +316,10 @@ final void handleRpcResponse(XdsResourceType<?> type, String versionInfo, List<A
       }
       responseReceived = true;
       respNonces.put(type, nonce);
-      xdsResponseHandler.handleResourceResponse(type, serverInfo, versionInfo, resources, nonce);
+      ProcessingTracker processingTracker = new ProcessingTracker(() -> request(1), syncContext);
+      xdsResponseHandler.handleResourceResponse(type, serverInfo, versionInfo, resources, nonce,
+          processingTracker);
+      processingTracker.onComplete();
     }
 
     final void handleRpcError(Throwable t) {
@@ -372,14 +377,15 @@ private void cleanUp() {
   }
 
   private final class AdsStreamV3 extends AbstractAdsStream {
-    private StreamObserver<DiscoveryRequest> requestWriter;
+    private ClientCallStreamObserver<DiscoveryRequest> requestWriter;
 
     @Override
     public boolean isReady() {
       return requestWriter != null && ((ClientCallStreamObserver<?>) requestWriter).isReady();
     }
 
     @Override
+    @SuppressWarnings("unchecked")
     void start() {
       AggregatedDiscoveryServiceGrpc.AggregatedDiscoveryServiceStub stub =
           AggregatedDiscoveryServiceGrpc.newStub(channel);
@@ -389,6 +395,7 @@ final class AdsClientResponseObserver
 
         @Override
         public void beforeStart(ClientCallStreamObserver<DiscoveryRequest> requestStream) {
+          requestStream.disableAutoRequestWithInitial(1);
           requestStream.setOnReadyHandler(ControlPlaneClient.this::readyHandler);
         }
 
@@ -437,7 +444,8 @@ public void run() {
         }
       }
 
-      requestWriter = stub.streamAggregatedResources(new AdsClientResponseObserver());
+      requestWriter = (ClientCallStreamObserver) stub.streamAggregatedResources(
+          new AdsClientResponseObserver());
     }
 
     @Override
@@ -467,6 +475,11 @@ void sendDiscoveryRequest(XdsResourceType<?> type, String versionInfo,
       }
     }
 
+    @Override
+    void request(int count) {
+      requestWriter.request(count);
+    }
+
     @Override
     void sendError(Exception error) {
       requestWriter.onError(error);

xds/src/main/java/io/grpc/xds/XdsClient.java

@@ -24,6 +24,7 @@
 import com.google.common.base.Splitter;
 import com.google.common.net.UrlEscapers;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.protobuf.Any;
 import io.grpc.Status;
 import io.grpc.xds.Bootstrapper.ServerInfo;
@@ -36,6 +37,8 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
 
 /**
@@ -305,10 +308,16 @@ TlsContextManager getTlsContextManager() {
    * Registers a data watcher for the given Xds resource.
    */
   <T extends ResourceUpdate> void watchXdsResource(XdsResourceType<T> type, String resourceName,
-                                                   ResourceWatcher<T> watcher) {
+                                                   ResourceWatcher<T> watcher,
+                                                   Executor executor) {
     throw new UnsupportedOperationException();
   }
 
+  <T extends ResourceUpdate> void watchXdsResource(XdsResourceType<T> type, String resourceName,
+                                                   ResourceWatcher<T> watcher) {
+    watchXdsResource(type, resourceName, watcher, MoreExecutors.directExecutor());
+  }
+
   /**
    * Unregisters the given resource watcher.
    */
@@ -353,11 +362,32 @@ Map<ServerInfo, LoadReportClient> getServerLrsClientMap() {
     throw new UnsupportedOperationException();
   }
 
+  static final class ProcessingTracker {
+    private final AtomicInteger pendingTask = new AtomicInteger(1);
+    private final Executor executor;
+    private final Runnable completionListener;
+
+    ProcessingTracker(Runnable completionListener, Executor executor) {
+      this.executor = executor;
+      this.completionListener = completionListener;
+    }
+
+    void startTask() {
+      pendingTask.incrementAndGet();
+    }
+
+    void onComplete() {
+      if (pendingTask.decrementAndGet() == 0) {
+        executor.execute(completionListener);
+      }
+    }
+  }
+
   interface XdsResponseHandler {
     /** Called when a xds response is received. */
     void handleResourceResponse(
         XdsResourceType<?> resourceType, ServerInfo serverInfo, String versionInfo,
-        List<Any> resources, String nonce);
+        List<Any> resources, String nonce, ProcessingTracker processingTracker);
 
     /** Called when the ADS stream is closed passively. */
     // Must be synchronized.