xds: make ControlPlaneClient.shutdown() synchronous on syncContext

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

@@ -114,21 +114,21 @@ final class ControlPlaneClient {
     logger.log(XdsLogLevel.INFO, "Created");
   }
 
+  // Must be synchronized.
   void shutdown() {
-    syncContext.execute(new Runnable() {
-      @Override
-      public void run() {
-        shutdown = true;
-        logger.log(XdsLogLevel.INFO, "Shutting down");
-        if (adsStream != null) {
-          adsStream.close(Status.CANCELLED.withDescription("shutdown").asException());
-        }
-        if (rpcRetryTimer != null && rpcRetryTimer.isPending()) {
-          rpcRetryTimer.cancel();
-        }
-        xdsTransport.shutdown();
-      }
-    });
+    syncContext.throwIfNotInThisSynchronizationContext();
+    if (shutdown) {
+      return;
+    }
+    shutdown = true;
+    logger.log(XdsLogLevel.INFO, "Shutting down");
+    if (adsStream != null) {
+      adsStream.close(Status.CANCELLED.withDescription("shutdown").asException());
+    }
+    if (rpcRetryTimer != null && rpcRetryTimer.isPending()) {
+      rpcRetryTimer.cancel();
+    }
+    xdsTransport.shutdown();
   }
 
   @Override

xds/src/main/java/io/grpc/xds/client/XdsClientImpl.java

