conn_pool: Remove startDrain() and replace it with an argument to drainConnections

envoy/common/conn_pool.h

@@ -36,6 +36,21 @@ class Cancellable {
   virtual void cancel(CancelPolicy cancel_policy) PURE;
 };
 
+/**
+ * Controls the behavior when draining a connection pool.
+ */
+enum class DrainBehavior {
+  // Starts draining a pool, by gracefully completing all requests and gracefully closing all
+  // connections, in preparation for deletion.
+  // can be immediately drained.
+  DrainAndDelete,
+  // Actively drain all existing connection pool connections. This can be used in cases where
+  // the connection pool is not being destroyed, but the caller wishes to make sure that
+  // all new streams take place on a new connection. For example, when a health check failure
+  // occurs.
+  DrainExistingConnections,
+};
+
 /**
  * An instance of a generic connection pool.
  */
@@ -59,20 +74,10 @@ class Instance {
   virtual bool isIdle() const PURE;
 
   /**
-   * Starts draining a pool, by gracefully completing all requests and gracefully closing all
-   * connections, in preparation for deletion. When the process completes, the function registered
-   * via `addIdleCallback()` is called. The callback may occur before this call returns if the pool
-   * can be immediately drained.
-   */
-  virtual void startDrain() PURE;
-
-  /**
-   * Actively drain all existing connection pool connections. This method can be used in cases
-   * where the connection pool is not being destroyed, but the caller wishes to make sure that
-   * all new streams take place on a new connection. For example, when a health check failure
-   * occurs.
+   * Drains the connections in a pool.
+   * @param drain_behavior A DrainBehavior that controls the behavior of the draining.
    */
-  virtual void drainConnections() PURE;
+  virtual void drainConnections(DrainBehavior drain_behavior) PURE;
 
   /**
    * @return Upstream::HostDescriptionConstSharedPtr the host for which connections are pooled.

source/common/conn_pool/conn_pool_base.cc

@@ -331,11 +331,6 @@ void ConnPoolImplBase::transitionActiveClientState(ActiveClient& client,
 
 void ConnPoolImplBase::addIdleCallbackImpl(Instance::IdleCb cb) { idle_callbacks_.push_back(cb); }
 
-void ConnPoolImplBase::startDrainImpl() {
-  is_draining_ = true;
-  checkForIdleAndCloseIdleConnsIfDraining();
-}
-
 void ConnPoolImplBase::closeIdleConnectionsForDrainingPool() {
   // Create a separate list of elements to close to avoid mutate-while-iterating problems.
   std::list<ActiveClient*> to_close;
@@ -359,7 +354,12 @@ void ConnPoolImplBase::closeIdleConnectionsForDrainingPool() {
   }
 }
 
-void ConnPoolImplBase::drainConnectionsImpl() {
+void ConnPoolImplBase::drainConnectionsImpl(DrainBehavior drain_behavior) {
+  if (drain_behavior == Envoy::ConnectionPool::DrainBehavior::DrainAndDelete) {
+    checkForIdleAndCloseIdleConnsIfDraining();
+    is_draining_ = true;
+    return;
+  }
   closeIdleConnectionsForDrainingPool();
 
   // closeIdleConnections() closes all connections in ready_clients_ with no active streams,

source/common/conn_pool/conn_pool_base.h

@@ -166,8 +166,7 @@ class ConnPoolImplBase : protected Logger::Loggable<Logger::Id::pool> {
   void addIdleCallbackImpl(Instance::IdleCb cb);
   // Returns true if the pool is idle.
   bool isIdleImpl() const;
-  void startDrainImpl();
-  void drainConnectionsImpl();
+  void drainConnectionsImpl(DrainBehavior drain_behavior);
   const Upstream::HostConstSharedPtr& host() const { return host_; }
   // Called if this pool is likely to be picked soon, to determine if it's worth preconnecting.
   bool maybePreconnectImpl(float global_preconnect_ratio);

source/common/http/conn_pool_base.h

@@ -61,8 +61,9 @@ class HttpConnPoolImplBase : public Envoy::ConnectionPool::ConnPoolImplBase,
   // ConnectionPool::Instance
   void addIdleCallback(IdleCb cb) override { addIdleCallbackImpl(cb); }
   bool isIdle() const override { return isIdleImpl(); }
-  void startDrain() override { startDrainImpl(); }
-  void drainConnections() override { drainConnectionsImpl(); }
+  void drainConnections(Envoy::ConnectionPool::DrainBehavior drain_behavior) override {
+    drainConnectionsImpl(drain_behavior);
+  }
   Upstream::HostDescriptionConstSharedPtr host() const override { return host_; }
   ConnectionPool::Cancellable* newStream(Http::ResponseDecoder& response_decoder,
                                          Http::ConnectionPool::Callbacks& callbacks) override;

source/common/http/conn_pool_grid.cc

@@ -296,24 +296,20 @@ void ConnectivityGrid::addIdleCallback(IdleCb cb) {
   idle_callbacks_.emplace_back(cb);
 }
 
-void ConnectivityGrid::startDrain() {
+void ConnectivityGrid::drainConnections(Envoy::ConnectionPool::DrainBehavior drain_behavior) {
   if (draining_) {
     // A drain callback has already been set, and only needs to happen once.
     return;
   }
 
-  // Note that no new pools can be created from this point on
-  // as createNextPool fast-fails if `draining_` is true.
-  draining_ = true;
-
-  for (auto& pool : pools_) {
-    pool->startDrain();
+  if (drain_behavior == Envoy::ConnectionPool::DrainBehavior::DrainAndDelete) {
+    // Note that no new pools can be created from this point on
+    // as createNextPool fast-fails if `draining_` is true.
+    draining_ = true;
   }
-}
 
-void ConnectivityGrid::drainConnections() {
   for (auto& pool : pools_) {
-    pool->drainConnections();
+    pool->drainConnections(drain_behavior);
   }
 }
 

source/common/http/conn_pool_grid.h

@@ -148,8 +148,7 @@ class ConnectivityGrid : public ConnectionPool::Instance,
                                          ConnectionPool::Callbacks& callbacks) override;
   void addIdleCallback(IdleCb cb) override;
   bool isIdle() const override;
-  void startDrain() override;
-  void drainConnections() override;
+  void drainConnections(Envoy::ConnectionPool::DrainBehavior drain_behavior) override;
   Upstream::HostDescriptionConstSharedPtr host() const override;
   bool maybePreconnect(float preconnect_ratio) override;
   absl::string_view protocolDescription() const override { return "connection grid"; }

source/common/tcp/conn_pool.h

@@ -150,9 +150,11 @@ class ConnPoolImpl : public Envoy::ConnectionPool::ConnPoolImplBase,
 
   void addIdleCallback(IdleCb cb) override { addIdleCallbackImpl(cb); }
   bool isIdle() const override { return isIdleImpl(); }
-  void startDrain() override { startDrainImpl(); }
-  void drainConnections() override {
-    drainConnectionsImpl();
+  void drainConnections(Envoy::ConnectionPool::DrainBehavior drain_behavior) override {
+    drainConnectionsImpl(drain_behavior);
+    if (drain_behavior == Envoy::ConnectionPool::DrainBehavior::DrainAndDelete) {
+      return;
+    }
     // Legacy behavior for the TCP connection pool marks all connecting clients
     // as draining.
     for (auto& connecting_client : connecting_clients_) {

source/common/tcp/original_conn_pool.cc

@@ -37,7 +37,13 @@ OriginalConnPoolImpl::~OriginalConnPoolImpl() {
   dispatcher_.clearDeferredDeleteList();
 }
 
-void OriginalConnPoolImpl::drainConnections() {
+void OriginalConnPoolImpl::drainConnections(Envoy::ConnectionPool::DrainBehavior drain_behavior) {
+  if (drain_behavior == Envoy::ConnectionPool::DrainBehavior::DrainAndDelete) {
+    is_draining_ = true;
+    checkForIdleAndCloseIdleConnsIfDraining();
+    return;
+  }
+
   ENVOY_LOG(debug, "draining connections");
   while (!ready_conns_.empty()) {
     ready_conns_.front()->conn_->close(Network::ConnectionCloseType::NoFlush);
@@ -70,11 +76,6 @@ void OriginalConnPoolImpl::closeConnections() {
 
 void OriginalConnPoolImpl::addIdleCallback(IdleCb cb) { idle_callbacks_.push_back(cb); }
 
-void OriginalConnPoolImpl::startDrain() {
-  is_draining_ = true;
-  checkForIdleAndCloseIdleConnsIfDraining();
-}
-
 void OriginalConnPoolImpl::assignConnection(ActiveConn& conn,
                                             ConnectionPool::Callbacks& callbacks) {
   ASSERT(conn.wrapper_ == nullptr);

source/common/tcp/original_conn_pool.h

@@ -34,8 +34,7 @@ class OriginalConnPoolImpl : Logger::Loggable<Logger::Id::pool>, public Connecti
   // ConnectionPool::Instance
   void addIdleCallback(IdleCb cb) override;
   bool isIdle() const override;
-  void startDrain() override;
-  void drainConnections() override;
+  void drainConnections(Envoy::ConnectionPool::DrainBehavior drain_behavior) override;
   void closeConnections() override;
   ConnectionPool::Cancellable* newConnection(ConnectionPool::Callbacks& callbacks) override;
   // The old pool does not implement preconnecting.

source/common/upstream/cluster_manager_impl.cc

@@ -1197,7 +1197,7 @@ void ClusterManagerImpl::ThreadLocalClusterManagerImpl::drainConnPools(
   // guarding deletion with `do_not_delete_` in the registered idle callback, and then checking
   // afterwards whether it is empty and deleting it if necessary.
   container.do_not_delete_ = true;
-  pools->startDrain();
+  pools->drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete);
   container.do_not_delete_ = false;
 
   if (container.pools_->size() == 0) {
@@ -1217,7 +1217,7 @@ void ClusterManagerImpl::ThreadLocalClusterManagerImpl::drainTcpConnPools(
 
   container.draining_ = true;
   for (auto pool : pools) {
-    pool->startDrain();
+    pool->drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete);
   }
 }
 
@@ -1389,7 +1389,8 @@ void ClusterManagerImpl::ThreadLocalClusterManagerImpl::drainAllConnPoolsWorker(
     const auto container = getHttpConnPoolsContainer(host);
     if (container != nullptr) {
       container->do_not_delete_ = true;
-      container->pools_->drainConnections();
+      container->pools_->drainConnections(
+          Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections);
       container->do_not_delete_ = false;
 
       if (container->pools_->size() == 0) {
@@ -1417,7 +1418,7 @@ void ClusterManagerImpl::ThreadLocalClusterManagerImpl::drainAllConnPoolsWorker(
             ClusterInfo::Features::CLOSE_CONNECTIONS_ON_HOST_HEALTH_FAILURE) {
           pool->closeConnections();
         } else {
-          pool->drainConnections();
+          pool->drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections);
         }
       }
     }

source/common/upstream/conn_pool_map.h

@@ -58,15 +58,10 @@ template <typename KEY_TYPE, typename POOL_TYPE> class ConnPoolMap {
    */
   void addIdleCallback(const IdleCb& cb);
 
-  /**
-   * See `Envoy::ConnectionPool::Instance::startDrain()`.
-   */
-  void startDrain();
-
   /**
    * See `Envoy::ConnectionPool::Instance::drainConnections()`.
    */
-  void drainConnections();
+  void drainConnections(ConnectionPool::DrainBehavior drain_behavior);
 
 private:
   /**

source/common/upstream/conn_pool_map_impl.h

@@ -99,7 +99,8 @@ void ConnPoolMap<KEY_TYPE, POOL_TYPE>::addIdleCallback(const IdleCb& cb) {
 }
 
 template <typename KEY_TYPE, typename POOL_TYPE>
-void ConnPoolMap<KEY_TYPE, POOL_TYPE>::startDrain() {
+void ConnPoolMap<KEY_TYPE, POOL_TYPE>::drainConnections(
+    ConnectionPool::DrainBehavior drain_behavior) {
   // Copy the `active_pools_` so that it is safe for the call to result
   // in deletion, and avoid iteration through a mutating container.
   std::vector<POOL_TYPE*> pools;
@@ -109,22 +110,7 @@ void ConnPoolMap<KEY_TYPE, POOL_TYPE>::startDrain() {
   }
 
   for (auto* pool : pools) {
-    pool->startDrain();
-  }
-}
-
-template <typename KEY_TYPE, typename POOL_TYPE>
-void ConnPoolMap<KEY_TYPE, POOL_TYPE>::drainConnections() {
-  // Copy the `active_pools_` so that it is safe for the call to result
-  // in deletion, and avoid iteration through a mutating container.
-  std::vector<POOL_TYPE*> pools;
-  pools.reserve(active_pools_.size());
-  for (auto& pool_pair : active_pools_) {
-    pools.push_back(pool_pair.second.get());
-  }
-
-  for (auto* pool : pools) {
-    pool->drainConnections();
+    pool->drainConnections(drain_behavior);
   }
 }
 

source/common/upstream/priority_conn_pool_map.h

@@ -51,15 +51,10 @@ template <typename KEY_TYPE, typename POOL_TYPE> class PriorityConnPoolMap {
    */
   void addIdleCallback(const IdleCb& cb);
 
-  /**
-   * See `Envoy::ConnectionPool::Instance::startDrain()`.
-   */
-  void startDrain();
-
   /**
    * See `Envoy::ConnectionPool::Instance::drainConnections()`.
    */
-  void drainConnections();
+  void drainConnections(ConnectionPool::DrainBehavior drain_behavior);
 
 private:
   size_t getPriorityIndex(ResourcePriority priority) const;

source/common/upstream/priority_conn_pool_map_impl.h

@@ -55,16 +55,10 @@ void PriorityConnPoolMap<KEY_TYPE, POOL_TYPE>::addIdleCallback(const IdleCb& cb)
 }
 
 template <typename KEY_TYPE, typename POOL_TYPE>
