redis: upstream client draining and temp host connection limit and cl…

api/envoy/config/filter/network/redis_proxy/v2/redis_proxy.proto

@@ -10,6 +10,7 @@ option go_package = "v2";
 import "envoy/api/v2/core/base.proto";
 
 import "google/protobuf/duration.proto";
+import "google/protobuf/wrappers.proto";
 
 import "validate/validate.proto";
 import "gogoproto/gogo.proto";
@@ -82,6 +83,21 @@ message RedisProxy {
     // If `max_buffer_size_before_flush` is set, but `buffer_flush_timeout` is not, the latter
     // defaults to 3ms.
     google.protobuf.Duration buffer_flush_timeout = 5 [(gogoproto.stdduration) = true];
+
+    // If `enable_redirection` is true, upstream connections to hosts that are not part of cluster
+    // configuration may be created. If these hosts are later discovered to be part of the cluster
+    // (e.g., redis cluster) or a host is removed from the cluster (whether `enable_redirection` is
+    // true or not), then the associated connections are drained and closed.
+    // `upstream_drain_poll_interval` controls how often the draining connections are checked to see
+    // if they are inactive before being closed; this interval defaults to 100ms.
+    google.protobuf.Duration upstream_drain_poll_interval = 6 [(gogoproto.stdduration) = true];
+
+    // max_upstream_unknown_connections` controls how many upstream connections to unknown hosts can
+    // be created at any given time by any given worker thread (see `enable_redirection` for more
+    // details). If the host is unknown and a connection cannot be created due to enforcing this
+    // limit, then redirection will fail and the original redirection error will be passed
+    // downstream unchanged. This limit defaults to 100.
+    google.protobuf.UInt32Value max_upstream_unknown_connections = 7;
   }
 
   // Network settings for the connection pool to the upstream clusters.

source/extensions/clusters/redis/redis_cluster.h

