From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Wed, 17 Jul 2019 14:19:59 +0000 (+0200)
Subject: dnsdist: Add support for periodically reloading the CDB KVS
X-Git-Tag: dnsdist-1.4.0-rc2~9^2~17
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=e54075131b85ed7a22f11911a505955453b7076b;p=pdns

dnsdist: Add support for periodically reloading the CDB KVS
---

diff --git a/pdns/dnsdist-lua-bindings.cc b/pdns/dnsdist-lua-bindings.cc
index dc6483646..e83cdcc06 100644
--- a/pdns/dnsdist-lua-bindings.cc
+++ b/pdns/dnsdist-lua-bindings.cc
@@ -739,11 +739,11 @@ void setupLuaBindings(bool client)
 #endif /* HAVE_LMDB */
 
 #ifdef HAVE_CDB
-  g_lua.writeFunction("newCDBKVStore", [client](const std::string& fname) {
+  g_lua.writeFunction("newCDBKVStore", [client](const std::string& fname, time_t refreshDelay) {
     if (client) {
       return std::shared_ptr<KeyValueStore>(nullptr);
     }
-    return std::shared_ptr<KeyValueStore>(new CDBKVStore(fname));
+    return std::shared_ptr<KeyValueStore>(new CDBKVStore(fname, refreshDelay));
   });
 #endif /* HAVE_CDB */
 
diff --git a/pdns/dnsdistdist/dnsdist-kvs.cc b/pdns/dnsdistdist/dnsdist-kvs.cc
index 2d0c5dbff..c36a326ae 100644
--- a/pdns/dnsdistdist/dnsdist-kvs.cc
+++ b/pdns/dnsdistdist/dnsdist-kvs.cc
@@ -23,6 +23,8 @@
 #include "dnsdist-kvs.hh"
 #include "dolog.hh"
 
+#include <sys/stat.h>
+
 std::vector<std::string> KeyValueLookupKeySourceIP::getKeys(const DNSQuestion& dq)
 {
   std::vector<std::string> result;
@@ -84,11 +86,63 @@ bool LMDBKVStore::getValue(const std::string& key, std::string& value)
 
 #ifdef HAVE_CDB
 
+CDBKVStore::CDBKVStore(const std::string& fname, time_t refreshDelay): d_fname(fname), d_refreshDelay(refreshDelay)
+{
+  pthread_rwlock_init(&d_lock, nullptr);
+  d_refreshing.clear();
+
+  time_t now = time(nullptr);
+  if (d_refreshDelay > 0) {
+    d_nextCheck = now + d_refreshDelay;
+  }
+
+  refreshDBIfNeeded(now);
+}
+
+void CDBKVStore::refreshDBIfNeeded(time_t now)
+{
+  if (d_refreshing.test_and_set()) {
+    /* someone else is already refreshing */
+    return;
+  }
+
+  try {
+    struct stat st;
+    if (stat(d_fname.c_str(), &st) == 0) {
+      if (st.st_mtime > d_mtime) {
+        auto newCDB = make_unique<CDB>(d_fname);
+        {
+          WriteLock wl(&d_lock);
+          d_cdb = std::move(newCDB);
+        }
+        d_mtime = st.st_mtime;
+      }
+    }
+    else {
+      warnlog("Error while retrieving the last modification time of CDB database '%s': %s", d_fname, stringerror());
+    }
+    d_nextCheck = now + d_refreshDelay;
+  }
+  catch(...) {
+    d_refreshing.clear();
+    throw;
+  }
+}
+
 bool CDBKVStore::getValue(const std::string& key, std::string& value)
 {
+  time_t now = time(nullptr);
+
   try {
-    if (d_cdb.findOne(key, value)) {
-      return true;
+    if (d_nextCheck != 0 && now >= d_nextCheck) {
+      refreshDBIfNeeded(now);
+    }
+
+    {
+      ReadLock rl(&d_lock);
+      if (d_cdb && d_cdb->findOne(key, value)) {
+        return true;
+      }
     }
   }
   catch(const std::exception& e) {
diff --git a/pdns/dnsdistdist/dnsdist-kvs.hh b/pdns/dnsdistdist/dnsdist-kvs.hh
index 5edf4794c..758104a94 100644
--- a/pdns/dnsdistdist/dnsdist-kvs.hh
+++ b/pdns/dnsdistdist/dnsdist-kvs.hh
@@ -145,15 +145,20 @@ private:
 class CDBKVStore: public KeyValueStore
 {
 public:
-  CDBKVStore(const std::string& fname): d_cdb(fname), d_fname(fname)
-  {
-  }
+  CDBKVStore(const std::string& fname, time_t refreshDelay);
 
   bool getValue(const std::string& key, std::string& value) override;
 
 private:
-  CDB d_cdb;
+  void refreshDBIfNeeded(time_t now);
+
+  std::unique_ptr<CDB> d_cdb{nullptr};
   std::string d_fname;
+  pthread_rwlock_t d_lock;
+  time_t d_mtime{0};
+  time_t d_nextCheck{0};
+  time_t d_refreshDelay{0};
+  std::atomic_flag d_refreshing;
 };
 
 #endif /* HAVE_LMDB */
diff --git a/pdns/dnsdistdist/test-dnsdistkvs_cc.cc b/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
index 7bc15aa67..7e43eb6a3 100644
--- a/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
+++ b/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
@@ -166,18 +166,25 @@ BOOST_AUTO_TEST_CASE(test_CDB) {
     writer.close();
   }
 
-  auto cdb = std::unique_ptr<KeyValueStore>(new CDBKVStore(db));
+  auto cdb = std::unique_ptr<KeyValueStore>(new CDBKVStore(db, 0));
   doKVSChecks(cdb, lc, rem, dq);
 
   /*
   std::string value;
   DTime dt;
   dt.set();
-  for (size_t idx = 0; idx < 10000000; idx++) {
+  auto lookupKey = make_unique<KeyValueLookupKeySourceIP>();
+  for (size_t idx = 0; idx < 100000000; idx++) {
     auto keys = lookupKey->getKeys(dq);
     for (const auto& key : keys) {
-      BOOST_CHECK_EQUAL(cdb->getValue(key, value), true);
-      BOOST_CHECK_EQUAL(value, "this is the value of the tag");
+      if (!cdb->getValue(key, value)) {
+        cerr<<"key not found"<<endl;
+        break;
+      }
+      if (value != "this is the value for the remote addr") {
+        cerr<<"invalid value found"<<endl;
+        break;
+      }
     }
   }
   cerr<<dt.udiff()/1000/1000<<endl;