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. It is invalid to create new streams or
+  // connections from this pool after draining a pool with this behavior.
+  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

@@ -229,6 +229,7 @@ void ConnPoolImplBase::onStreamClosed(Envoy::ConnectionPool::ActiveClient& clien
 }
 
 ConnectionPool::Cancellable* ConnPoolImplBase::newStreamImpl(AttachContext& context) {
+  ASSERT(!is_draining_for_deletion_);
   ASSERT(!deferred_deleting_);
 
   ASSERT(static_cast<ssize_t>(connecting_stream_capacity_) ==
@@ -331,11 +332,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 +355,12 @@ void ConnPoolImplBase::closeIdleConnectionsForDrainingPool() {
   }
 }
 
-void ConnPoolImplBase::drainConnectionsImpl() {
+void ConnPoolImplBase::drainConnectionsImpl(DrainBehavior drain_behavior) {
+  if (drain_behavior == Envoy::ConnectionPool::DrainBehavior::DrainAndDelete) {
+    is_draining_for_deletion_ = true;
+    checkForIdleAndCloseIdleConnsIfDraining();
+    return;
+  }
   closeIdleConnectionsForDrainingPool();
 
   // closeIdleConnections() closes all connections in ready_clients_ with no active streams,
@@ -387,12 +388,13 @@ bool ConnPoolImplBase::isIdleImpl() const {
 }
 
 void ConnPoolImplBase::checkForIdleAndCloseIdleConnsIfDraining() {
-  if (is_draining_) {
+  if (is_draining_for_deletion_) {
     closeIdleConnectionsForDrainingPool();
   }
 
   if (isIdleImpl()) {
-    ENVOY_LOG(debug, "invoking idle callbacks - is_draining_={}", is_draining_);
+    ENVOY_LOG(debug, "invoking idle callbacks - is_draining_for_deletion_={}",
+              is_draining_for_deletion_);
     for (const Instance::IdleCb& cb : idle_callbacks_) {
       cb();
     }

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);
@@ -335,7 +334,7 @@ class ConnPoolImplBase : protected Logger::Loggable<Logger::Id::pool> {
 
   // Whether the connection pool is currently in the process of closing
   // all connections so that it can be gracefully deleted.
-  bool is_draining_{false};
+  bool is_draining_for_deletion_{false};
 
   // True iff this object is in the deferred delete list.
   bool deferred_deleting_{false};

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

@@ -3,6 +3,7 @@
 #include <functional>
 #include <vector>
 
+#include "envoy/common/conn_pool.h"
 #include "envoy/event/dispatcher.h"
 #include "envoy/upstream/resource_manager.h"
 #include "envoy/upstream/upstream.h"
@@ -58,15 +59,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(Envoy::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(
+    Envoy::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(Envoy::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