balancer: reduce lock scope in exact balance by splitting incNumConnections (#43862)

Commit Message:
Move the resource-limit bookkeeping (openConnections().inc()) outside
the critical section in
ExactConnectionBalancerImpl::pickTargetHandler(). Only the raw
connection counter increment, needed for balancing correctness, is held
under the lock; the ResourceLimit gauge update is deferred to after lock
release.

This is achieved by adding two optional virtual methods to
BalancedConnectionHandler with safe defaults:
- incConnectionCount(): increments only the balancing counter (defaults
to incNumConnections() for backward compatibility)
- postConnectionAccepted(): performs resource-limit bookkeeping
(defaults to no-op)

ActiveTcpListener overrides both to split the work that was previously
done in a single incNumConnections() call.

Additional Description:
Risk Level:
Testing:
Docs Changes:
Release Notes:
Platform Specific Features:
[Optional Runtime guard:]
[Optional Fixes #Issue]
[Optional Fixes commit #PR or SHA]
[Optional Deprecated:]
[Optional [API
Considerations](https://github.com/envoyproxy/envoy/blob/main/api/review_checklist.md):]

Signed-off-by: William Dauchy <william.dauchy@datadoghq.com>

Author: wdauchy

contrib/dlb/source/connection_balancer_impl.h

@@ -43,7 +43,8 @@ class DlbBalancedConnectionHandlerImpl : public Envoy::Network::BalancedConnecti
 
   // Only for override, those are never used.
   uint64_t numConnections() const override { return 0; }
-  void incNumConnections() override {}
+  void preIncNumConnections() override {}
+  void postIncNumConnections() override {}
 
 private:
   Envoy::Network::BalancedConnectionHandler& handler_;

envoy/network/connection_balancer.h

@@ -19,13 +19,20 @@ class BalancedConnectionHandler {
   virtual uint64_t numConnections() const PURE;
 
   /**
-   * Increment the number of connections within the handler. This must be called by a connection
-   * balancer implementation prior to a connection being picked via pickTargetHandler(). This makes
-   * sure that connection counts are accurate during connection transfer (i.e., that the target
-   * balancer accounts for the incoming connection). This is done by the balancer vs. the
-   * connection handler to account for different locking needs inside the balancer.
+   * Increment the connection count used for balancing decisions. Balancer implementations must call
+   * this prior to returning from pickTargetHandler() to ensure accurate connection counts during
+   * connection transfer. For balancers that hold a lock (e.g., ExactConnectionBalancerImpl), this
+   * should be called inside the critical section, while postIncNumConnections() is called after
+   * the lock is released.
    */
-  virtual void incNumConnections() PURE;
+  virtual void preIncNumConnections() PURE;
+
+  /**
+   * Perform any resource-limit bookkeeping after a connection has been assigned by the balancer.
+   * This is the counterpart to preIncNumConnections() and must be called outside any balancer lock
+   * to avoid holding the lock longer than necessary.
+   */
+  virtual void postIncNumConnections() PURE;
 
   /**
    * Post a connected socket to this connection handler. This is used for cross-thread connection
@@ -69,8 +76,9 @@ class ConnectionBalancer {
    * @return current_handler if the connection should stay bound to the current handler, or a
    *         different handler if the connection should be rebalanced.
    *
-   * NOTE: It is the responsibility of the balancer to call incNumConnections() on the returned
-   *       balancer. See the comments above for more explanation.
+   * NOTE: It is the responsibility of the balancer to call preIncNumConnections() and
+   *       postIncNumConnections() on the returned handler. See the comments above for more
+   *       explanation.
    */
   virtual BalancedConnectionHandler&
   pickTargetHandler(BalancedConnectionHandler& current_handler) PURE;

source/common/listener_manager/active_tcp_listener.h

@@ -68,9 +68,13 @@ class ActiveTcpListener final : public Network::TcpListenerCallbacks,
 
   // Network::BalancedConnectionHandler
   uint64_t numConnections() const override { return num_listener_connections_; }
+  void preIncNumConnections() override { ++num_listener_connections_; }
+  void postIncNumConnections() override { config_->openConnections().inc(); }
+
+  // ActiveStreamListenerBase
   void incNumConnections() override {
-    ++num_listener_connections_;
-    config_->openConnections().inc();
+    preIncNumConnections();
+    postIncNumConnections();
   }
   void post(Network::ConnectionSocketPtr&& socket) override;
   void onAcceptWorker(Network::ConnectionSocketPtr&& socket,

source/common/network/connection_balancer_impl.cc

@@ -32,9 +32,10 @@ ExactConnectionBalancerImpl::pickTargetHandler(BalancedConnectionHandler&) {
       }
     }
 
-    min_connection_handler->incNumConnections(); // NOLINT(clang-analyzer-core.CallAndMessage)
+    min_connection_handler->preIncNumConnections(); // NOLINT(clang-analyzer-core.CallAndMessage)
   }
 
+  min_connection_handler->postIncNumConnections();
   return *min_connection_handler;
 }
 

source/common/network/connection_balancer_impl.h

@@ -56,7 +56,8 @@ class NopConnectionBalancerImpl : public ConnectionBalancer {
   BalancedConnectionHandler&
   pickTargetHandler(BalancedConnectionHandler& current_handler) override {
     // In the NOP case just increment the connection count and return the current handler.
-    current_handler.incNumConnections();
+    current_handler.preIncNumConnections();
+    current_handler.postIncNumConnections();
     return current_handler;
   }
 };

source/extensions/bootstrap/internal_listener/active_internal_listener.h

@@ -86,7 +86,7 @@ class ActiveInternalListener : public Network::InternalListener,
   }
   void onAccept(Network::ConnectionSocketPtr&& socket) override;
 
-  // Network::BalancedConnectionHandler
+  // ActiveStreamListenerBase
   void incNumConnections() override { config_->openConnections().inc(); }
   void decNumConnections() override { config_->openConnections().dec(); }
 

test/server/active_tcp_listener_test.cc

@@ -701,9 +701,11 @@ TEST_F(ActiveTcpListenerTest, RedirectedRebalancer) {
 
   // 1. Listener1 re-balance. Set the balance target to the the active listener itself.
   EXPECT_CALL(balancer1, pickTargetHandler(_))
-      .WillOnce(testing::DoAll(
-          testing::WithArg<0>(Invoke([](auto& target) { target.incNumConnections(); })),
-          ReturnRef(*active_listener1)));
+      .WillOnce(testing::DoAll(testing::WithArg<0>(Invoke([](auto& target) {
+                                 target.preIncNumConnections();
+                                 target.postIncNumConnections();
+                               })),
+                               ReturnRef(*active_listener1)));
 
   EXPECT_CALL(listener_config1, filterChainFactory())
       .WillRepeatedly(ReturnRef(filter_chain_factory_));
@@ -732,9 +734,11 @@ TEST_F(ActiveTcpListenerTest, RedirectedRebalancer) {
 
   // 3. Listener2 re-balance. Set the balance target to the the active listener itself.
   EXPECT_CALL(balancer2, pickTargetHandler(_))
-      .WillOnce(testing::DoAll(
-          testing::WithArg<0>(Invoke([](auto& target) { target.incNumConnections(); })),
-          ReturnRef(*active_listener2)));
+      .WillOnce(testing::DoAll(testing::WithArg<0>(Invoke([](auto& target) {
+                                 target.preIncNumConnections();
+                                 target.postIncNumConnections();
+                               })),
+                               ReturnRef(*active_listener2)));
 
   EXPECT_CALL(listener_config2, filterChainFactory())
       .WillRepeatedly(ReturnRef(filter_chain_factory_));
@@ -810,9 +814,11 @@ TEST_F(ActiveTcpListenerTest, SkipRedirection) {
 
   // 1. Listener1 re-balance. Set the balance target to the the active listener itself.
   EXPECT_CALL(balancer1, pickTargetHandler(_))
-      .WillOnce(testing::DoAll(
-          testing::WithArg<0>(Invoke([](auto& target) { target.incNumConnections(); })),
-          ReturnRef(*active_listener1)));
+      .WillOnce(testing::DoAll(testing::WithArg<0>(Invoke([](auto& target) {
+                                 target.preIncNumConnections();
+                                 target.postIncNumConnections();
+                               })),
+                               ReturnRef(*active_listener1)));
 
   EXPECT_CALL(listener_config1, filterChainFactory())
       .WillRepeatedly(ReturnRef(filter_chain_factory_));
@@ -901,7 +907,8 @@ TEST_F(ActiveTcpListenerTest, Rebalance) {
   // active_listener1 re-balance. Set the balance target to the the active_listener2.
   EXPECT_CALL(balancer1, pickTargetHandler(_))
       .WillOnce(testing::DoAll(testing::WithArg<0>(Invoke([&active_listener2](auto&) {
-                                 active_listener2->incNumConnections();
+                                 active_listener2->preIncNumConnections();
+                                 active_listener2->postIncNumConnections();
                                })),
                                ReturnRef(*active_listener2)));
 

test/server/connection_handler_test.cc

@@ -515,7 +515,8 @@ TEST_F(ConnectionHandlerTest, RemoveListenerDuringRebalance) {
     post_cb = std::move(cb);
   });
   Network::MockConnectionSocket* connection = new NiceMock<Network::MockConnectionSocket>();
-  current_handler->incNumConnections();
+  current_handler->preIncNumConnections();
+  current_handler->postIncNumConnections();
 #ifndef NDEBUG
   EXPECT_CALL(*access_log_, log(_, _));
 #endif
@@ -791,7 +792,8 @@ TEST_F(ConnectionHandlerTest, RebalanceWithMultiAddressListener) {
   EXPECT_CALL(*access_log_, log(_, _));
   EXPECT_CALL(manager_, findFilterChain(_, _)).WillOnce(Return(nullptr));
 
-  current_handler1->incNumConnections();
+  current_handler1->preIncNumConnections();
+  current_handler1->postIncNumConnections();
   listener_callbacks1->onAccept(std::make_unique<NiceMock<Network::MockConnectionSocket>>());
 
   // Send connection to the second listener, expect mock_connection_balancer2 will be called.
@@ -801,7 +803,8 @@ TEST_F(ConnectionHandlerTest, RebalanceWithMultiAddressListener) {
   EXPECT_CALL(*access_log_, log(_, _));
   EXPECT_CALL(manager_, findFilterChain(_, _)).WillOnce(Return(nullptr));
 
-  current_handler2->incNumConnections();
+  current_handler2->preIncNumConnections();
+  current_handler2->postIncNumConnections();
   listener_callbacks2->onAccept(std::make_unique<NiceMock<Network::MockConnectionSocket>>());
 
   EXPECT_CALL(*mock_connection_balancer1, unregisterHandler(_));
@@ -2328,7 +2331,8 @@ TEST_F(ConnectionHandlerTest, TcpListenerInplaceUpdate) {
       << "new listener should be inplace added and callback should not change";
 
   Network::MockConnectionSocket* connection = new NiceMock<Network::MockConnectionSocket>();
-  current_handler->incNumConnections();
+  current_handler->preIncNumConnections();
+  current_handler->postIncNumConnections();
 
   EXPECT_CALL(*mock_connection_balancer, pickTargetHandler(_))
       .WillOnce(ReturnRef(*current_handler));