From 585cb5909c2cac8b2c1539074d2bbe5ad3f9986a Mon Sep 17 00:00:00 2001
From: Guohao Li <gehafearless@apache.org>
Date: Sun, 4 Feb 2024 17:00:59 +0800
Subject: [PATCH] feat(fqdn): Limit the cache size on dns_resolve

---
 src/runtime/rpc/dns_resolver.cpp | 17 +++++++++++++----
 src/runtime/rpc/dns_resolver.h   |  2 --
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/runtime/rpc/dns_resolver.cpp b/src/runtime/rpc/dns_resolver.cpp
index 06bf4aca17..2e52d1d68d 100644
--- a/src/runtime/rpc/dns_resolver.cpp
+++ b/src/runtime/rpc/dns_resolver.cpp
@@ -29,6 +29,7 @@
 #include "runtime/rpc/group_address.h"
 #include "runtime/rpc/group_host_port.h"
 #include "utils/autoref_ptr.h"
+#include "utils/flags.h"
 #include "utils/fmt_logging.h"
 
 METRIC_DEFINE_gauge_int64(server,
@@ -48,6 +49,12 @@ METRIC_DEFINE_percentile_int64(server,
                                "The duration of resolving a host port by DNS lookup");
 namespace dsn {
 
+DSN_DEFINE_int32(network,
+                 max_count_for_dns_cache,
+                 100,
+                 "The size of dns_cache on dns_resolver. The part exceeding the "
+                 "cache size needs to be resolved from the system level each time");
+
 dns_resolver::dns_resolver()
     : METRIC_VAR_INIT_server(dns_resolver_cache_size),
       METRIC_VAR_INIT_server(dns_resolver_resolve_duration_ns),
@@ -95,10 +102,12 @@ error_s dns_resolver::resolve_addresses(const host_port &hp, std::vector<rpc_add
                       resolved_addresses[0]);
         }
 
-        utils::auto_write_lock l(_lock);
-        const auto it = _dns_cache.insert(std::make_pair(hp, resolved_addresses[0]));
-        if (it.second) {
-            METRIC_VAR_INCREMENT(dns_resolver_cache_size);
+        if (dsn_likely(_dns_cache.size() < FLAGS_max_count_for_dns_cache)) {
+            utils::auto_write_lock l(_lock);
+            const auto it = _dns_cache.insert(std::make_pair(hp, resolved_addresses[0]));
+            if (it.second) {
+                METRIC_VAR_INCREMENT(dns_resolver_cache_size);
+            }
         }
     }
 
diff --git a/src/runtime/rpc/dns_resolver.h b/src/runtime/rpc/dns_resolver.h
index aadb3287ff..925e801c03 100644
--- a/src/runtime/rpc/dns_resolver.h
+++ b/src/runtime/rpc/dns_resolver.h
@@ -35,8 +35,6 @@ namespace dsn {
 // the resolved result list.
 // If some host_port's rpc_address changes, you need to restart the Pegasus process to make it take
 // effect.
-// TODO(yingchun): Now the cache is unlimited, the cache size may be huge. Implement an expiration
-// mechanism to limit the cache size and make it possible to update the resolve result.
 class dns_resolver
 {
 public: