dnsdist: Add constants and helper functions to the IDState

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index 359988cfe13e123bf7e4fe804e26ec6a234e7b04..c2b63c08ab3b61cb5a2a327d31b32f63db744bf6 100644 (file)
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -546,7 +546,7 @@ try {
         IDState* ids = &dss->idStates[queryId];
         int64_t usageIndicator = ids->usageIndicator;
 
-        if(usageIndicator < 0) {
+        if(!IDState::isInUse(usageIndicator)) {
           /* the corresponding state is marked as not in use, meaning that:
              - it was already cleaned up by another thread and the state is gone ;
              - we already got a response for this query and this one is a duplicate.
@@ -572,7 +572,7 @@ try {
         bool isDoH = du != nullptr;
         /* atomically mark the state as available, but only if it has not been altered
            in the meantime */
-        if (ids->usageIndicator.compare_exchange_strong(usageIndicator, -1)) {
+        if (ids->tryMarkUnused(usageIndicator)) {
           /* clear the potential DOHUnit asap, it's ours now
            and since we just marked the state as unused,
            someone could overwrite it. */
@@ -1591,15 +1591,13 @@ static void processUDPQuery(ClientState& cs, LocalHolders& holders, const struct
     /* that means that the state was in use, possibly with an allocated
        DOHUnit that we will need to handle, but we can't touch it before
        confirming that we now own this state */
-    if (ids->usageIndicator != -1) {
+    if (ids->isInUse()) {
       du = ids->du;
     }
 
     /* we atomically replace the value, we now own this state */
-    auto generation = ids->generation++;
-    int64_t oldUsage = ids->usageIndicator.exchange(generation);
-    if(oldUsage < 0) {
-      /* the value was -1, meaning that the state was not in use.
+    if (!ids->markAsUsed()) {
+      /* the state was not in use.
          we reset 'du' because it might have still been in use when we read it. */
       du = nullptr;
       ++ss->outstanding;
@@ -2093,14 +2091,14 @@ static void healthChecksThread()
       
       for(IDState& ids  : dss->idStates) { // timeouts
         int64_t usageIndicator = ids.usageIndicator;
-        if(usageIndicator >=0 && ids.age++ > g_udpTimeout) {
-          /* We set usageIndicator to -1 as soon as possible
+        if(IDState::isInUse(usageIndicator) && ids.age++ > g_udpTimeout) {
+          /* We mark the state as unused as soon as possible
              to limit the risk of racing with the
              responder thread.
           */
           auto oldDU = ids.du;
 
-          if (!ids.usageIndicator.compare_exchange_strong(usageIndicator, -1)) {
+          if (!ids.tryMarkUnused(usageIndicator)) {
             /* this state has been altered in the meantime,
                don't go anywhere near it */
             continue;

diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index 72ec1a143e58c57efddf6dbbe43001dd2254c047..a9078c700ee5e4f4ea8ea067d6b307050bd07953 100644 (file)
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -551,6 +551,39 @@ struct IDState
     tempFailureTTL = orig.tempFailureTTL;
   }
 
+  static const int64_t unusedIndicator = -1;
+
+  static bool isInUse(int64_t usageIndicator)
+  {
+    return usageIndicator != unusedIndicator;
+  }
+
+  bool isInUse() const
+  {
+    return usageIndicator != unusedIndicator;
+  }
+
+  /* return true if the value has been successfully replaced meaning that
+     no-one updated the usage indicator in the meantime */
+  bool tryMarkUnused(int64_t expectedUsageIndicator)
+  {
+    return usageIndicator.compare_exchange_strong(expectedUsageIndicator, unusedIndicator);
+  }
+
+  /* mark as unused no matter what, return true if the state was in use before */
+  bool markAsUsed()
+  {
+    auto currentGeneration = generation++;
+    return markAsUsed(currentGeneration);
+  }
+
+  /* mark as unused no matter what, return true if the state was in use before */
+  bool markAsUsed(int64_t currentGeneration)
+  {
+    int64_t oldUsage = usageIndicator.exchange(currentGeneration);
+    return oldUsage != unusedIndicator;
+  }
+
   /* We use this value to detect whether this state is in use.
      For performance reasons we don't want to use a lock here, but that means
      we need to be very careful when modifying this value. Modifications happen
@@ -584,8 +617,8 @@ struct IDState
      wrapping around if necessary, and we set an atomic signed 64-bit value, so that we still have -1
      when the state is unused and the value of our counter otherwise.
   */
-  std::atomic<int64_t> usageIndicator{-1};  // set to <0 to indicate this state is empty   // 4
-  std::atomic<uint32_t> generation{0}; // increased every time a state is used, to be able to detect an ABA issue
+  std::atomic<int64_t> usageIndicator{unusedIndicator};  // set to unusedIndicator to indicate this state is empty   // 8
+  std::atomic<uint32_t> generation{0}; // increased every time a state is used, to be able to detect an ABA issue    // 4
   ComboAddress origRemote;                                    // 28
   ComboAddress origDest;                                      // 28
   StopWatch sentTime;                                         // 16

diff --git a/pdns/dnsdistdist/doh.cc b/pdns/dnsdistdist/doh.cc
index 29b5665b158584e34f6ef8ea7ae72481ca730dd0..f03a0bc64bed0d5953e123b5e6f647b0a86329e8 100644 (file)
--- a/pdns/dnsdistdist/doh.cc
+++ b/pdns/dnsdistdist/doh.cc
@@ -238,7 +238,7 @@ static int processDOHQuery(DOHUnit* du)
     IDState* ids = &ss->idStates[idOffset];
     ids->age = 0;
     DOHUnit* oldDU = nullptr;
-    if (ids->usageIndicator != -1) {
+    if (ids->isInUse()) {
       /* that means that the state was in use, possibly with an allocated
          DOHUnit that we will need to handle, but we can't touch it before
          confirming that we now own this state */
@@ -247,19 +247,18 @@ static int processDOHQuery(DOHUnit* du)
 
     /* we atomically replace the value, we now own this state */
     int64_t generation = ids->generation++;
-    int64_t oldUsage = ids->usageIndicator.exchange(generation);
-    if (oldUsage < 0) {
-      /* the value was -1, meaning that the state was not in use.
+    if (!ids->markAsUsed(generation)) {
+      /* the state was not in use.
          we reset 'oldDU' because it might have still been in use when we read it. */
       oldDU = nullptr;
       ++ss->outstanding;
     }
     else {
+      ids->du = nullptr;
+      /* we are reusing a state, no change in outstanding but if there was an existing DOHUnit we need
+         to handle it because it's about to be overwritten. */
       ++ss->reuseds;
       ++g_stats.downstreamTimeouts;
-      /* we are reusing a state, if there was an existing DOHUnit we need
-         to handle it because it's about to be overwritten. */
-      ids->du = nullptr;
       handleDOHTimeout(oldDU);
     }
 
@@ -295,7 +294,7 @@ static int processDOHQuery(DOHUnit* du)
       /* we are about to handle the error, make sure that
          this pointer is not accessed when the state is cleaned,
          but first check that it still belongs to us */
-      if (ids->usageIndicator.compare_exchange_strong(generation, -1)) {
+      if (ids->tryMarkUnused(generation)) {
         ids->du = nullptr;
         --ss->outstanding;
       }