From 307994e76533a10ec7d951ea42c143a78b81f143 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Thu, 6 Feb 2020 17:26:49 +0100
Subject: [PATCH] auth: Prepare the caches' buckets in advance

---
 pdns/auth-packetcache.cc |  8 ++++++++
 pdns/auth-packetcache.hh |  5 +++++
 pdns/auth-querycache.cc  |  8 ++++++++
 pdns/auth-querycache.hh  |  5 +++++
 pdns/common_startup.cc   |  1 +
 pdns/dbdnsseckeeper.cc   | 16 +++++++++++++---
 pdns/dnsseckeeper.hh     |  5 +++++
 7 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/pdns/auth-packetcache.cc b/pdns/auth-packetcache.cc
index 094c1f63ba..7759c8c071 100644
--- a/pdns/auth-packetcache.cc
+++ b/pdns/auth-packetcache.cc
@@ -58,6 +58,14 @@ AuthPacketCache::~AuthPacketCache()
   }
 }
 
+void AuthPacketCache::MapCombo::reserve(size_t numberOfEntries)
+{
+#if BOOST_VERSION >= 105600
+  WriteLock wl(&d_mut);
+  d_map.get<HashTag>().reserve(numberOfEntries);
+#endif /* BOOST_VERSION >= 105600 */
+}
+
 bool AuthPacketCache::get(DNSPacket& p, DNSPacket& cached)
 {
   if(!d_ttl) {
diff --git a/pdns/auth-packetcache.hh b/pdns/auth-packetcache.hh
index 91b41528bc..5fa324c991 100644
--- a/pdns/auth-packetcache.hh
+++ b/pdns/auth-packetcache.hh
@@ -66,6 +66,9 @@ public:
   void setMaxEntries(uint64_t maxEntries) 
   {
     d_maxEntries = maxEntries;
+    for (auto& shard : d_maps) {
+      shard.reserve(maxEntries / d_maps.size());
+    }
   }
   void setTTL(uint32_t ttl)
   {
@@ -115,6 +118,8 @@ private:
     MapCombo(const MapCombo&) = delete; 
     MapCombo& operator=(const MapCombo&) = delete;
 
+    void reserve(size_t numberOfEntries);
+
     pthread_rwlock_t d_mut;
     cmap_t d_map;
   };
diff --git a/pdns/auth-querycache.cc b/pdns/auth-querycache.cc
index e790e9e165..b9b13602c1 100644
--- a/pdns/auth-querycache.cc
+++ b/pdns/auth-querycache.cc
@@ -59,6 +59,14 @@ AuthQueryCache::~AuthQueryCache()
   }
 }
 
+void AuthQueryCache::MapCombo::reserve(size_t numberOfEntries)
+{
+#if BOOST_VERSION >= 105600
+  WriteLock wl(&d_mut);
+  d_map.get<HashTag>().reserve(numberOfEntries);
+#endif /* BOOST_VERSION >= 105600 */
+}
+
 // called from ueberbackend
 bool AuthQueryCache::getEntry(const DNSName &qname, const QType& qtype, vector<DNSZoneRecord>& value, int zoneID)
 {
diff --git a/pdns/auth-querycache.hh b/pdns/auth-querycache.hh
index 2dd27c835d..5fc4c55e30 100644
--- a/pdns/auth-querycache.hh
+++ b/pdns/auth-querycache.hh
@@ -56,6 +56,9 @@ public:
   void setMaxEntries(uint64_t maxEntries)
   {
     d_maxEntries = maxEntries;
+    for (auto& shard : d_maps) {
+      shard.reserve(maxEntries / d_maps.size());
+    }
   }
 private:
 
@@ -98,6 +101,8 @@ private:
     MapCombo(const MapCombo &) = delete; 
     MapCombo & operator=(const MapCombo &) = delete;
 
+    void reserve(size_t numberOfEntries);
+
     pthread_rwlock_t d_mut;
     cmap_t d_map;
   };
