dnsdist: Refactor some very similar functions in the TCP limits code

diff --git a/pdns/dnsdistdist/dnsdist-concurrent-connections.cc b/pdns/dnsdistdist/dnsdist-concurrent-connections.cc
index bf3169a168b6027b8bde732c6021f0b854c87ccc..c2f2c646575f5cbc0575f090f653bc682869a8d2 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-concurrent-connections.cc
+++ b/pdns/dnsdistdist/dnsdist-concurrent-connections.cc
@@ -269,12 +269,8 @@ void IncomingConcurrentTCPConnectionsManager::banClientFor(const ComboAddress& f
   vinfolog("Banned TCP client %s for %d seconds", from.toStringWithPort(), seconds);
 }
 
-void IncomingConcurrentTCPConnectionsManager::accountClosedTCPConnection(const ComboAddress& from)
+static void editEntryIfPresent(const ComboAddress& from, const std::function<void(const ClientEntry& entry)>& callback)
 {
-  const auto maxConnsPerClient = dnsdist::configuration::getImmutableConfiguration().d_maxTCPConnectionsPerClient;
-  if (maxConnsPerClient == 0) {
-    return;
-  }
   auto addr = getRange(from);
   auto shardID = getShardID(addr);
   {
@@ -283,47 +279,44 @@ void IncomingConcurrentTCPConnectionsManager::accountClosedTCPConnection(const C
     if (it == db->end()) {
       return;
     }
-    auto& count = it->d_concurrentConnections;
-    count--;
+    callback(*it);
   }
 }
 
+void IncomingConcurrentTCPConnectionsManager::accountClosedTCPConnection(const ComboAddress& from)
+{
+  const auto maxConnsPerClient = dnsdist::configuration::getImmutableConfiguration().d_maxTCPConnectionsPerClient;
+  if (maxConnsPerClient == 0) {
+    return;
+  }
+  editEntryIfPresent(from, [](const ClientEntry& entry) {
+    auto& count = entry.d_concurrentConnections;
+    count--;
+  });
+}
+
 void IncomingConcurrentTCPConnectionsManager::accountTLSNewSession(const ComboAddress& from)
 {
-  const auto maxRate = dnsdist::configuration::getImmutableConfiguration().d_maxTLSNewSessionsRatePerClient > 0;
+  const auto maxRate = dnsdist::configuration::getImmutableConfiguration().d_maxTLSNewSessionsRatePerClient;
   if (maxRate == 0) {
     return;
   }
-  auto addr = getRange(from);
-  auto shardID = getShardID(addr);
-  {
-    auto db = s_tcpClientsConnectionMetrics.at(shardID).lock();
-    auto it = db->find(addr);
-    if (it == db->end()) {
-      return;
-    }
-    auto& count = getCurrentClientActivity(*it, time(nullptr)).tlsNewSessions;
+  editEntryIfPresent(from, [](const ClientEntry& entry) {
+    auto& count = getCurrentClientActivity(entry, time(nullptr)).tlsNewSessions;
     count++;
-  }
+  });
 }
 
 void IncomingConcurrentTCPConnectionsManager::accountTLSResumedSession(const ComboAddress& from)
 {
-  const auto maxRate = dnsdist::configuration::getImmutableConfiguration().d_maxTLSResumedSessionsRatePerClient > 0;
+  const auto maxRate = dnsdist::configuration::getImmutableConfiguration().d_maxTLSResumedSessionsRatePerClient;
   if (maxRate == 0) {
     return;
   }
-  auto addr = getRange(from);
-  auto shardID = getShardID(addr);
-  {
-    auto db = s_tcpClientsConnectionMetrics.at(shardID).lock();
-    auto it = db->find(addr);
-    if (it == db->end()) {
-      return;
-    }
-    auto& count = getCurrentClientActivity(*it, time(nullptr)).tlsResumedSessions;
+  editEntryIfPresent(from, [](const ClientEntry& entry) {
+    auto& count = getCurrentClientActivity(entry, time(nullptr)).tlsResumedSessions;
     count++;
-  }
+  });
 }
 
 }