upstream: change pool APIs

include/envoy/upstream/thread_local_cluster.h

@@ -8,6 +8,52 @@
 namespace Envoy {
 namespace Upstream {
 
+// HttpPoolData returns information about a given pool as well as a function
+// to create streams on that pool.
+struct HttpPoolData {
+  using OnNewStreamFn = std::function<void()>;
+
+  HttpPoolData(OnNewStreamFn on_new_stream, Http::ConnectionPool::Instance* pool)
+      : on_new_stream_(on_new_stream), pool_(pool) {}
+
+  Envoy::Http::ConnectionPool::Cancellable*
+  newStream(Http::ResponseDecoder& response_decoder,
+            Envoy::Http::ConnectionPool::Callbacks& callbacks) {
+    on_new_stream_();
+    return pool_->newStream(response_decoder, callbacks);
+  }
+
+  Upstream::HostDescriptionConstSharedPtr host() const { return pool_->host(); }
+
+private:
+  friend class HttpPoolDataPeer;
+
+  OnNewStreamFn on_new_stream_;
+  Http::ConnectionPool::Instance* pool_;
+};
+
+// Tcp pool returns information about a given pool, as well as a function to
+// create connections on that pool.
+struct TcpPoolData {
+  using OnNewConnectionFn = std::function<void()>;
+
+  TcpPoolData(OnNewConnectionFn on_new_connection, Tcp::ConnectionPool::Instance* pool)
+      : on_new_connection_(on_new_connection), pool_(pool) {}
+
+  Envoy::Tcp::ConnectionPool::Cancellable*
+  newConnection(Envoy::Tcp::ConnectionPool::Callbacks& callbacks) {
+    on_new_connection_();
+    return pool_->newConnection(callbacks);
+  }
+
+  Upstream::HostDescriptionConstSharedPtr host() const { return pool_->host(); }
+
+private:
+  friend class TcpPoolDataPeer;
+  OnNewConnectionFn on_new_connection_;
+  Tcp::ConnectionPool::Instance* pool_;
+};
+
 /**
  * A thread local cluster instance that can be used for direct load balancing and host set
  * interactions. In general, an instance of ThreadLocalCluster can only be safely used in the
@@ -42,10 +88,11 @@ class ThreadLocalCluster {
    * @param priority the connection pool priority.
    * @param downstream_protocol the downstream protocol (if one exists) to use in protocol
    *        selection.
-   * @param context the optional load balancer context.
-   * @return the connection pool or nullptr if there is no host available in the cluster.
+   * @param context the optional load balancer context. Must continue to be
+   *        valid until newConnection is called on the pool (if it is to be called).
+   * @return the connection pool data or nullopt if there is no host available in the cluster.
    */
-  virtual Http::ConnectionPool::Instance*
+  virtual absl::optional<HttpPoolData>
   httpConnPool(ResourcePriority priority, absl::optional<Http::Protocol> downstream_protocol,
                LoadBalancerContext* context) PURE;
 
@@ -55,11 +102,12 @@ class ThreadLocalCluster {
    * is used is the one defined on the cluster when it was created.
    *
    * @param priority the connection pool priority.
-   * @param context the optional load balancer context.
-   * @return the connection pool or nullptr if there is no host available in the cluster.
+   * @param context the optional load balancer context. Must continue to be
+   *        valid until newConnection is called on the pool (if it is to be called).
+   * @return the connection pool data or nullopt if there is no host available in the cluster.
    */
-  virtual Tcp::ConnectionPool::Instance* tcpConnPool(ResourcePriority priority,
-                                                     LoadBalancerContext* context) PURE;
+  virtual absl::optional<TcpPoolData> tcpConnPool(ResourcePriority priority,
+                                                  LoadBalancerContext* context) PURE;
 
   /**
    * Allocate a load balanced TCP connection for a cluster. The created connection is already

source/common/tcp_proxy/upstream.cc

@@ -150,7 +150,7 @@ TcpConnPool::TcpConnPool(Upstream::ThreadLocalCluster& thread_local_cluster,
                          Upstream::LoadBalancerContext* context,
                          Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks)
     : upstream_callbacks_(upstream_callbacks) {
-  conn_pool_ = thread_local_cluster.tcpConnPool(Upstream::ResourcePriority::Default, context);
+  conn_pool_data_ = thread_local_cluster.tcpConnPool(Upstream::ResourcePriority::Default, context);
 }
 
 TcpConnPool::~TcpConnPool() {
@@ -165,7 +165,7 @@ void TcpConnPool::newStream(GenericConnectionPoolCallbacks& callbacks) {
   // valid connection handle. If newConnection fails inline it may result in attempting to
   // select a new host, and a recursive call to initializeUpstreamConnection. In this case the
   // first call to newConnection will return null and the inner call will persist.
-  Tcp::ConnectionPool::Cancellable* handle = conn_pool_->newConnection(*this);
+  Tcp::ConnectionPool::Cancellable* handle = conn_pool_data_.value().newConnection(*this);
   if (handle) {
     ASSERT(upstream_handle_ == nullptr);
     upstream_handle_ = handle;
@@ -195,8 +195,8 @@ HttpConnPool::HttpConnPool(Upstream::ThreadLocalCluster& thread_local_cluster,
                            Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks,
                            Http::CodecClient::Type type)
     : config_(config), type_(type), upstream_callbacks_(upstream_callbacks) {
-  conn_pool_ = thread_local_cluster.httpConnPool(Upstream::ResourcePriority::Default, absl::nullopt,
-                                                 context);
+  conn_pool_data_ = thread_local_cluster.httpConnPool(Upstream::ResourcePriority::Default,
+                                                      absl::nullopt, context);
 }
 
 HttpConnPool::~HttpConnPool() {
@@ -215,7 +215,7 @@ void HttpConnPool::newStream(GenericConnectionPoolCallbacks& callbacks) {
     upstream_ = std::make_unique<Http2Upstream>(upstream_callbacks_, config_);
   }
   Tcp::ConnectionPool::Cancellable* handle =
-      conn_pool_->newStream(upstream_->responseDecoder(), *this);
+      conn_pool_data_.value().newStream(upstream_->responseDecoder(), *this);
   if (handle != nullptr) {
     upstream_handle_ = handle;
   }

source/common/tcp_proxy/upstream.h

@@ -5,6 +5,7 @@
 #include "envoy/tcp/conn_pool.h"
 #include "envoy/tcp/upstream.h"
 #include "envoy/upstream/load_balancer.h"
+#include "envoy/upstream/thread_local_cluster.h"
 #include "envoy/upstream/upstream.h"
 
 #include "common/common/dump_state_utils.h"
@@ -21,7 +22,7 @@ class TcpConnPool : public GenericConnPool, public Tcp::ConnectionPool::Callback
               Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks);
   ~TcpConnPool() override;
 
-  bool valid() const { return conn_pool_ != nullptr; }
+  bool valid() const { return conn_pool_data_.has_value(); }
 
   // GenericConnPool
   void newStream(GenericConnectionPoolCallbacks& callbacks) override;
@@ -33,7 +34,7 @@ class TcpConnPool : public GenericConnPool, public Tcp::ConnectionPool::Callback
                    Upstream::HostDescriptionConstSharedPtr host) override;
 
 private:
-  Tcp::ConnectionPool::Instance* conn_pool_{};
+  absl::optional<Upstream::TcpPoolData> conn_pool_data_{};
   Tcp::ConnectionPool::Cancellable* upstream_handle_{};
   GenericConnectionPoolCallbacks* callbacks_{};
   Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks_;
@@ -53,7 +54,9 @@ class HttpConnPool : public GenericConnPool, public Http::ConnectionPool::Callba
   ~HttpConnPool() override;
 
   // HTTP/3 upstreams are not supported at the moment.
-  bool valid() const { return conn_pool_ != nullptr && type_ <= Http::CodecClient::Type::HTTP2; }
+  bool valid() const {
+    return conn_pool_data_.has_value() && type_ <= Http::CodecClient::Type::HTTP2;
+  }
 
   // GenericConnPool
   void newStream(GenericConnectionPoolCallbacks& callbacks) override;
@@ -98,7 +101,7 @@ class HttpConnPool : public GenericConnPool, public Http::ConnectionPool::Callba
                           Ssl::ConnectionInfoConstSharedPtr ssl_info);
   const TunnelingConfig config_;
   Http::CodecClient::Type type_;
-  Http::ConnectionPool::Instance* conn_pool_{};
+  absl::optional<Upstream::HttpPoolData> conn_pool_data_{};
   Http::ConnectionPool::Cancellable* upstream_handle_{};
   GenericConnectionPoolCallbacks* callbacks_{};
   Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks_;

source/common/upstream/cluster_manager_impl.cc

@@ -904,41 +904,45 @@ void ClusterManagerImpl::maybePreconnect(
   }
 }
 
-Http::ConnectionPool::Instance*
+absl::optional<HttpPoolData>
 ClusterManagerImpl::ThreadLocalClusterManagerImpl::ClusterEntry::httpConnPool(
     ResourcePriority priority, absl::optional<Http::Protocol> protocol,
     LoadBalancerContext* context) {
   // Select a host and create a connection pool for it if it does not already exist.
-  auto ret = connPool(priority, protocol, context, false);
-
-  // Now see if another host should be preconnected.
-  // httpConnPool is called immediately before a call for newStream. newStream doesn't
-  // have the load balancer context needed to make selection decisions so preconnecting must be
-  // performed here in anticipation of the new stream.
-  // TODO(alyssawilk) refactor to have one function call and return a pair, so this invariant is
-  // code-enforced.
-  maybePreconnect(*this, parent_.cluster_manager_state_, [this, &priority, &protocol, &context]() {
-    return connPool(priority, protocol, context, true);
-  });
-  return ret;
+  auto pool = connPool(priority, protocol, context, false);
+  if (pool == nullptr) {
+    return absl::nullopt;
+  }
+
+  HttpPoolData data(
+      [this, priority, protocol, context]() -> void {
+        // Now that a new stream is being established, attempt to preconnect.
+        maybePreconnect(*this, parent_.cluster_manager_state_,
+                        [this, &priority, &protocol, &context]() {
+                          return connPool(priority, protocol, context, true);
+                        });
+      },
+      pool);
+  return data;
 }
 
-Tcp::ConnectionPool::Instance*
+absl::optional<TcpPoolData>
 ClusterManagerImpl::ThreadLocalClusterManagerImpl::ClusterEntry::tcpConnPool(
     ResourcePriority priority, LoadBalancerContext* context) {
   // Select a host and create a connection pool for it if it does not already exist.
-  auto ret = tcpConnPool(priority, context, false);
-
-  // tcpConnPool is called immediately before a call for newConnection. newConnection
-  // doesn't have the load balancer context needed to make selection decisions so preconnecting must
-  // be performed here in anticipation of the new connection.
-  // TODO(alyssawilk) refactor to have one function call and return a pair, so this invariant is
-  // code-enforced.
-  // Now see if another host should be preconnected.
-  maybePreconnect(*this, parent_.cluster_manager_state_,
-                  [this, &priority, &context]() { return tcpConnPool(priority, context, true); });
-
-  return ret;
+  auto pool = tcpConnPool(priority, context, false);
+  if (pool == nullptr) {
+    return absl::nullopt;
+  }
+
+  TcpPoolData data(
+      [this, priority, context]() -> void {
+        maybePreconnect(*this, parent_.cluster_manager_state_, [this, &priority, &context]() {
+          return tcpConnPool(priority, context, true);
+        });
+      },
+      pool);
+  return data;
 }
 
 void ClusterManagerImpl::postThreadLocalDrainConnections(const Cluster& cluster,

source/common/upstream/cluster_manager_impl.h

@@ -404,11 +404,11 @@ class ClusterManagerImpl : public ClusterManager, Logger::Loggable<Logger::Id::u
       const PrioritySet& prioritySet() override { return priority_set_; }
       ClusterInfoConstSharedPtr info() override { return cluster_info_; }
       LoadBalancer& loadBalancer() override { return *lb_; }
-      Http::ConnectionPool::Instance*
-      httpConnPool(ResourcePriority priority, absl::optional<Http::Protocol> downstream_protocol,
-                   LoadBalancerContext* context) override;
-      Tcp::ConnectionPool::Instance* tcpConnPool(ResourcePriority priority,
-                                                 LoadBalancerContext* context) override;
+      absl::optional<HttpPoolData> httpConnPool(ResourcePriority priority,
+                                                absl::optional<Http::Protocol> downstream_protocol,
+                                                LoadBalancerContext* context) override;
+      absl::optional<TcpPoolData> tcpConnPool(ResourcePriority priority,
+                                              LoadBalancerContext* context) override;
       Host::CreateConnectionData tcpConn(LoadBalancerContext* context) override;
       Http::AsyncClient& httpAsyncClient() override;
 

source/extensions/filters/network/dubbo_proxy/router/router_impl.cc

@@ -65,9 +65,8 @@ FilterStatus Router::onMessageDecoded(MessageMetadataSharedPtr metadata, Context
     return FilterStatus::StopIteration;
   }
 
-  Tcp::ConnectionPool::Instance* conn_pool =
-      cluster->tcpConnPool(Upstream::ResourcePriority::Default, this);
-  if (!conn_pool) {
+  auto conn_pool_data = cluster->tcpConnPool(Upstream::ResourcePriority::Default, this);
+  if (!conn_pool_data) {
     callbacks_->sendLocalReply(
         AppException(
             ResponseStatus::ServerError,
@@ -112,8 +111,9 @@ FilterStatus Router::onMessageDecoded(MessageMetadataSharedPtr metadata, Context
     upstream_request_buffer_.move(ctx->originMessage(), ctx->messageSize());
   }
 
-  upstream_request_ = std::make_unique<UpstreamRequest>(
-      *this, *conn_pool, metadata, callbacks_->serializationType(), callbacks_->protocolType());
+  upstream_request_ = std::make_unique<UpstreamRequest>(*this, *conn_pool_data, metadata,
+                                                        callbacks_->serializationType(),
+                                                        callbacks_->protocolType());
   return upstream_request_->start();
 }
 
@@ -188,11 +188,11 @@ void Router::cleanup() {
   }
 }
 
-Router::UpstreamRequest::UpstreamRequest(Router& parent, Tcp::ConnectionPool::Instance& pool,
+Router::UpstreamRequest::UpstreamRequest(Router& parent, Upstream::TcpPoolData& pool_data,
                                          MessageMetadataSharedPtr& metadata,
                                          SerializationType serialization_type,
                                          ProtocolType protocol_type)
-    : parent_(parent), conn_pool_(pool), metadata_(metadata),
+    : parent_(parent), conn_pool_data_(pool_data), metadata_(metadata),
       protocol_(
           NamedProtocolConfigFactory::getFactory(protocol_type).createProtocol(serialization_type)),
       request_complete_(false), response_started_(false), response_complete_(false),
@@ -201,7 +201,7 @@ Router::UpstreamRequest::UpstreamRequest(Router& parent, Tcp::ConnectionPool::In
 Router::UpstreamRequest::~UpstreamRequest() = default;
 
 FilterStatus Router::UpstreamRequest::start() {
-  Tcp::ConnectionPool::Cancellable* handle = conn_pool_.newConnection(*this);
+  Tcp::ConnectionPool::Cancellable* handle = conn_pool_data_.newConnection(*this);
   if (handle) {
     // Pause while we wait for a connection.
     conn_pool_handle_ = handle;

source/extensions/filters/network/dubbo_proxy/router/router_impl.h

@@ -47,7 +47,7 @@ class Router : public Tcp::ConnectionPool::UpstreamCallbacks,
 
 private:
   struct UpstreamRequest : public Tcp::ConnectionPool::Callbacks {
-    UpstreamRequest(Router& parent, Tcp::ConnectionPool::Instance& pool,
+    UpstreamRequest(Router& parent, Upstream::TcpPoolData& pool_data,
                     MessageMetadataSharedPtr& metadata, SerializationType serialization_type,
                     ProtocolType protocol_type);
     ~UpstreamRequest() override;
@@ -69,7 +69,7 @@ class Router : public Tcp::ConnectionPool::UpstreamCallbacks,
     void onResetStream(ConnectionPool::PoolFailureReason reason);
 
     Router& parent_;
-    Tcp::ConnectionPool::Instance& conn_pool_;
+    Upstream::TcpPoolData conn_pool_data_;
     MessageMetadataSharedPtr metadata_;
 
     Tcp::ConnectionPool::Cancellable* conn_pool_handle_{};

source/extensions/filters/network/rocketmq_proxy/router/router_impl.cc

@@ -106,17 +106,16 @@ void RouterImpl::sendRequestToUpstream(ActiveMessage& active_message) {
     return;
   }
 
-  Tcp::ConnectionPool::Instance* conn_pool =
-      cluster->tcpConnPool(Upstream::ResourcePriority::Default, this);
-  if (!conn_pool) {
+  auto data = cluster->tcpConnPool(Upstream::ResourcePriority::Default, this);
+  if (!data) {
     ENVOY_LOG(warn, "No host available for cluster {}. Opaque: {}", cluster_name, opaque);
     active_message.onError("No host available");
     reset();
     return;
   }
 
   upstream_request_ = std::make_unique<UpstreamRequest>(*this);
-  Tcp::ConnectionPool::Cancellable* cancellable = conn_pool->newConnection(*upstream_request_);
+  Tcp::ConnectionPool::Cancellable* cancellable = data.value().newConnection(*upstream_request_);
   if (cancellable) {
     handle_ = cancellable;
     ENVOY_LOG(trace, "No connection is available for now. Create a cancellable handle. Opaque: {}",

source/extensions/filters/network/thrift_proxy/router/router_impl.cc

@@ -281,9 +281,8 @@ FilterStatus Router::messageBegin(MessageMetadataSharedPtr metadata) {
     passthrough_supported_ = true;
   }
 
-  Tcp::ConnectionPool::Instance* conn_pool =
-      cluster->tcpConnPool(Upstream::ResourcePriority::Default, this);
-  if (!conn_pool) {
+  auto conn_pool_data = cluster->tcpConnPool(Upstream::ResourcePriority::Default, this);
+  if (!conn_pool_data) {
     stats_.no_healthy_upstream_.inc();
     callbacks_->sendLocalReply(
         AppException(AppExceptionType::InternalError,
@@ -303,7 +302,7 @@ FilterStatus Router::messageBegin(MessageMetadataSharedPtr metadata) {
   }
 
   upstream_request_ =
-      std::make_unique<UpstreamRequest>(*this, *conn_pool, metadata, transport, protocol);
+      std::make_unique<UpstreamRequest>(*this, *conn_pool_data, metadata, transport, protocol);
   return upstream_request_->start();
 }
 
@@ -422,10 +421,10 @@ void Router::convertMessageBegin(MessageMetadataSharedPtr metadata) {
 
 void Router::cleanup() { upstream_request_.reset(); }
 
-Router::UpstreamRequest::UpstreamRequest(Router& parent, Tcp::ConnectionPool::Instance& pool,
+Router::UpstreamRequest::UpstreamRequest(Router& parent, Upstream::TcpPoolData& pool_data,
                                          MessageMetadataSharedPtr& metadata,
                                          TransportType transport_type, ProtocolType protocol_type)
-    : parent_(parent), conn_pool_(pool), metadata_(metadata),
+    : parent_(parent), conn_pool_data_(pool_data), metadata_(metadata),
       transport_(NamedTransportConfigFactory::getFactory(transport_type).createTransport()),
       protocol_(NamedProtocolConfigFactory::getFactory(protocol_type).createProtocol()),
       request_complete_(false), response_started_(false), response_complete_(false) {}
@@ -437,7 +436,7 @@ Router::UpstreamRequest::~UpstreamRequest() {
 }
 
 FilterStatus Router::UpstreamRequest::start() {
-  Tcp::ConnectionPool::Cancellable* handle = conn_pool_.newConnection(*this);
+  Tcp::ConnectionPool::Cancellable* handle = conn_pool_data_.newConnection(*this);
   if (handle) {
     // Pause while we wait for a connection.
     conn_pool_handle_ = handle;