diff --git a/pdns/common_startup.cc b/pdns/common_startup.cc
index c8c24ede78..feae5c4330 100644
--- a/pdns/common_startup.cc
+++ b/pdns/common_startup.cc
@@ -540,6 +540,7 @@ void mainthread()
    PC.setTTL(::arg().asNum("cache-ttl"));
    PC.setMaxEntries(::arg().asNum("max-packet-cache-entries"));
    QC.setMaxEntries(::arg().asNum("max-cache-entries"));
+   DNSSECKeeper::setMaxEntries(::arg().asNum("max-cache-entries"));
 
    if (!PC.enabled() && ::arg().mustDo("log-dns-queries")) {
      g_log<<Logger::Warning<<"Packet cache disabled, logging queries without HIT/MISS"<<endl;
diff --git a/pdns/dbdnsseckeeper.cc b/pdns/dbdnsseckeeper.cc
index a3bd4ac45c..e7f42a1477 100644
--- a/pdns/dbdnsseckeeper.cc
+++ b/pdns/dbdnsseckeeper.cc
@@ -51,6 +51,7 @@ pthread_rwlock_t DNSSECKeeper::s_metacachelock = PTHREAD_RWLOCK_INITIALIZER;
 pthread_rwlock_t DNSSECKeeper::s_keycachelock = PTHREAD_RWLOCK_INITIALIZER;
 AtomicCounter DNSSECKeeper::s_ops;
 time_t DNSSECKeeper::s_last_prune;
+size_t DNSSECKeeper::s_maxEntries = 0;
 
 bool DNSSECKeeper::doesDNSSEC()
 {
@@ -860,12 +861,21 @@ void DNSSECKeeper::cleanup()
   if(now.tv_sec - s_last_prune > (time_t)(30)) {
     {
         WriteLock l(&s_metacachelock);
-        pruneCollection<SequencedTag>(*this, s_metacache, ::arg().asNum("max-cache-entries"));
+        pruneCollection<SequencedTag>(*this, s_metacache, s_maxEntries);
     }
     {
         WriteLock l(&s_keycachelock);
-        pruneCollection<SequencedTag>(*this, s_keycache, ::arg().asNum("max-cache-entries"));
+        pruneCollection<SequencedTag>(*this, s_keycache, s_maxEntries);
     }
-    s_last_prune=time(0);
+    s_last_prune = time(nullptr);
   }
 }
+
+void DNSSECKeeper::setMaxEntries(size_t maxEntries)
+{
+  s_maxEntries = maxEntries;
+#if BOOST_VERSION >= 105600
+  WriteLock wl(&s_keycachelock);
+  s_keycache.get<KeyCacheTag>().reserve(s_maxEntries);
+#endif /* BOOST_VERSION >= 105600 */
+}
diff --git a/pdns/dnsseckeeper.hh b/pdns/dnsseckeeper.hh
index 92df79203d..c2fa1c0ab6 100644
--- a/pdns/dnsseckeeper.hh
+++ b/pdns/dnsseckeeper.hh
@@ -235,6 +235,9 @@ public:
   void getSoaEdit(const DNSName& zname, std::string& value);
   bool unSecureZone(const DNSName& zone, std::string& error, std::string& info);
   bool rectifyZone(const DNSName& zone, std::string& error, std::string& info, bool doTransaction);
+
+  static void setMaxEntries(size_t maxEntries);
+
 private:
 
 
@@ -277,6 +280,7 @@ private:
       sequenced<tag<SequencedTag>>
     >
   > keycache_t;
+
   typedef multi_index_container<
     METACacheEntry,
     indexed_by<
@@ -298,6 +302,7 @@ private:
   static pthread_rwlock_t s_keycachelock;
   static AtomicCounter s_ops;
   static time_t s_last_prune;
+  static size_t s_maxEntries;
 
 public:
   void preRemoval(const KeyCacheEntry&)
-- 
2.47.2