Close remote connections only after all active server calls have completed

service/src/main/java/org/whispersystems/textsecuregcm/grpc/ChannelShutdownInterceptor.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.grpc;
+
+import io.grpc.Context;
+import io.grpc.ForwardingServerCallListener;
+import io.grpc.Grpc;
+import io.grpc.Metadata;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.Status;
+import io.netty.channel.local.LocalAddress;
+import org.whispersystems.textsecuregcm.grpc.net.GrpcClientConnectionManager;
+
+/**
+ * Then channel shutdown interceptor rejects new requests if a channel is shutting down and works in tandem with
+ * {@link GrpcClientConnectionManager} to maintain an active call count for each channel otherwise.
+ */
+public class ChannelShutdownInterceptor implements ServerInterceptor {
+
+  private final GrpcClientConnectionManager grpcClientConnectionManager;
+
+  public ChannelShutdownInterceptor(final GrpcClientConnectionManager grpcClientConnectionManager) {
+    this.grpcClientConnectionManager = grpcClientConnectionManager;
+  }
+
+  @Override
+  public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(final ServerCall<ReqT, RespT> call,
+      final Metadata headers,
+      final ServerCallHandler<ReqT, RespT> next) {
+
+    if (!grpcClientConnectionManager.handleServerCallStart(call)) {
+      // Don't allow new calls if the connection is getting ready to close
+      return ServerInterceptorUtil.closeWithStatus(call, Status.UNAVAILABLE);
+    }
+
+    return new ForwardingServerCallListener.SimpleForwardingServerCallListener<>(next.startCall(call, headers)) {
+      @Override
+      public void onComplete() {
+        grpcClientConnectionManager.handleServerCallComplete(call);
+        super.onComplete();
+      }
+
+      @Override
+      public void onCancel() {
+        grpcClientConnectionManager.handleServerCallComplete(call);
+        super.onCancel();
+      }
+    };
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/GrpcClientConnectionManager.java

@@ -27,6 +27,7 @@ import org.whispersystems.textsecuregcm.auth.DisconnectionRequestListener;
 import org.whispersystems.textsecuregcm.auth.grpc.AuthenticatedDevice;
 import org.whispersystems.textsecuregcm.grpc.ChannelNotFoundException;
 import org.whispersystems.textsecuregcm.grpc.RequestAttributes;
+import org.whispersystems.textsecuregcm.util.ClosableEpoch;
 
 /**
  * A client connection manager associates a local connection to a local gRPC server with a remote connection through a
@@ -58,6 +59,10 @@ public class GrpcClientConnectionManager implements DisconnectionRequestListener
   public static final AttributeKey<RequestAttributes> REQUEST_ATTRIBUTES_KEY =
       AttributeKey.valueOf(GrpcClientConnectionManager.class, "requestAttributes");
 
+  @VisibleForTesting
+  static final AttributeKey<ClosableEpoch> EPOCH_ATTRIBUTE_KEY =
+      AttributeKey.valueOf(GrpcClientConnectionManager.class, "epoch");
+
   private static final Logger log = LoggerFactory.getLogger(GrpcClientConnectionManager.class);
 
   /**
@@ -107,6 +112,39 @@ public class GrpcClientConnectionManager implements DisconnectionRequestListener
     return requestAttributes;
   }
 
+  /**
+   * Handles the start of a server call, incrementing the active call count for the remote channel associated with the
+   * given server call.
+   *
+   * @param serverCall the server call to start
+   *
+   * @return {@code true} if the call should start normally or {@code false} if the call should be aborted because the
+   * underlying channel is closing
+   */
+  public boolean handleServerCallStart(final ServerCall<?, ?> serverCall) {
+    try {
+      return getRemoteChannel(serverCall).attr(EPOCH_ATTRIBUTE_KEY).get().tryArrive();
+    } catch (final ChannelNotFoundException e) {
+      // This would only happen if the channel had already closed, which is certainly possible. In this case, the call
+      // should certainly not proceed.
+      return false;
+    }
+  }
+
+  /**
+   * Handles completion (successful or not) of a server call, decrementing the active call count for the remote channel
+   * associated with the given server call.
+   *
+   * @param serverCall the server call to complete
+   */
+  public void handleServerCallComplete(final ServerCall<?, ?> serverCall) {
+    try {
+      getRemoteChannel(serverCall).attr(EPOCH_ATTRIBUTE_KEY).get().depart();
+    } catch (final ChannelNotFoundException ignored) {
+      // In practice, we'd only get here if the channel has already closed, so we can just ignore the exception
+    }
+  }
+
   /**
    * Closes any client connections to this host associated with the given authenticated device.
    *
@@ -119,10 +157,13 @@ public class GrpcClientConnectionManager implements DisconnectionRequestListener
     final List<Channel> channelsToClose =
         new ArrayList<>(remoteChannelsByAuthenticatedDevice.getOrDefault(authenticatedDevice, Collections.emptyList()));
 
-    channelsToClose.forEach(channel ->
-        channel.writeAndFlush(new CloseWebSocketFrame(ApplicationWebSocketCloseReason.REAUTHENTICATION_REQUIRED
-                .toWebSocketCloseStatus("Reauthentication required")))
-            .addListener(ChannelFutureListener.CLOSE_ON_FAILURE));
+    channelsToClose.forEach(channel -> channel.attr(EPOCH_ATTRIBUTE_KEY).get().close());
+  }
+
+  private static void closeRemoteChannel(final Channel channel) {
+    channel.writeAndFlush(new CloseWebSocketFrame(ApplicationWebSocketCloseReason.REAUTHENTICATION_REQUIRED
+            .toWebSocketCloseStatus("Reauthentication required")))
+        .addListener(ChannelFutureListener.CLOSE_ON_FAILURE);
   }
 
   @VisibleForTesting
@@ -200,6 +241,9 @@ public class GrpcClientConnectionManager implements DisconnectionRequestListener
     maybeAuthenticatedDevice.ifPresent(authenticatedDevice ->
         remoteChannel.attr(GrpcClientConnectionManager.AUTHENTICATED_DEVICE_ATTRIBUTE_KEY).set(authenticatedDevice));
 
+    remoteChannel.attr(EPOCH_ATTRIBUTE_KEY)
+        .set(new ClosableEpoch(() -> closeRemoteChannel(remoteChannel)));
+
     remoteChannelsByLocalAddress.put(localChannel.localAddress(), remoteChannel);
 
     getAuthenticatedDevice(remoteChannel).ifPresent(authenticatedDevice ->