From a196aaffc475ba01c0ea508bb8a5be8912ed44d0 Mon Sep 17 00:00:00 2001
From: Matthias Killat <matthias.killat{at}gmx.de>
Date: Mon, 16 Nov 2020 20:51:45 +0100
Subject: [PATCH] fixed invalidation of token

---
 include/waitset/waitset.hpp   | 33 +++++++++++++--------------------
 include/waitset/waittoken.hpp |  9 ++++-----
 test_waitset.cpp              | 26 ++++++++++++++------------
 3 files changed, 31 insertions(+), 37 deletions(-)

diff --git a/include/waitset/waitset.hpp b/include/waitset/waitset.hpp
index 22feecb..74e9ccc 100644
--- a/include/waitset/waitset.hpp
+++ b/include/waitset/waitset.hpp
@@ -63,26 +63,6 @@ class WaitSet
         return WaitToken(node);
     }
 
-    bool remove(WaitToken &token)
-    {
-        auto id = token.id();
-        token.invalidate();
-        return remove(id);
-    }
-
-    bool remove(id_t id)
-    {
-        std::lock_guard g(m_nodesMutex);
-        auto &node = m_nodes[id];
-
-        //only remove it if no token references it anymore
-        if (node.numReferences() <= 0)
-        {
-            return m_nodes.remove(id);
-        }
-        return false;
-    }
-
     //todo: we could also add a timed wait but in theory this can be done with a condition that is set to true by a timer
 
     //we can only have one waiter for proper operation (concurrent condition result reset would cause problems!)
@@ -182,6 +162,19 @@ class WaitSet
 
     //we do not want to add this to the container itself, we need a scoped lock for iteration
     std::mutex m_nodesMutex;
+
+    bool remove(id_t id)
+    {
+        std::lock_guard g(m_nodesMutex);
+        auto &node = m_nodes[id];
+
+        //only remove it if no token references it anymore
+        if (node.numReferences() <= 0)
+        {
+            return m_nodes.remove(id);
+        }
+        return false;
+    }
 };
 
 //can only be defined when WaitSet is fully defined
diff --git a/include/waitset/waittoken.hpp b/include/waitset/waittoken.hpp
index 967ae4f..5607b25 100644
--- a/include/waitset/waittoken.hpp
+++ b/include/waitset/waittoken.hpp
@@ -24,10 +24,7 @@ class WaitToken
     {
         if (&rhs != this)
         {
-            if (isValid())
-            {
-                m_waitNode->decrementRefCount();
-            }
+            invalidate();
 
             m_waitNode = rhs.m_waitNode;
             if (isValid())
@@ -47,6 +44,7 @@ class WaitToken
     {
         if (&rhs != this)
         {
+            invalidate();
             m_waitNode = rhs.m_waitNode;
             rhs.m_waitNode = nullptr;
         }
@@ -87,7 +85,8 @@ class WaitToken
     {
         if (isValid())
         {
-            if (m_waitNode->decrementRefCount())
+            m_waitNode->decrementRefCount();
+            if (m_waitNode->numReferences() <= 0)
             {
                 m_waitNode->tryDelete();
             }
diff --git a/test_waitset.cpp b/test_waitset.cpp
index a574f51..92aec12 100644
--- a/test_waitset.cpp
+++ b/test_waitset.cpp
@@ -98,19 +98,21 @@ int main(int argc, char **argv)
     auto maybeToken = waitSet.add(always_true);
     if (!maybeToken.has_value())
     {
-        waitSet.remove(token3); //a copy of token1, if we do not remove it we cannot free the node (shared by token1 and token3)
-        std::cout << "could not get another token" << std::endl;
-        if (waitSet.remove(token1))
+        token1.invalidate();
+        token3.invalidate(); //invalidating both tokens (referencing the same node) will delete the node from the waitset
+
+        //we should have space again
+        //maybeToken = waitSet.add(always_false);
+        maybeToken = waitSet.add(always_true);
+        if (maybeToken.has_value())
+        {
+            token1 = *maybeToken;
+            token3 = token1;
+            std::cout << "regenerated token1 and its copy token3" << std::endl;
+        }
+        else
         {
-            //we should have space again
-            //maybeToken = waitSet.add(always_false);
-            maybeToken = waitSet.add(always_true);
-            if (maybeToken.has_value())
-            {
-                token1 = *maybeToken;
-                token3 = token1;
-                std::cout << "regenerated token1 and its copy token3" << std::endl;
-            }
+            std::cout << "could not get another token" << std::endl;
         }
     }