-void PriorityConnPoolMap<KEY_TYPE, POOL_TYPE>::startDrain() {
+void PriorityConnPoolMap<KEY_TYPE, POOL_TYPE>::drainConnections(
+    ConnectionPool::DrainBehavior drain_behavior) {
   for (auto& pool_map : conn_pool_maps_) {
-    pool_map->startDrain();
-  }
-}
-
-template <typename KEY_TYPE, typename POOL_TYPE>
-void PriorityConnPoolMap<KEY_TYPE, POOL_TYPE>::drainConnections() {
-  for (auto& pool_map : conn_pool_maps_) {
-    pool_map->drainConnections();
+    pool_map->drainConnections(drain_behavior);
   }
 }
 

test/common/conn_pool/conn_pool_base_test.cc

@@ -227,7 +227,7 @@ TEST_F(ConnPoolImplBaseTest, PoolIdleCallbackTriggeredRemoteClose) {
   clients_.back()->onEvent(Network::ConnectionEvent::RemoteClose);
 
   EXPECT_CALL(idle_pool_callback, Call());
-  pool_.startDrainImpl();
+  pool_.drainConnectionsImpl(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete);
 }
 
 // Local close simulates what would happen for an idle timeout on a connection.
@@ -255,7 +255,7 @@ TEST_F(ConnPoolImplBaseTest, PoolIdleCallbackTriggeredLocalClose) {
   clients_.back()->onEvent(Network::ConnectionEvent::LocalClose);
 
   EXPECT_CALL(idle_pool_callback, Call());
-  pool_.startDrainImpl();
+  pool_.drainConnectionsImpl(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete);
 }
 
 } // namespace ConnectionPool