@@ -227,6 +227,10 @@ class RedisCluster : public Upstream::BaseDynamicClusterImpl {
     bool enableRedirection() const override { return false; }
     uint32_t maxBufferSizeBeforeFlush() const override { return 0; }
     std::chrono::milliseconds bufferFlushTimeoutInMs() const override { return buffer_timeout_; }
+    std::chrono::milliseconds upstreamDrainPollIntervalInMs() const override {
+      return std::chrono::milliseconds(0);
+    }
+    uint32_t maxUpstreamUnknownConnections() const override { return 0; }
 
     // Extensions::NetworkFilters::Common::Redis::Client::PoolCallbacks
     void onResponse(NetworkFilters::Common::Redis::RespValuePtr&& value) override;

source/extensions/filters/network/common/redis/client.h

@@ -76,6 +76,12 @@ class Client : public Event::DeferredDeletable {
    */
   virtual void addConnectionCallbacks(Network::ConnectionCallbacks& callbacks) PURE;
 
+  /**
+   * Called to determine if the client has pending requests.
+   * @return bool true if the client is processing requests or false if it is currently idle.
+   */
+  virtual bool active() PURE;
+
   /**
    * Closes the underlying network connection.
    */
@@ -134,6 +140,18 @@ class Config {
    * @return timeout for batching commands for a single upstream host.
    */
   virtual std::chrono::milliseconds bufferFlushTimeoutInMs() const PURE;
+
+  /**
+   * @return interval for checking the state of draining upstream connections.
+   *
+   */
+  virtual std::chrono::milliseconds upstreamDrainPollIntervalInMs() const PURE;
+
+  /**
+   * @return the maximum number of upstream connections to unknown hosts when enableRedirection() is
+   * true.
+   */
+  virtual uint32_t maxUpstreamUnknownConnections() const PURE;
 };
 
 /**

source/extensions/filters/network/common/redis/client_impl.cc

@@ -16,9 +16,12 @@ ConfigImpl::ConfigImpl(
           config.max_buffer_size_before_flush()), // This is a scalar, so default is zero.
       buffer_flush_timeout_(PROTOBUF_GET_MS_OR_DEFAULT(
           config, buffer_flush_timeout,
-          3)) // Default timeout is 3ms. If max_buffer_size_before_flush is zero, this is not used
-              // as the buffer is flushed on each request immediately.
-{}
+          3)), // Default timeout is 3ms. If max_buffer_size_before_flush is zero, this is not used
+               // as the buffer is flushed on each request immediately.
+      upstream_drain_poll_interval_(PROTOBUF_GET_MS_OR_DEFAULT(
+          config, upstream_drain_poll_interval, 100)), // The default interval is 100ms.
+      max_upstream_unknown_connections_(
+          PROTOBUF_GET_WRAPPED_OR_DEFAULT(config, max_upstream_unknown_connections, 100)) {}
 
 ClientPtr ClientImpl::create(Upstream::HostConstSharedPtr host, Event::Dispatcher& dispatcher,
                              EncoderPtr&& encoder, DecoderFactory& decoder_factory,

source/extensions/filters/network/common/redis/client_impl.h

@@ -45,13 +45,21 @@ class ConfigImpl : public Config {
   std::chrono::milliseconds bufferFlushTimeoutInMs() const override {
     return buffer_flush_timeout_;
   }
+  std::chrono::milliseconds upstreamDrainPollIntervalInMs() const override {
+    return upstream_drain_poll_interval_;
+  }
+  uint32_t maxUpstreamUnknownConnections() const override {
+    return max_upstream_unknown_connections_;
+  }
 
 private:
   const std::chrono::milliseconds op_timeout_;
   const bool enable_hashtagging_;
   const bool enable_redirection_;
   const uint32_t max_buffer_size_before_flush_;
   const std::chrono::milliseconds buffer_flush_timeout_;
+  const std::chrono::milliseconds upstream_drain_poll_interval_;
+  const uint32_t max_upstream_unknown_connections_;
 };
 
 class ClientImpl : public Client, public DecoderCallbacks, public Network::ConnectionCallbacks {
@@ -68,6 +76,7 @@ class ClientImpl : public Client, public DecoderCallbacks, public Network::Conne
   }
   void close() override;
   PoolRequest* makeRequest(const RespValue& request, PoolCallbacks& callbacks) override;
+  bool active() override { return !pending_requests_.empty(); }
   void flushBufferAndResetTimer();
 
 private:

source/extensions/filters/network/redis_proxy/conn_pool_impl.cc

@@ -48,7 +48,8 @@ InstanceImpl::makeRequestToHost(const std::string& host_address,
 
 InstanceImpl::ThreadLocalPool::ThreadLocalPool(InstanceImpl& parent, Event::Dispatcher& dispatcher,
                                                std::string cluster_name)
-    : parent_(parent), dispatcher_(dispatcher), cluster_name_(std::move(cluster_name)) {
+    : parent_(parent), dispatcher_(dispatcher), cluster_name_(std::move(cluster_name)),
+      drain_timer_(dispatcher.createTimer([this]() -> void { drainClients(); })) {
 
   cluster_update_handle_ = parent_.cm_.addThreadLocalClusterUpdateCallbacks(*this);
   Upstream::ThreadLocalCluster* cluster = parent_.cm_.get(cluster_name_);
@@ -69,6 +70,9 @@ InstanceImpl::ThreadLocalPool::~ThreadLocalPool() {
   while (!client_map_.empty()) {
     client_map_.begin()->second->redis_client_->close();
   }
+  while (!clients_to_drain_.empty()) {
+    (*clients_to_drain_.begin())->redis_client_->close();
+  }
 }
 
 void InstanceImpl::ThreadLocalPool::onClusterAddOrUpdateNonVirtual(
@@ -86,8 +90,9 @@ void InstanceImpl::ThreadLocalPool::onClusterAddOrUpdateNonVirtual(
   cluster_ = &cluster;
   ASSERT(host_set_member_update_cb_handle_ == nullptr);
   host_set_member_update_cb_handle_ = cluster_->prioritySet().addMemberUpdateCb(
-      [this](const std::vector<Upstream::HostSharedPtr>&,
+      [this](const std::vector<Upstream::HostSharedPtr>& hosts_added,
              const std::vector<Upstream::HostSharedPtr>& hosts_removed) -> void {
+        onHostsAdded(hosts_added);
         onHostsRemoved(hosts_removed);
       });
 
@@ -106,25 +111,73 @@ void InstanceImpl::ThreadLocalPool::onClusterRemoval(const std::string& cluster_
 
   // Treat cluster removal as a removal of all hosts. Close all connections and fail all pending
   // requests.
+  if (host_set_member_update_cb_handle_ != nullptr) {
+    host_set_member_update_cb_handle_->remove();
+    host_set_member_update_cb_handle_ = nullptr;
+  }
   while (!client_map_.empty()) {
     client_map_.begin()->second->redis_client_->close();
   }
+  while (!clients_to_drain_.empty()) {
+    (*clients_to_drain_.begin())->redis_client_->close();
+  }
 
   cluster_ = nullptr;
-  host_set_member_update_cb_handle_ = nullptr;
   host_address_map_.clear();
 }
 
+void InstanceImpl::ThreadLocalPool::onHostsAdded(
+    const std::vector<Upstream::HostSharedPtr>& hosts_added) {
+  for (const auto& host : hosts_added) {
+    std::string host_address = host->address()->asString();
+    // Insert new host into address map, possibly overwriting a previous host's entry.
+    host_address_map_[host_address] = host;
+    for (const auto& created_host : created_hosts_) {
+      if (created_host->address()->asString() == host_address) {
+        // Remove our "temporary" host created in makeRequestToHost().
+        onHostsRemoved({created_host});
+        created_hosts_.remove(created_host);
+        break;
+      }
+    }
+  }
+}
+
 void InstanceImpl::ThreadLocalPool::onHostsRemoved(
     const std::vector<Upstream::HostSharedPtr>& hosts_removed) {
   for (const auto& host : hosts_removed) {
     auto it = client_map_.find(host);
     if (it != client_map_.end()) {
-      // We don't currently support any type of draining for redis connections. If a host is gone,
-      // we just close the connection. This will fail any pending requests.
-      it->second->redis_client_->close();
+      if (it->second->redis_client_->active()) {
+        // Put the ThreadLocalActiveClient to the side to drain.
+        clients_to_drain_.push_back(std::move(it->second));
+        client_map_.erase(it);
+        if (!drain_timer_->enabled()) {
+          drain_timer_->enableTimer(parent_.config_.upstreamDrainPollIntervalInMs());
+        }
+      } else {
+        // There are no pending requests so close the connection.
+        it->second->redis_client_->close();
+      }
     }
-    host_address_map_.erase(host->address()->asString());
+    // There is the possibility that multiple hosts with the same address
+    // are registered in host_address_map_ given that hosts may be created
+    // upon redirection or supplied as part of the cluster's definition.
+    try {
+      if (host_address_map_.at(host->address()->asString()) == host) {
+        host_address_map_.erase(host->address()->asString());
+      }
+    } catch (std::out_of_range&) {
+    }
+  }
+}
+
+void InstanceImpl::ThreadLocalPool::drainClients() {
+  while (!clients_to_drain_.empty() && !(*clients_to_drain_.begin())->redis_client_->active()) {
+    (*clients_to_drain_.begin())->redis_client_->close();
+  }
+  if (!clients_to_drain_.empty()) {
+    drain_timer_->enableTimer(parent_.config_.upstreamDrainPollIntervalInMs());
   }
 }
 
@@ -163,12 +216,6 @@ InstanceImpl::ThreadLocalPool::makeRequest(const std::string& key,
 
   ThreadLocalActiveClientPtr& client = threadLocalActiveClient(host);
 
-  // Keep host_address_map_ in sync with client_map_.
-  auto host_cached_by_address = host_address_map_.find(host->address()->asString());
-  if (host_cached_by_address == host_address_map_.end()) {
-    host_address_map_[host->address()->asString()] = host;
-  }
-
   return client->redis_client_->makeRequest(request, callbacks);
 }
 
@@ -211,9 +258,10 @@ InstanceImpl::ThreadLocalPool::makeRequestToHost(const std::string& host_address
   auto it = host_address_map_.find(host_address_map_key);
   if (it == host_address_map_.end()) {
     // This host is not known to the cluster manager. Create a new host and insert it into the map.
-    // TODO(msukalski): Add logic to track the number of these "unknown" host connections,
-    // cap the number of these connections, and implement time-out and cleaning logic, etc.
-
+    if (created_hosts_.size() == parent_.config_.maxUpstreamUnknownConnections()) {
+      // Too many upstream connections to unknown hosts have been created.
+      return nullptr;
+    }
     if (!ipv6) {
       // Only create an IPv4 address instance if we need a new Upstream::HostImpl.
       const auto ip_port = absl::string_view(host_address).substr(colon_pos + 1);
@@ -233,6 +281,7 @@ InstanceImpl::ThreadLocalPool::makeRequestToHost(const std::string& host_address
         envoy::api::v2::endpoint::Endpoint::HealthCheckConfig::default_instance(), 0,
         envoy::api::v2::core::HealthStatus::UNKNOWN)};
     host_address_map_[host_address_map_key] = new_host;
+    created_hosts_.push_back(new_host);
     it = host_address_map_.find(host_address_map_key);
   }
 
@@ -245,9 +294,19 @@ void InstanceImpl::ThreadLocalActiveClient::onEvent(Network::ConnectionEvent eve
   if (event == Network::ConnectionEvent::RemoteClose ||
       event == Network::ConnectionEvent::LocalClose) {
     auto client_to_delete = parent_.client_map_.find(host_);
-    ASSERT(client_to_delete != parent_.client_map_.end());
-    parent_.dispatcher_.deferredDelete(std::move(client_to_delete->second->redis_client_));
-    parent_.client_map_.erase(client_to_delete);
+    if (client_to_delete != parent_.client_map_.end()) {
+      parent_.dispatcher_.deferredDelete(std::move(redis_client_));
+      parent_.client_map_.erase(client_to_delete);
+    } else {
+      for (auto it = parent_.clients_to_drain_.begin(); it != parent_.clients_to_drain_.end();
+           it++) {
+        if ((*it).get() == this) {
+          parent_.dispatcher_.deferredDelete(std::move(redis_client_));
+          parent_.clients_to_drain_.erase(it);
+          break;
+        }
+      }
+    }
   }
 }
 

source/extensions/filters/network/redis_proxy/conn_pool_impl.h

@@ -84,7 +84,9 @@ class InstanceImpl : public Instance {
     makeRequestToHost(const std::string& host_address, const Common::Redis::RespValue& request,
                       Common::Redis::Client::PoolCallbacks& callbacks);
     void onClusterAddOrUpdateNonVirtual(Upstream::ThreadLocalCluster& cluster);
+    void onHostsAdded(const std::vector<Upstream::HostSharedPtr>& hosts_added);
     void onHostsRemoved(const std::vector<Upstream::HostSharedPtr>& hosts_removed);
+    void drainClients();
 
     // Upstream::ClusterUpdateCallbacks
     void onClusterAddOrUpdate(Upstream::ThreadLocalCluster& cluster) override {
@@ -101,6 +103,9 @@ class InstanceImpl : public Instance {
     Envoy::Common::CallbackHandle* host_set_member_update_cb_handle_{};
     std::unordered_map<std::string, Upstream::HostConstSharedPtr> host_address_map_;
     std::string auth_password_;
+    std::list<Upstream::HostSharedPtr> created_hosts_;
+    std::list<ThreadLocalActiveClientPtr> clients_to_drain_;
+    Event::TimerPtr drain_timer_;
   };
 
   struct LbContextImpl : public Upstream::LoadBalancerContextBase {

source/extensions/health_checkers/redis/redis.h

@@ -77,6 +77,12 @@ class RedisHealthChecker : public Upstream::HealthCheckerImplBase {
       return std::chrono::milliseconds(1);
     }
 
+    // Redirection connection and draining controls.
+    std::chrono::milliseconds upstreamDrainPollIntervalInMs() const override {
+      return std::chrono::milliseconds(0);
+    }
+    uint32_t maxUpstreamUnknownConnections() const override { return 0; }
+
     // Extensions::NetworkFilters::Common::Redis::Client::PoolCallbacks
     void onResponse(NetworkFilters::Common::Redis::RespValuePtr&& value) override;
     void onFailure() override;

test/extensions/clusters/redis/redis_cluster_test.cc

@@ -375,6 +375,8 @@ class RedisClusterTest : public testing::Test,
     RedisCluster::RedisDiscoverySession discovery_session(*cluster_, *this);
     EXPECT_FALSE(discovery_session.enableHashtagging());
     EXPECT_EQ(discovery_session.bufferFlushTimeoutInMs(), std::chrono::milliseconds(0));
+    EXPECT_EQ(discovery_session.upstreamDrainPollIntervalInMs(), std::chrono::milliseconds(0));
+    EXPECT_EQ(discovery_session.maxUpstreamUnknownConnections(), 0);
 
     NetworkFilters::Common::Redis::RespValue dummy_value;
     dummy_value.type(NetworkFilters::Common::Redis::RespType::Error);

test/extensions/filters/network/common/redis/client_impl_test.cc

@@ -80,6 +80,7 @@ class RedisClientImplTest : public testing::Test, public Common::Redis::DecoderF
                                  *config_);
     EXPECT_EQ(1UL, host_->cluster_.stats_.upstream_cx_total_.value());
     EXPECT_EQ(1UL, host_->stats_.cx_total_.value());
+    EXPECT_EQ(false, client_->active());
 
     // NOP currently.
     upstream_connection_->runHighWatermarkCallbacks();
@@ -95,6 +96,7 @@ class RedisClientImplTest : public testing::Test, public Common::Redis::DecoderF
     Common::Redis::RespValuePtr response1{new Common::Redis::RespValue()};
     response1->type(Common::Redis::RespType::SimpleString);
     response1->asString() = "OK";
+    EXPECT_EQ(true, client_->active());
     ClientImpl* client_impl = dynamic_cast<ClientImpl*>(client_.get());
     EXPECT_NE(client_impl, nullptr);
     client_impl->onRespValue(std::move(response1));
@@ -154,6 +156,10 @@ class ConfigBufferSizeGTSingleRequest : public Config {
   std::chrono::milliseconds bufferFlushTimeoutInMs() const override {
     return std::chrono::milliseconds(1);
   }
+  std::chrono::milliseconds upstreamDrainPollIntervalInMs() const override {
+    return std::chrono::milliseconds(0);
+  }
+  uint32_t maxUpstreamUnknownConnections() const override { return 0; }
 };
 
 TEST_F(RedisClientImplTest, BatchWithTimerFiring) {
@@ -455,6 +461,10 @@ class ConfigOutlierDisabled : public Config {
   std::chrono::milliseconds bufferFlushTimeoutInMs() const override {
     return std::chrono::milliseconds(0);
   }
+  std::chrono::milliseconds upstreamDrainPollIntervalInMs() const override {
+    return std::chrono::milliseconds(0);
+  }
+  uint32_t maxUpstreamUnknownConnections() const override { return 0; }
 };
 
 TEST_F(RedisClientImplTest, OutlierDisabled) {

test/extensions/filters/network/common/redis/mocks.h

@@ -69,6 +69,7 @@ class MockClient : public Client {
   }
 
   MOCK_METHOD1(addConnectionCallbacks, void(Network::ConnectionCallbacks& callbacks));
+  MOCK_METHOD0(active, bool());
   MOCK_METHOD0(close, void());
   MOCK_METHOD2(makeRequest,
                PoolRequest*(const Common::Redis::RespValue& request, PoolCallbacks& callbacks));

test/extensions/filters/network/common/redis/test_utils.h

@@ -17,11 +17,12 @@ namespace Client {
 
 inline envoy::config::filter::network::redis_proxy::v2::RedisProxy::ConnPoolSettings
 createConnPoolSettings(int64_t millis = 20, bool hashtagging = true,
-                       bool redirection_support = true) {
+                       bool redirection_support = true, uint32_t max_unknown_conns = 100) {
   envoy::config::filter::network::redis_proxy::v2::RedisProxy::ConnPoolSettings setting{};
   setting.mutable_op_timeout()->CopyFrom(Protobuf::util::TimeUtil::MillisecondsToDuration(millis));
   setting.set_enable_hashtagging(hashtagging);
   setting.set_enable_redirection(redirection_support);
+  setting.mutable_max_upstream_unknown_connections()->set_value(max_unknown_conns);
   return setting;
 }