@@ -574,6 +574,11 @@ private <T extends ResourceUpdate> void handleResourceUpdate(
       XdsResourceType.Args args, List<Any> resources, XdsResourceType<T> xdsResourceType,
       boolean isFirstResponse, ProcessingTracker processingTracker) {
     ControlPlaneClient controlPlaneClient = serverCpClientMap.get(args.serverInfo);
+    if (controlPlaneClient == null) {
+      logger.log(XdsLogLevel.DEBUG,
+          "Couldn't find CPC for {0}, so skipping resource update", args.serverInfo);
+      return;
+    }
 
     if (isFirstResponse) {
       shutdownLowerPriorityCpcs(controlPlaneClient);

xds/src/test/java/io/grpc/xds/XdsClientFallbackTest.java

@@ -22,15 +22,20 @@
 import static org.mockito.AdditionalAnswers.delegatesTo;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.protobuf.Any;
+import io.envoyproxy.envoy.service.discovery.v3.DiscoveryRequest;
+import io.envoyproxy.envoy.service.discovery.v3.DiscoveryResponse;
 import io.grpc.ChannelCredentials;
 import io.grpc.Grpc;
 import io.grpc.MetricRecorder;
@@ -591,6 +596,112 @@ public void used_then_mainServerRestart_fallbackServerUp() {
     assertThat(getLrsServerInfo("localhost:" + fallbackServer.getServer().getPort())).isNull();
   }
 
+  @Test
+  public void concurrentResponsesOnMainAndFallback_shutsDownFallbackWithoutNpe() {
+    @SuppressWarnings("unchecked")
+    XdsTransportFactory.StreamingCall<DiscoveryRequest, DiscoveryResponse> mainCall1 =
+        mock(XdsTransportFactory.StreamingCall.class);
+    @SuppressWarnings("unchecked")
+    XdsTransportFactory.StreamingCall<DiscoveryRequest, DiscoveryResponse> mainCall2 =
+        mock(XdsTransportFactory.StreamingCall.class);
+    @SuppressWarnings("unchecked")
+    XdsTransportFactory.StreamingCall<DiscoveryRequest, DiscoveryResponse> fallbackCall =
+        mock(XdsTransportFactory.StreamingCall.class);
+    when(mainCall1.isReady()).thenReturn(true);
+    when(mainCall2.isReady()).thenReturn(true);
+    when(fallbackCall.isReady()).thenReturn(true);
+
+    XdsTransportFactory.XdsTransport mainTransport = mock(XdsTransportFactory.XdsTransport.class);
+    XdsTransportFactory.XdsTransport fallbackTransport =
+        mock(XdsTransportFactory.XdsTransport.class);
+    when(mainTransport.<DiscoveryRequest, DiscoveryResponse>createStreamingCall(
+            any(), any(), any()))
+        .thenReturn(mainCall1)
+        .thenReturn(mainCall2);
+    when(fallbackTransport.<DiscoveryRequest, DiscoveryResponse>createStreamingCall(
+            any(), any(), any()))
+        .thenReturn(fallbackCall);
+
+    XdsTransportFactory customTransportFactory =
+        serverInfo ->
+            serverInfo.target().equals("main-uri") ? mainTransport : fallbackTransport;
+
+    XdsClientImpl client =
+        CommonBootstrapperTestUtils.createXdsClient(
+            Arrays.asList("main-uri", "fallback-uri"),
+            customTransportFactory,
+            fakeClock,
+            new ExponentialBackoffPolicy.Provider(),
+            MessagePrinter.INSTANCE,
+            xdsClientMetricReporter);
+
+    client.watchXdsResource(XdsListenerResource.getInstance(), MAIN_SERVER, ldsWatcher);
+
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<XdsTransportFactory.EventHandler<DiscoveryResponse>> mainHandler1Captor =
+        ArgumentCaptor.forClass(XdsTransportFactory.EventHandler.class);
+    verify(mainCall1).start(mainHandler1Captor.capture());
+
+    // Fail main's initial ADS stream before any response, activating fallback.
+    mainHandler1Captor.getValue().onStatusReceived(Status.UNAVAILABLE);
+
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<XdsTransportFactory.EventHandler<DiscoveryResponse>> fallbackHandlerCaptor =
+        ArgumentCaptor.forClass(XdsTransportFactory.EventHandler.class);
+    verify(fallbackCall).start(fallbackHandlerCaptor.capture());
+    fallbackHandlerCaptor.getValue().onReady();
+
+    // Advance clock so main's retry timer fires and opens a second ADS stream while fallback's
+    // stream is also open.
+    fakeClock.forwardTime(2, TimeUnit.SECONDS);
+
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<XdsTransportFactory.EventHandler<DiscoveryResponse>> mainHandler2Captor =
+        ArgumentCaptor.forClass(XdsTransportFactory.EventHandler.class);
+    verify(mainCall2).start(mainHandler2Captor.capture());
+
+    DiscoveryResponse mainResponse =
+        DiscoveryResponse.newBuilder()
+            .setTypeUrl(ADS_TYPE_URL_LDS)
+            .setVersionInfo("1")
+            .setNonce("nonce-main")
+            .addResources(Any.pack(ControlPlaneRule.buildClientListener(MAIN_SERVER, RDS_NAME)))
+            .build();
+    DiscoveryResponse fallbackResponse =
+        DiscoveryResponse.newBuilder()
+            .setTypeUrl(ADS_TYPE_URL_LDS)
+            .setVersionInfo("1")
+            .setNonce("nonce-fallback")
+            .addResources(
+                Any.pack(ControlPlaneRule.buildClientListener(MAIN_SERVER, FALLBACK_RDS_NAME)))
+            .build();
+
+    // When mainCall2 sends its initial DiscoveryRequest (inside XdsClientImpl's syncContext),
+    // enqueue onRecvMessage for both main and fallback before syncContext finishes draining.
+    // Main's first response runs shutdownLowerPriorityCpcs(mainCpc), and fallback's already-queued
+    // onRecvMessage runs immediately after in the same drain loop.
+    doAnswer(
+            invocation -> {
+              mainHandler2Captor.getValue().onRecvMessage(mainResponse);
+              fallbackHandlerCaptor.getValue().onRecvMessage(fallbackResponse);
+              return null;
+            })
+        .doNothing()
+        .when(mainCall2)
+        .sendMessage(any());
+
+    mainHandler2Captor.getValue().onReady();
+
+    verify(ldsWatcher, timeout(5000))
+        .onResourceChanged(
+            StatusOr.fromValue(LdsUpdate.forApiListener(MAIN_HTTP_CONNECTION_MANAGER)));
+    verify(fallbackCall).sendError(any());
+    verify(fallbackTransport).shutdown();
+    verify(fallbackCall, times(1)).sendMessage(any());
+
+    client.shutdown();
+  }
+
   /**
    * gRFC A95: for an authority with fallback_on_reachability_only, falling back must happen when
    * the primary server becomes unreachable, even though every subscribed resource is cached.

xds/src/test/java/io/grpc/xds/client/ControlPlaneClientTest.java

@@ -273,6 +273,38 @@ public void cancelBeforeAck_sendsEmptyUnsubscribe() {
     assertThat(unsub.getResourceNamesList()).isEmpty();
   }
 
+  @Test
+  public void shutdown_closesStreamImmediatelyAndSkipsQueuedRecvMessage() {
+    Map<String, XdsResourceType<?>> subscribedTypes = ImmutableMap.of(CDS_TYPE_URL, cdsType);
+    when(resourceStore.getSubscribedResourceTypesWithTypeUrl()).thenReturn(subscribedTypes);
+    when(resourceStore.getSubscribedResources(serverInfo, cdsType))
+        .thenReturn(ImmutableList.of("foo-cluster"));
+
+    syncContext.execute(cpc::sendDiscoveryRequests);
+    verify(streamingCall).start(handlerCaptor.capture());
+    handlerCaptor.getValue().onReady();
+
+    DiscoveryResponse response =
+        DiscoveryResponse.newBuilder()
+            .setTypeUrl(CDS_TYPE_URL)
+            .setVersionInfo("1")
+            .setNonce("nonce-1")
+            .build();
+
+    // Simulate a response arriving while syncContext is draining a task that shuts down this CPC
+    // (e.g. XdsClientImpl.shutdownLowerPriorityCpcs when a higher-priority server responds first).
+    syncContext.execute(
+        () -> {
+          handlerCaptor.getValue().onRecvMessage(response);
+          cpc.shutdown();
+        });
+
+    verify(streamingCall).sendError(any());
+    verify(xdsTransport).shutdown();
+    verify(responseHandler, never())
+        .handleResourceResponse(any(), any(), any(), any(), any(), any(Boolean.class), any());
+  }
+
   private static DiscoveryRequest argThatTypeUrlIs(String typeUrl) {
     return argThat(req -> req != null && typeUrl.equals(req.getTypeUrl()));
   }