dns cache: Replace copy-on-add/remove with lock-guarded global map (envoyproxy#14349)

* Change both dns cache and dfp cluster to use global maps guarded by
mutexes. Previously, each of these maps would be copied on each
add/remove, which led to performance issues as the maps grew.

* Change worker DNS cache callback container to use a map from
hostnames to lists of callbacks. This ensures that callbacks are only
invoked for the hosts that are completely ready (e.g. the dynamic
forward proxy's host info has been populated).

* Hold the PrimaryHostInfo objects by shared_ptr to reduce the time
required for the main thread to hold the map lock.

Signed-off-by: Craig Radcliffe <craig.radcliffe@broadcom.com>

Author: Craig Radcliffe

source/common/common/cleanup.h

@@ -55,4 +55,42 @@ template <class T> class RaiiListElement {
   bool cancelled_;
 };
 
+// RAII helper class to add an element to a std::list held inside an absl::flat_hash_map on
+// construction and erase it on destruction, unless the cancel method has been called. If the list
+// is empty after removal of the element, the destructor will also remove the list from the map.
+template <class Key, class Value> class RaiiMapOfListElement {
+public:
+  using MapOfList = absl::flat_hash_map<Key, std::list<Value>>;
+
+  template <typename ConvertibleToKey>
+  RaiiMapOfListElement(MapOfList& map, const ConvertibleToKey& key, Value value)
+      : map_(map), list_(map_.try_emplace(key).first->second), key_(key), cancelled_(false) {
+    it_ = list_.emplace(list_.begin(), value);
+  }
+
+  virtual ~RaiiMapOfListElement() {
+    if (!cancelled_) {
+      erase();
+    }
+  }
+
+  void cancel() { cancelled_ = true; }
+
+private:
+  void erase() {
+    ASSERT(!cancelled_);
+    list_.erase(it_);
+    if (list_.empty()) {
+      map_.erase(key_);
+    }
+    cancelled_ = true;
+  }
+
+  MapOfList& map_;
+  std::list<Value>& list_;
+  // Because of absl::flat_hash_map iterator instability we have to keep a copy of the key
+  const Key key_;
+  typename MapOfList::mapped_type::iterator it_;
+  bool cancelled_;
+};
 } // namespace Envoy

source/extensions/clusters/dynamic_forward_proxy/cluster.cc

@@ -25,8 +25,7 @@ Cluster::Cluster(
                                        added_via_api, factory_context.dispatcher().timeSource()),
       dns_cache_manager_(cache_manager_factory.get()),
       dns_cache_(dns_cache_manager_->getCache(config.dns_cache_config())),
-      update_callbacks_handle_(dns_cache_->addUpdateCallbacks(*this)), local_info_(local_info),
-      host_map_(std::make_shared<HostInfoMap>()) {
+      update_callbacks_handle_(dns_cache_->addUpdateCallbacks(*this)), local_info_(local_info) {
   // Block certain TLS context parameters that don't make sense on a cluster-wide scale. We will
   // support these parameters dynamically in the future. This is not an exhaustive list of
   // parameters that don't make sense but should be the most obvious ones that a user might set
@@ -44,123 +43,116 @@ Cluster::Cluster(
 
 void Cluster::startPreInit() {
   // If we are attaching to a pre-populated cache we need to initialize our hosts.
-  auto existing_hosts = dns_cache_->hosts();
-  if (!existing_hosts.empty()) {
-    std::shared_ptr<HostInfoMap> new_host_map;
-    std::unique_ptr<Upstream::HostVector> hosts_added;
-    for (const auto& existing_host : existing_hosts) {
-      addOrUpdateWorker(existing_host.first, existing_host.second, new_host_map, hosts_added);
-    }
-    swapAndUpdateMap(new_host_map, *hosts_added, {});
+  std::unique_ptr<Upstream::HostVector> hosts_added;
+  dns_cache_->iterateHostMap(
+      [&](absl::string_view host, const Common::DynamicForwardProxy::DnsHostInfoSharedPtr& info) {
+        addOrUpdateHost(host, info, hosts_added);
+      });
+  if (hosts_added) {
+    updatePriorityState(*hosts_added, {});
   }
-
   onPreInitComplete();
 }
 
-void Cluster::addOrUpdateWorker(
-    const std::string& host,
+void Cluster::addOrUpdateHost(
+    absl::string_view host,
     const Extensions::Common::DynamicForwardProxy::DnsHostInfoSharedPtr& host_info,
-    std::shared_ptr<HostInfoMap>& new_host_map,
     std::unique_ptr<Upstream::HostVector>& hosts_added) {
-  // We should never get a host with no address from the cache.
-  ASSERT(host_info->address() != nullptr);
-
-  // NOTE: Right now we allow a DNS cache to be shared between multiple clusters. Though we have
-  // connection/request circuit breakers on the cluster, we don't have any way to control the
-  // maximum hosts on a cluster. We currently assume that host data shared via shared pointer is a
-  // marginal memory cost above that already used by connections and requests, so relying on
-  // connection/request circuit breakers is sufficient. We may have to revisit this in the future.
-
-  HostInfoMapSharedPtr current_map = getCurrentHostMap();
-  const auto host_map_it = current_map->find(host);
-  if (host_map_it != current_map->end()) {
-    // If we only have an address change, we can do that swap inline without any other updates.
-    // The appropriate R/W locking is in place to allow this. The details of this locking are:
-    //  - Hosts are not thread local, they are global.
-    //  - We take a read lock when reading the address and a write lock when changing it.
-    //  - Address updates are very rare.
-    //  - Address reads are only done when a connection is being made and a "real" host
-    //    description is created or the host is queried via the admin endpoint. Both of
-    //    these operations are relatively rare and the read lock is held for a short period
-    //    of time.
-    //
-    // TODO(mattklein123): Right now the dynamic forward proxy / DNS cache works similar to how
-    //                     logical DNS works, meaning that we only store a single address per
-    //                     resolution. It would not be difficult to also expose strict DNS
-    //                     semantics, meaning the cache would expose multiple addresses and the
-    //                     cluster would create multiple logical hosts based on those addresses.
-    //                     We will leave this is a follow up depending on need.
-    ASSERT(host_info == host_map_it->second.shared_host_info_);
-    ASSERT(host_map_it->second.shared_host_info_->address() !=
-           host_map_it->second.logical_host_->address());
-    ENVOY_LOG(debug, "updating dfproxy cluster host address '{}'", host);
-    host_map_it->second.logical_host_->setNewAddress(host_info->address(), dummy_lb_endpoint_);
-    return;
-  }
+  Upstream::LogicalHostSharedPtr emplaced_host;
+  {
+    absl::WriterMutexLock lock{&host_map_lock_};
+    // We should never get a host with no address from the cache.
+    ASSERT(host_info->address() != nullptr);
+
+    // NOTE: Right now we allow a DNS cache to be shared between multiple clusters. Though we have
+    // connection/request circuit breakers on the cluster, we don't have any way to control the
+    // maximum hosts on a cluster. We currently assume that host data shared via shared pointer is a
+    // marginal memory cost above that already used by connections and requests, so relying on
+    // connection/request circuit breakers is sufficient. We may have to revisit this in the future.
+    const auto host_map_it = host_map_.find(host);
+    if (host_map_it != host_map_.end()) {
+      // If we only have an address change, we can do that swap inline without any other updates.
+      // The appropriate R/W locking is in place to allow this. The details of this locking are:
+      //  - Hosts are not thread local, they are global.
+      //  - We take a read lock when reading the address and a write lock when changing it.
+      //  - Address updates are very rare.
+      //  - Address reads are only done when a connection is being made and a "real" host
+      //    description is created or the host is queried via the admin endpoint. Both of
+      //    these operations are relatively rare and the read lock is held for a short period
+      //    of time.
+      //
+      // TODO(mattklein123): Right now the dynamic forward proxy / DNS cache works similar to how
+      //                     logical DNS works, meaning that we only store a single address per
+      //                     resolution. It would not be difficult to also expose strict DNS
+      //                     semantics, meaning the cache would expose multiple addresses and the
+      //                     cluster would create multiple logical hosts based on those addresses.
+      //                     We will leave this is a follow up depending on need.
+      ASSERT(host_info == host_map_it->second.shared_host_info_);
+      ASSERT(host_map_it->second.shared_host_info_->address() !=
+             host_map_it->second.logical_host_->address());
+      ENVOY_LOG(debug, "updating dfproxy cluster host address '{}'", host);
+      host_map_it->second.logical_host_->setNewAddress(host_info->address(), dummy_lb_endpoint_);
+      return;
+    }
 
-  ENVOY_LOG(debug, "adding new dfproxy cluster host '{}'", host);
+    ENVOY_LOG(debug, "adding new dfproxy cluster host '{}'", host);
 
-  if (new_host_map == nullptr) {
-    new_host_map = std::make_shared<HostInfoMap>(*current_map);
+    emplaced_host = host_map_
+                        .try_emplace(host, host_info,
+                                     std::make_shared<Upstream::LogicalHost>(
+                                         info(), std::string{host}, host_info->address(),
+                                         dummy_locality_lb_endpoint_, dummy_lb_endpoint_, nullptr,
+                                         time_source_))
+                        .first->second.logical_host_;
   }
-  const auto emplaced =
-      new_host_map->try_emplace(host, host_info,
-                                std::make_shared<Upstream::LogicalHost>(
-                                    info(), host, host_info->address(), dummy_locality_lb_endpoint_,
-                                    dummy_lb_endpoint_, nullptr, time_source_));
+
+  ASSERT(emplaced_host);
   if (hosts_added == nullptr) {
     hosts_added = std::make_unique<Upstream::HostVector>();
   }
-  hosts_added->emplace_back(emplaced.first->second.logical_host_);
+  hosts_added->emplace_back(emplaced_host);
 }
 
 void Cluster::onDnsHostAddOrUpdate(
     const std::string& host,
     const Extensions::Common::DynamicForwardProxy::DnsHostInfoSharedPtr& host_info) {
-  std::shared_ptr<HostInfoMap> new_host_map;
+  ENVOY_LOG(debug, "Adding host info for {}", host);
+
   std::unique_ptr<Upstream::HostVector> hosts_added;
-  addOrUpdateWorker(host, host_info, new_host_map, hosts_added);
+  addOrUpdateHost(host, host_info, hosts_added);
   if (hosts_added != nullptr) {
-    ASSERT(!new_host_map->empty());
     ASSERT(!hosts_added->empty());
-    // Swap in the new map. This will be picked up when the per-worker LBs are recreated via
-    // the host set update.
-    swapAndUpdateMap(new_host_map, *hosts_added, {});
+    updatePriorityState(*hosts_added, {});
   }
 }
 
-void Cluster::swapAndUpdateMap(const HostInfoMapSharedPtr& new_hosts_map,
-                               const Upstream::HostVector& hosts_added,
-                               const Upstream::HostVector& hosts_removed) {
-  {
-    absl::WriterMutexLock lock(&host_map_lock_);
-    host_map_ = new_hosts_map;
-  }
-
+void Cluster::updatePriorityState(const Upstream::HostVector& hosts_added,
+                                  const Upstream::HostVector& hosts_removed) {
   Upstream::PriorityStateManager priority_state_manager(*this, local_info_, nullptr);
   priority_state_manager.initializePriorityFor(dummy_locality_lb_endpoint_);
-  for (const auto& host : (*new_hosts_map)) {
-    priority_state_manager.registerHostForPriority(host.second.logical_host_,
-                                                   dummy_locality_lb_endpoint_);
+  {
+    absl::ReaderMutexLock lock{&host_map_lock_};
+    for (const auto& host : host_map_) {
+      priority_state_manager.registerHostForPriority(host.second.logical_host_,
+                                                     dummy_locality_lb_endpoint_);
+    }
   }
   priority_state_manager.updateClusterPrioritySet(
       0, std::move(priority_state_manager.priorityState()[0].first), hosts_added, hosts_removed,
       absl::nullopt, absl::nullopt);
 }
 
 void Cluster::onDnsHostRemove(const std::string& host) {
-  HostInfoMapSharedPtr current_map = getCurrentHostMap();
-  const auto host_map_it = current_map->find(host);
-  ASSERT(host_map_it != current_map->end());
-  const auto new_host_map = std::make_shared<HostInfoMap>(*current_map);
   Upstream::HostVector hosts_removed;
-  hosts_removed.emplace_back(host_map_it->second.logical_host_);
-  new_host_map->erase(host);
-  ENVOY_LOG(debug, "removing dfproxy cluster host '{}'", host);
-
-  // Swap in the new map. This will be picked up when the per-worker LBs are recreated via
-  // the host set update.
-  swapAndUpdateMap(new_host_map, {}, hosts_removed);
+  {
+    absl::WriterMutexLock lock{&host_map_lock_};
+    const auto host_map_it = host_map_.find(host);
+    ASSERT(host_map_it != host_map_.end());
+    hosts_removed.emplace_back(host_map_it->second.logical_host_);
+    host_map_.erase(host);
+    ENVOY_LOG(debug, "removing dfproxy cluster host '{}'", host);
+  }
+  updatePriorityState({}, hosts_removed);
 }
 
 Upstream::HostConstSharedPtr
@@ -179,13 +171,15 @@ Cluster::LoadBalancer::chooseHost(Upstream::LoadBalancerContext* context) {
   if (host.empty()) {
     return nullptr;
   }
-
-  const auto host_it = host_map_->find(host);
-  if (host_it == host_map_->end()) {
-    return nullptr;
-  } else {
-    host_it->second.shared_host_info_->touch();
-    return host_it->second.logical_host_;
+  {
+    absl::ReaderMutexLock lock{&cluster_.host_map_lock_};
+    const auto host_it = cluster_.host_map_.find(host);
+    if (host_it == cluster_.host_map_.end()) {
+      return nullptr;
+    } else {
+      host_it->second.shared_host_info_->touch();
+      return host_it->second.logical_host_;
+    }
   }
 }
 

source/extensions/clusters/dynamic_forward_proxy/cluster.h

@@ -52,10 +52,9 @@ class Cluster : public Upstream::BaseDynamicClusterImpl,
   };
 
   using HostInfoMap = absl::flat_hash_map<std::string, HostInfo>;
-  using HostInfoMapSharedPtr = std::shared_ptr<const HostInfoMap>;
 
   struct LoadBalancer : public Upstream::LoadBalancer {
-    LoadBalancer(const HostInfoMapSharedPtr& host_map) : host_map_(host_map) {}
+    LoadBalancer(const Cluster& cluster) : cluster_(cluster) {}
 
     // Upstream::LoadBalancer
     Upstream::HostConstSharedPtr chooseHost(Upstream::LoadBalancerContext* context) override;
@@ -64,16 +63,14 @@ class Cluster : public Upstream::BaseDynamicClusterImpl,
       return nullptr;
     }
 
-    const HostInfoMapSharedPtr host_map_;
+    const Cluster& cluster_;
   };
 
   struct LoadBalancerFactory : public Upstream::LoadBalancerFactory {
     LoadBalancerFactory(Cluster& cluster) : cluster_(cluster) {}
 
     // Upstream::LoadBalancerFactory
-    Upstream::LoadBalancerPtr create() override {
-      return std::make_unique<LoadBalancer>(cluster_.getCurrentHostMap());
-    }
+    Upstream::LoadBalancerPtr create() override { return std::make_unique<LoadBalancer>(cluster_); }
 
     Cluster& cluster_;
   };
@@ -90,19 +87,15 @@ class Cluster : public Upstream::BaseDynamicClusterImpl,
     Cluster& cluster_;
   };
 
-  HostInfoMapSharedPtr getCurrentHostMap() {
-    absl::ReaderMutexLock lock(&host_map_lock_);
-    return host_map_;
-  }
-
   void
-  addOrUpdateWorker(const std::string& host,
-                    const Extensions::Common::DynamicForwardProxy::DnsHostInfoSharedPtr& host_info,
-                    std::shared_ptr<HostInfoMap>& new_host_map,
-                    std::unique_ptr<Upstream::HostVector>& hosts_added);
-  void swapAndUpdateMap(const HostInfoMapSharedPtr& new_hosts_map,
-                        const Upstream::HostVector& hosts_added,
-                        const Upstream::HostVector& hosts_removed);
+  addOrUpdateHost(absl::string_view host,
+                  const Extensions::Common::DynamicForwardProxy::DnsHostInfoSharedPtr& host_info,
+                  std::unique_ptr<Upstream::HostVector>& hosts_added)
+      ABSL_LOCKS_EXCLUDED(host_map_lock_);
+
+  void updatePriorityState(const Upstream::HostVector& hosts_added,
+                           const Upstream::HostVector& hosts_removed)
+      ABSL_LOCKS_EXCLUDED(host_map_lock_);
 
   const Extensions::Common::DynamicForwardProxy::DnsCacheManagerSharedPtr dns_cache_manager_;
   const Extensions::Common::DynamicForwardProxy::DnsCacheSharedPtr dns_cache_;
@@ -112,8 +105,8 @@ class Cluster : public Upstream::BaseDynamicClusterImpl,
   const envoy::config::endpoint::v3::LbEndpoint dummy_lb_endpoint_;
   const LocalInfo::LocalInfo& local_info_;
 
-  absl::Mutex host_map_lock_;
-  HostInfoMapSharedPtr host_map_ ABSL_GUARDED_BY(host_map_lock_);
+  mutable absl::Mutex host_map_lock_;
+  HostInfoMap host_map_ ABSL_GUARDED_BY(host_map_lock_);
 
   friend class ClusterFactory;
   friend class ClusterTest;

source/extensions/common/dynamic_forward_proxy/dns_cache.h

@@ -23,7 +23,7 @@ class DnsHostInfo {
    * Returns the host's currently resolved address. This address may change periodically due to
    * async re-resolution.
    */
-  virtual Network::Address::InstanceConstSharedPtr address() PURE;
+  virtual Network::Address::InstanceConstSharedPtr address() const PURE;
 
   /**
    * Returns the host that was actually resolved via DNS. If port was originally specified it will
@@ -172,10 +172,14 @@ class DnsCache {
    */
   virtual AddUpdateCallbacksHandlePtr addUpdateCallbacks(UpdateCallbacks& callbacks) PURE;
 
+  using IterateHostMapCb = std::function<void(absl::string_view, const DnsHostInfoSharedPtr&)>;
+
   /**
-   * @return all hosts currently stored in the cache.
+   * Iterates over all entries in the cache, calling a callback for each entry
+   *
+   * @param iterate_callback the callback to invoke for each entry in the cache
    */
-  virtual absl::flat_hash_map<std::string, DnsHostInfoSharedPtr> hosts() PURE;
+  virtual void iterateHostMap(IterateHostMapCb iterate_callback) PURE;
 
   /**
    * Retrieve the DNS host info of a given host currently stored in the cache.