test/common/http/conn_pool_grid_test.cc

@@ -382,20 +382,23 @@ TEST_F(ConnectivityGridTest, TestCancel) {
 
 // Make sure drains get sent to all active pools.
 TEST_F(ConnectivityGridTest, Drain) {
-  grid_.drainConnections();
+  grid_.drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections);
 
   // Synthetically create a pool.
   grid_.createNextPool();
   {
-    EXPECT_CALL(*grid_.first(), drainConnections());
-    grid_.drainConnections();
+    EXPECT_CALL(*grid_.first(),
+                drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections));
+    grid_.drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections);
   }
 
   grid_.createNextPool();
   {
-    EXPECT_CALL(*grid_.first(), drainConnections());
-    EXPECT_CALL(*grid_.second(), drainConnections());
-    grid_.drainConnections();
+    EXPECT_CALL(*grid_.first(),
+                drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections));
+    EXPECT_CALL(*grid_.second(),
+                drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections));
+    grid_.drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainExistingConnections);
   }
 }
 
@@ -411,16 +414,22 @@ TEST_F(ConnectivityGridTest, DrainCallbacks) {
 
   // The first time a drain is started, both pools should start draining.
   {
-    EXPECT_CALL(*grid_.first(), startDrain());
-    EXPECT_CALL(*grid_.second(), startDrain());
-    grid_.startDrain();
+    EXPECT_CALL(*grid_.first(),
+                drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete));
+    EXPECT_CALL(*grid_.second(),
+                drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete));
+    grid_.drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete);
   }
 
   // The second time, the pools will not see any change.
   {
-    EXPECT_CALL(*grid_.first(), startDrain()).Times(0);
-    EXPECT_CALL(*grid_.second(), startDrain()).Times(0);
-    grid_.startDrain();
+    EXPECT_CALL(*grid_.first(),
+                drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete))
+        .Times(0);
+    EXPECT_CALL(*grid_.second(),
+                drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete))
+        .Times(0);
+    grid_.drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete);
   }
   {
     // Notify the grid the second pool has been drained. This should not be
@@ -481,7 +490,7 @@ TEST_F(ConnectivityGridTest, NoDrainOnTeardown) {
 
   {
     grid_.addIdleCallback([&drain_received]() -> void { drain_received = true; });
-    grid_.startDrain();
+    grid_.drainConnections(Envoy::ConnectionPool::DrainBehavior::DrainAndDelete);
   }
 
   grid_.setDestroying(); // Fake being in the destructor.