From f907c798eb6ed7184af888a3b447a4901c002d65 Mon Sep 17 00:00:00 2001
From: hanwei <chuxiking@163.com>
Date: Tue, 19 Sep 2023 15:40:05 +0800
Subject: [PATCH] feature: add support to config dns client cache size

---
 crates/shadowsocks-service/src/config.rs           |  7 +++++++
 .../src/local/dns/client_cache.rs                  |  2 +-
 crates/shadowsocks-service/src/local/dns/server.rs | 14 ++++++++++----
 crates/shadowsocks-service/src/local/mod.rs        |  2 ++
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/crates/shadowsocks-service/src/config.rs b/crates/shadowsocks-service/src/config.rs
index ccb3a8b2631f..577ff96c339a 100644
--- a/crates/shadowsocks-service/src/config.rs
+++ b/crates/shadowsocks-service/src/config.rs
@@ -860,6 +860,11 @@ pub struct LocalConfig {
     /// Sending DNS query through proxy to this address
     #[cfg(feature = "local-dns")]
     pub remote_dns_addr: Option<Address>,
+    // client cache size
+    // if a lot of `create connection` observed in log,
+    // increase the size
+    #[cfg(feature = "local-dns")]
+    pub client_cache_size: Option<usize>,
 
     /// Tun interface's name
     ///
@@ -919,6 +924,8 @@ impl LocalConfig {
             local_dns_addr: None,
             #[cfg(feature = "local-dns")]
             remote_dns_addr: None,
+            #[cfg(feature = "local-dns")]
+            client_cache_size: Some(5),
 
             #[cfg(feature = "local-tun")]
             tun_interface_name: None,
diff --git a/crates/shadowsocks-service/src/local/dns/client_cache.rs b/crates/shadowsocks-service/src/local/dns/client_cache.rs
index 558fb1827174..4e919a194247 100644
--- a/crates/shadowsocks-service/src/local/dns/client_cache.rs
+++ b/crates/shadowsocks-service/src/local/dns/client_cache.rs
@@ -196,7 +196,7 @@ impl DnsClientCache {
                 }
             }
             Entry::Vacant(vac) => {
-                let mut q = VecDeque::with_capacity(5);
+                let mut q = VecDeque::with_capacity(self.max_client_per_addr);
                 q.push_back(client);
                 vac.insert(q);
             }
diff --git a/crates/shadowsocks-service/src/local/dns/server.rs b/crates/shadowsocks-service/src/local/dns/server.rs
index 08f68eaab6f6..01a8b7839db5 100644
--- a/crates/shadowsocks-service/src/local/dns/server.rs
+++ b/crates/shadowsocks-service/src/local/dns/server.rs
@@ -53,6 +53,7 @@ pub struct DnsBuilder {
     remote_addr: Address,
     bind_addr: ServerAddr,
     balancer: PingBalancer,
+    client_cache_size: usize,
 }
 
 impl DnsBuilder {
@@ -62,9 +63,10 @@ impl DnsBuilder {
         local_addr: NameServerAddr,
         remote_addr: Address,
         balancer: PingBalancer,
+        client_cache_size: usize,
     ) -> DnsBuilder {
         let context = ServiceContext::new();
-        DnsBuilder::with_context(Arc::new(context), bind_addr, local_addr, remote_addr, balancer)
+        DnsBuilder::with_context(Arc::new(context), bind_addr, local_addr, remote_addr, balancer, client_cache_size)
     }
 
     /// Create with an existed `context`
@@ -74,6 +76,7 @@ impl DnsBuilder {
         local_addr: NameServerAddr,
         remote_addr: Address,
         balancer: PingBalancer,
+        client_cache_size: usize,
     ) -> DnsBuilder {
         DnsBuilder {
             context,
@@ -82,6 +85,7 @@ impl DnsBuilder {
             remote_addr,
             bind_addr,
             balancer,
+            client_cache_size,
         }
     }
 
@@ -92,7 +96,8 @@ impl DnsBuilder {
 
     /// Build DNS server
     pub async fn build(self) -> io::Result<Dns> {
-        let client = Arc::new(DnsClient::new(self.context.clone(), self.balancer, self.mode));
+        let client = Arc::new(DnsClient::new(self.context.clone(), self.balancer, self.mode,
+                                             self.client_cache_size));
 
         let local_addr = Arc::new(self.local_addr);
         let remote_addr = Arc::new(self.remote_addr);
@@ -589,10 +594,11 @@ struct DnsClient {
 }
 
 impl DnsClient {
-    fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mode: Mode) -> DnsClient {
+    fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mode: Mode,
+           client_cache_size: usize) -> DnsClient {
         DnsClient {
             context,
-            client_cache: DnsClientCache::new(5),
+            client_cache: DnsClientCache::new(client_cache_size),
             mode,
             balancer,
             attempts: 2,
diff --git a/crates/shadowsocks-service/src/local/mod.rs b/crates/shadowsocks-service/src/local/mod.rs
index 942488cd1311..c2d0fb6455aa 100644
--- a/crates/shadowsocks-service/src/local/mod.rs
+++ b/crates/shadowsocks-service/src/local/mod.rs
@@ -353,6 +353,7 @@ impl Server {
                     let mut server_builder = {
                         let local_addr = local_config.local_dns_addr.expect("missing local_dns_addr");
                         let remote_addr = local_config.remote_dns_addr.expect("missing remote_dns_addr");
+                        let client_cache_size = local_config.client_cache_size.unwrap();
 
                         DnsBuilder::with_context(
                             context.clone(),
@@ -360,6 +361,7 @@ impl Server {
                             local_addr.clone(),
                             remote_addr.clone(),
                             balancer,
+                            client_cache_size,
                         )
                     };
                     server_builder.set_mode(local_config.mode);