server: handle call cancellation gracefully (#80)

sschepens · snowp · commit c56c848b5d6b · 2018-12-02T13:38:46.000-06:00
Immediately cancel watches when responseObserver is cancelled. Catch cancellation exceptions when responding watches. Correctly handle cancellations on error. Fixes #79 Signed-off-by: Sebastian Schepens <sebastian.schepens@mercadolibre.com>
diff --git a/server/src/main/java/io/envoyproxy/controlplane/server/DiscoveryServer.java b/server/src/main/java/io/envoyproxy/controlplane/server/DiscoveryServer.java
@@ -17,6 +17,8 @@
 import io.envoyproxy.controlplane.cache.Response;
 import io.envoyproxy.controlplane.cache.Watch;
 import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
 import java.util.Collections;
 import java.util.List;
@@ -152,123 +154,158 @@ private StreamObserver<DiscoveryRequest> createRequestHandler(
 
     callbacks.forEach(cb -> cb.onStreamOpen(streamId, defaultTypeUrl));
 
-    return new StreamObserver<DiscoveryRequest>() {
+    final DiscoveryRequestStreamObserver requestStreamObserver =
+        new DiscoveryRequestStreamObserver(defaultTypeUrl, responseObserver, streamId, ads);
 
-      private final Map<String, Watch> watches = new ConcurrentHashMap<>(Resources.TYPE_URLS.size());
-      private final Map<String, DiscoveryResponse> latestResponse =
-          new ConcurrentHashMap<>(Resources.TYPE_URLS.size());
-      private final Map<String, Set<String>> ackedResources = new ConcurrentHashMap<>(Resources.TYPE_URLS.size());
+    if (responseObserver instanceof ServerCallStreamObserver) {
+      ((ServerCallStreamObserver) responseObserver).setOnCancelHandler(requestStreamObserver::onCancelled);
+    }
 
-      private AtomicLong streamNonce = new AtomicLong();
+    return requestStreamObserver;
+  }
 
-      @Override
-      public void onNext(DiscoveryRequest request) {
-        String nonce = request.getResponseNonce();
-        String requestTypeUrl = request.getTypeUrl();
-
-        if (defaultTypeUrl.equals(ANY_TYPE_URL)) {
-          if (requestTypeUrl.isEmpty()) {
-            responseObserver.onError(
-                Status.UNKNOWN
-                    .withDescription(String.format("[%d] type URL is required for ADS", streamId))
-                    .asRuntimeException());
-
-            return;
-          }
-        } else if (requestTypeUrl.isEmpty()) {
-          requestTypeUrl = defaultTypeUrl;
+  private class DiscoveryRequestStreamObserver implements StreamObserver<DiscoveryRequest> {
+
+    private final Map<String, Watch> watches;
+    private final Map<String, DiscoveryResponse> latestResponse;
+    private final Map<String, Set<String>> ackedResources;
+    private final String defaultTypeUrl;
+    private final StreamObserver<DiscoveryResponse> responseObserver;
+    private final long streamId;
+    private final boolean ads;
+
+    private AtomicLong streamNonce;
+
+    public DiscoveryRequestStreamObserver(String defaultTypeUrl, StreamObserver<DiscoveryResponse> responseObserver,
+                                          long streamId, boolean ads) {
+      this.defaultTypeUrl = defaultTypeUrl;
+      this.responseObserver = responseObserver;
+      this.streamId = streamId;
+      this.ads = ads;
+      watches = new ConcurrentHashMap<>(Resources.TYPE_URLS.size());
+      latestResponse = new ConcurrentHashMap<>(Resources.TYPE_URLS.size());
+      ackedResources = new ConcurrentHashMap<>(Resources.TYPE_URLS.size());
+      streamNonce = new AtomicLong();
+    }
+
+    @Override
+    public void onNext(DiscoveryRequest request) {
+      String nonce = request.getResponseNonce();
+      String requestTypeUrl = request.getTypeUrl();
+
+      if (defaultTypeUrl.equals(ANY_TYPE_URL)) {
+        if (requestTypeUrl.isEmpty()) {
+          responseObserver.onError(
+              Status.UNKNOWN
+                  .withDescription(String.format("[%d] type URL is required for ADS", streamId))
+                  .asRuntimeException());
+
+          return;
         }
+      } else if (requestTypeUrl.isEmpty()) {
+        requestTypeUrl = defaultTypeUrl;
+      }
 
-        LOGGER.info("[{}] request {}[{}] with nonce {} from version {}",
-            streamId,
-            requestTypeUrl,
-            String.join(", ", request.getResourceNamesList()),
-            nonce,
-            request.getVersionInfo());
-
-        callbacks.forEach(cb -> cb.onStreamRequest(streamId, request));
-
-        for (String typeUrl : Resources.TYPE_URLS) {
-          DiscoveryResponse response = latestResponse.get(typeUrl);
-          String resourceNonce = response == null ? null : response.getNonce();
-
-          if (requestTypeUrl.equals(typeUrl) && (isNullOrEmpty(resourceNonce)
-              || resourceNonce.equals(nonce))) {
-            if (!request.hasErrorDetail() && response != null) {
-              Set<String> ackedResourcesForType = response.getResourcesList()
-                  .stream()
-                  .map(Resources::getResourceName)
-                  .collect(Collectors.toSet());
-              ackedResources.put(typeUrl, ackedResourcesForType);
-            }
+      LOGGER.info("[{}] request {}[{}] with nonce {} from version {}",
+          streamId,
+          requestTypeUrl,
+          String.join(", ", request.getResourceNamesList()),
+          nonce,
+          request.getVersionInfo());
+
+      callbacks.forEach(cb -> cb.onStreamRequest(streamId, request));
+
+      for (String typeUrl : Resources.TYPE_URLS) {
+        DiscoveryResponse response = latestResponse.get(typeUrl);
+        String resourceNonce = response == null ? null : response.getNonce();
+
+        if (requestTypeUrl.equals(typeUrl) && (isNullOrEmpty(resourceNonce)
+            || resourceNonce.equals(nonce))) {
+          if (!request.hasErrorDetail() && response != null) {
+            Set<String> ackedResourcesForType = response.getResourcesList()
+                .stream()
+                .map(Resources::getResourceName)
+                .collect(Collectors.toSet());
+            ackedResources.put(typeUrl, ackedResourcesForType);
+          }
 
-            watches.compute(typeUrl, (t, oldWatch) -> {
-              if (oldWatch != null) {
-                oldWatch.cancel();
-              }
+          watches.compute(typeUrl, (t, oldWatch) -> {
+            if (oldWatch != null) {
+              oldWatch.cancel();
+            }
 
-              return configWatcher.createWatch(
-                  ads,
-                  request,
-                  ackedResources.getOrDefault(typeUrl, Collections.emptySet()),
-                  r -> send(r, typeUrl));
-            });
+            return configWatcher.createWatch(
+                ads,
+                request,
+                ackedResources.getOrDefault(typeUrl, Collections.emptySet()),
+                r -> send(r, typeUrl));
+          });
 
-            return;
-          }
+          return;
         }
       }
+    }
 
-      @Override
-      public void onError(Throwable t) {
-        if (!Status.fromThrowable(t).getCode().equals(Status.CANCELLED.getCode())) {
-          LOGGER.error("[{}] stream closed with error", streamId, t);
-        }
-
-        try {
-          callbacks.forEach(cb -> cb.onStreamCloseWithError(streamId, defaultTypeUrl, t));
-          responseObserver.onError(Status.fromThrowable(t).asException());
-        } finally {
-          cancel();
-        }
+    @Override
+    public void onError(Throwable t) {
+      if (!Status.fromThrowable(t).getCode().equals(Status.CANCELLED.getCode())) {
+        LOGGER.error("[{}] stream closed with error", streamId, t);
       }
 
-      @Override
-      public void onCompleted() {
-        LOGGER.info("[{}] stream closed", streamId);
-
-        try {
-          callbacks.forEach(cb -> cb.onStreamClose(streamId, defaultTypeUrl));
-          responseObserver.onCompleted();
-        } finally {
-          cancel();
-        }
+      try {
+        callbacks.forEach(cb -> cb.onStreamCloseWithError(streamId, defaultTypeUrl, t));
+        responseObserver.onError(Status.fromThrowable(t).asException());
+      } finally {
+        cancel();
       }
+    }
+
+    @Override
+    public void onCompleted() {
+      LOGGER.info("[{}] stream closed", streamId);
 
-      private void cancel() {
-        watches.values().forEach(Watch::cancel);
+      try {
+        callbacks.forEach(cb -> cb.onStreamClose(streamId, defaultTypeUrl));
+        responseObserver.onCompleted();
+      } finally {
+        cancel();
       }
+    }
 
-      private void send(Response response, String typeUrl) {
-        String nonce = Long.toString(streamNonce.getAndIncrement());
+    void onCancelled() {
+      LOGGER.info("[{}] stream cancelled", streamId);
+      cancel();
+    }
 
-        DiscoveryResponse discoveryResponse = DiscoveryResponse.newBuilder()
-            .setVersionInfo(response.version())
-            .addAllResources(response.resources().stream().map(Any::pack).collect(Collectors.toList()))
-            .setTypeUrl(typeUrl)
-            .setNonce(nonce)
-            .build();
+    private void cancel() {
+      watches.values().forEach(Watch::cancel);
+    }
 
-        LOGGER.info("[{}] response {} with nonce {} version {}", streamId, typeUrl, nonce, response.version());
+    private void send(Response response, String typeUrl) {
+      String nonce = Long.toString(streamNonce.getAndIncrement());
 
-        callbacks.forEach(cb -> cb.onStreamResponse(streamId, response.request(), discoveryResponse));
+      DiscoveryResponse discoveryResponse = DiscoveryResponse.newBuilder()
+          .setVersionInfo(response.version())
+          .addAllResources(response.resources().stream().map(Any::pack).collect(Collectors.toList()))
+          .setTypeUrl(typeUrl)
+          .setNonce(nonce)
+          .build();
 
-        // Store the latest response *before* we send the response. This ensures that by the time the request
-        // is processed the map is guaranteed to be updated. Doing it afterwards leads to a race conditions
-        // which may see the incoming request arrive before the map is updated, failing the nonce check erroneously.
-        latestResponse.put(typeUrl, discoveryResponse);
+      LOGGER.info("[{}] response {} with nonce {} version {}", streamId, typeUrl, nonce, response.version());
+
+      callbacks.forEach(cb -> cb.onStreamResponse(streamId, response.request(), discoveryResponse));
+
+      // Store the latest response *before* we send the response. This ensures that by the time the request
+      // is processed the map is guaranteed to be updated. Doing it afterwards leads to a race conditions
+      // which may see the incoming request arrive before the map is updated, failing the nonce check erroneously.
+      latestResponse.put(typeUrl, discoveryResponse);
+      try {
         responseObserver.onNext(discoveryResponse);
+      } catch (StatusRuntimeException e) {
+        if (!Status.CANCELLED.getCode().equals(e.getStatus().getCode())) {
+          throw e;
+        }
       }
-    };
+    }
   }
 }
diff --git a/server/src/test/java/io/envoyproxy/controlplane/server/DiscoveryServerTest.java b/server/src/test/java/io/envoyproxy/controlplane/server/DiscoveryServerTest.java
@@ -804,6 +804,67 @@ public void callbackOnError_doesNotLogError_whenCancelled() {
     }
   }
 
+  @Test
+  public void testCallbacksOnCancelled() throws InterruptedException, ClassNotFoundException {
+    final CountDownLatch streamCloseWithErrorLatch = new CountDownLatch(1);
+    final CountDownLatch watchCreated = new CountDownLatch(1);
+    AtomicReference<Watch> watch = new AtomicReference<>();
+
+    MockDiscoveryServerCallbacks callbacks = new MockDiscoveryServerCallbacks() {
+      @Override
+      public void onStreamCloseWithError(long streamId, String typeUrl, Throwable error) {
+        // watch should already be closed by the time we report a stream close error
+        assertThat(watch.get().isCancelled()).isTrue();
+        super.onStreamCloseWithError(streamId, typeUrl, error);
+        streamCloseWithErrorLatch.countDown();
+      }
+    };
+
+    MockConfigWatcher configWatcher = new MockConfigWatcher(false, createResponses()) {
+      @Override
+      public Watch createWatch(boolean ads, DiscoveryRequest request, Set<String> knownResources,
+                               Consumer<Response> responseConsumer) {
+        watchCreated.countDown();
+        watch.set(super.createWatch(ads, request, knownResources, responseConsumer));
+        return watch.get();
+      }
+    };
+    DiscoveryServer server = new DiscoveryServer(callbacks, configWatcher);
+
+    grpcServer.getServiceRegistry().addService(server.getClusterDiscoveryServiceImpl());
+
+    ClusterDiscoveryServiceStub stub = ClusterDiscoveryServiceGrpc.newStub(grpcServer.getChannel());
+
+    MockDiscoveryResponseObserver responseObserver = new MockDiscoveryResponseObserver();
+
+    StreamObserver<DiscoveryRequest> requestObserver = stub.streamClusters(responseObserver);
+
+    requestObserver.onNext(DiscoveryRequest.newBuilder()
+        .setNode(NODE)
+        .setResponseNonce("1")
+        .setTypeUrl(Resources.CLUSTER_TYPE_URL)
+        .setVersionInfo(VERSION)
+        .build());
+
+    if (!watchCreated.await(1, TimeUnit.SECONDS)) {
+      fail("failed to execute watchCreated callback before timeout");
+    }
+
+    requestObserver.onError(Status.CANCELLED.asException());
+
+    if (!streamCloseWithErrorLatch.await(1, TimeUnit.SECONDS)) {
+      fail("failed to execute onStreamCloseWithError callback before timeout");
+    }
+
+    callbacks.assertThatNoErrors();
+
+    assertThat(callbacks.streamCloseCount).hasValue(0);
+    assertThat(callbacks.streamCloseWithErrorCount).hasValue(1);
+    assertThat(callbacks.streamOpenCount).hasValue(1);
+    assertThat(callbacks.streamRequestCount).hasValue(1);
+    assertThat(callbacks.streamResponseCount).hasValue(1);
+  }
+
   private static Table<String, String, Collection<? extends Message>> createResponses() {
     return ImmutableTable.<String, String, Collection<? extends Message>>builder()
         .put(Resources.CLUSTER_TYPE_URL, VERSION, ImmutableList.of(CLUSTER))
