rpc/connection_cache: protect spawned tasks with gate

Prevents the sharded connection_cache from being stopped while there are in-flight tasks for it.
ballard26 · Oct 4, 2023 · 5d1ae81 · 5d1ae81
1 parent 8649216
commit 5d1ae81
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 15 deletions.
diff --git a/src/v/rpc/connection_cache.cc b/src/v/rpc/connection_cache.cc
@@ -331,10 +331,12 @@ ss::future<> connection_cache::reset_client_backoff(
         return ss::now();
     }
 
-    return container().invoke_on(
-      *shard, [node_id](rpc::connection_cache& cache) mutable {
-          cache._cache.reset_client_backoff(node_id);
-      });
+    return ss::with_gate(_gate, [this, node_id, shard] {
+        return container().invoke_on(
+          *shard, [node_id](rpc::connection_cache& cache) mutable {
+              cache._cache.reset_client_backoff(node_id);
+          });
+    });
 }
 
 } // namespace rpc
diff --git a/src/v/rpc/connection_cache.h b/src/v/rpc/connection_cache.h
@@ -21,6 +21,7 @@
 #include "rpc/types.h"
 #include "utils/mutex.h"
 
+#include <seastar/core/gate.hh>
 #include <seastar/core/sharded.hh>
 #include <seastar/core/shared_ptr.hh>
 
@@ -141,17 +142,25 @@ class connection_cache final
               rpc::make_error_code(errc::missing_node_rpc_client));
         }
 
-        return container().invoke_on(
-          *shard,
-          [node_id, f = std::forward<Func>(f), connection_timeout](
-            connection_cache& cache) mutable {
-              if (cache.is_shutting_down()) {
-                  return ss::futurize<ret_t>::convert(
-                    rpc::make_error_code(errc::shutting_down));
-              }
-
-              return cache._cache.with_node_client<Protocol, Func>(
-                node_id, connection_timeout, std::forward<Func>(f));
+        return ss::with_gate(
+          _gate,
+          [this,
+           node_id,
+           connection_timeout,
+           shard,
+           f = std::forward<Func>(f)]() mutable {
+              return container().invoke_on(
+                *shard,
+                [node_id, f = std::forward<Func>(f), connection_timeout](
+                  connection_cache& cache) mutable {
+                    if (cache.is_shutting_down()) {
+                        return ss::futurize<ret_t>::convert(
+                          rpc::make_error_code(errc::shutting_down));
+                    }
+
+                    return cache._cache.with_node_client<Protocol, Func>(
+                      node_id, connection_timeout, std::forward<Func>(f));
+                });
           });
     }