thread pool: add ExecutionContext

include/opendht/thread_pool.h

@@ -93,4 +93,71 @@ private:
     void schedule();
 };
 
+class OPENDHT_PUBLIC ExecutionContext {
+public:
+    ExecutionContext(ThreadPool& pool)
+     : threadPool_(pool), state_(std::make_shared<SharedState>())
+    {}
+
+    ~ExecutionContext() {
+        state_->destroy();
+    }
+
+    /** Wait for ongoing tasks to complete execution and drop other pending tasks */
+    void stop() {
+        state_->destroy(false);
+    }
+
+    void run(std::function<void()>&& task) {
+        std::lock_guard<std::mutex> lock(state_->mtx);
+        if (state_->shutdown_) return;
+        state_->pendingTasks++;
+        threadPool_.get().run([task = std::move(task), state = state_] {
+            state->run(task);
+        });
+    }
+
+private:
+    struct SharedState {
+        std::mutex mtx {};
+        std::condition_variable cv {};
+        unsigned pendingTasks {0};
+        unsigned ongoingTasks {0};
+        /** When true, prevents new tasks to be scheduled */
+        bool shutdown_ {false};
+        /** When true, prevents scheduled tasks to be executed */
+        std::atomic_bool destroyed {false};
+
+        void destroy(bool wait = true) {
+            std::unique_lock<std::mutex> lock(mtx);
+            if (destroyed) return;
+            if (wait) {
+                cv.wait(lock, [this] { return pendingTasks == 0 && ongoingTasks == 0; });
+            }
+            shutdown_ = true;
+            if (not wait) {
+                cv.wait(lock, [this] { return ongoingTasks == 0; });
+            }
+            destroyed = true;
+        }
+
+        void run(const std::function<void()>& task) {
+            {
+                std::lock_guard<std::mutex> lock(mtx);
+                pendingTasks--;
+                ongoingTasks++;
+            }
+                if (destroyed) return;
+                task();
+            {
+                std::lock_guard<std::mutex> lock(mtx);
+                ongoingTasks--;
+                cv.notify_all();
+            }
+        }
+    };
+    std::reference_wrapper<ThreadPool> threadPool_;
+    std::shared_ptr<SharedState> state_;
+};
+
 }

src/dht.cpp

@@ -873,11 +873,6 @@ Dht::listenTo(const InfoHash& id, sa_family_t af, ValueCallback cb, Value::Filte
 size_t
 Dht::listen(const InfoHash& id, ValueCallback cb, Value::Filter f, Where where)
 {
-    if (not id) {
-        if (logger_)
-            logger_->w(id, "Listen called with invalid key");
-        return 0;
-    }
     scheduler.syncTime();
 
     auto token = ++listener_token;
@@ -960,9 +955,7 @@ struct GetStatus : public OpStatus {
 void
 Dht::put(const InfoHash& id, Sp<Value> val, DoneCallback callback, time_point created, bool permanent)
 {
-    if (not id or not val) {
-        if (logger_)
-            logger_->w(id, "Put called with invalid key or value");
+    if (not val) {
         if (callback)
             callback(false, {});
         return;
@@ -1032,13 +1025,6 @@ bool callbackWrapper(Cb get_cb, DoneCallback done_cb, const std::vector<Sp<T>>&
 void
 Dht::get(const InfoHash& id, GetCallback getcb, DoneCallback donecb, Value::Filter&& filter, Where&& where)
 {
-    if (not id) {
-        if (logger_)
-            logger_->w(id, "Get called with invalid key");
-        if (donecb)
-            donecb(false, {});
-        return;
-    }
     scheduler.syncTime();
 
     auto op = std::make_shared<GetStatus<std::map<Value::Id, Sp<Value>>>>();
@@ -1079,13 +1065,6 @@ Dht::get(const InfoHash& id, GetCallback getcb, DoneCallback donecb, Value::Filt
 
 void Dht::query(const InfoHash& id, QueryCallback cb, DoneCallback done_cb, Query&& q)
 {
-    if (not id) {
-        if (logger_)
-            logger_->w(id, "Query called with invalid key");
-        if (done_cb)
-            done_cb(false, {});
-        return;
-    }
     scheduler.syncTime();
     auto op = std::make_shared<GetStatus<std::vector<Sp<FieldValueIndex>>>>();
     auto f = q.where.getFilter();

src/dhtrunner.cpp

@@ -294,11 +294,6 @@ DhtRunner::shutdown(ShutdownCallback cb, bool stop) {
         if (dht_)
             dht_->shutdown(onShutdown, stop);
     });
-    // Force execution of pending ops
-    if (getStatus() == NodeStatus::Connecting) {
-        status4 = NodeStatus::Disconnected;
-        status6 = NodeStatus::Disconnected;
-    }
     cv.notify_all();
 }
 
@@ -373,7 +368,6 @@ DhtRunner::join()
     {
         std::lock_guard<std::mutex> lck(dht_mtx);
         resetDht();
-        std::lock_guard<std::mutex> lck2(storage_mtx);
         status4 = NodeStatus::Disconnected;
         status6 = NodeStatus::Disconnected;
     }
@@ -709,11 +703,8 @@ DhtRunner::loop_()
     NodeStatus nstatus4 = dht->updateStatus(AF_INET);
     NodeStatus nstatus6 = dht->updateStatus(AF_INET6);
     if (nstatus4 != status4 || nstatus6 != status6) {
-        {
-            std::lock_guard<std::mutex> lck(storage_mtx);
-            status4 = nstatus4;
-            status6 = nstatus6;
-        }
+        status4 = nstatus4;
+        status6 = nstatus6;
         if (statusCb)
             statusCb(status4, status6);
     }

tests/threadpooltester.cpp

@@ -59,7 +59,7 @@ ThreadPoolTester::testExecutor()
     auto executor8 = std::make_shared<dht::Executor>(pool, 8);
 
     constexpr unsigned N = 64 * 1024;
-    std::atomic_uint count1 {0};
+    unsigned count1 {0};
     std::atomic_uint count4 {0};
     std::atomic_uint count8 {0};
     for (unsigned i=0; i<N; i++) {
@@ -69,7 +69,7 @@ ThreadPoolTester::testExecutor()
     }
 
     auto start = clock::now();
-    while ((count1.load() != N ||
+    while ((count1 != N ||
             count4.load() != N ||
             count8.load() != N) && clock::now() - start < std::chrono::seconds(20))
     {
@@ -78,11 +78,28 @@ ThreadPoolTester::testExecutor()
     executor1.reset();
     executor4.reset();
     executor8.reset();
-    CPPUNIT_ASSERT_EQUAL(N, count1.load());
+    CPPUNIT_ASSERT_EQUAL(N, count1);
     CPPUNIT_ASSERT_EQUAL(N, count4.load());
     CPPUNIT_ASSERT_EQUAL(N, count8.load());
 }
 
+void
+ThreadPoolTester::testContext()
+{
+    std::atomic_uint count {0};
+    constexpr unsigned N = 64 * 1024;
+
+    {
+        dht::ExecutionContext ctx(dht::ThreadPool::computation());
+        for (unsigned i=0; i<N; i++) {
+            ctx.run([&] { count++; });
+        }
+    }
+
+    CPPUNIT_ASSERT_EQUAL(N, count.load());
+
+}
+
 void
 ThreadPoolTester::tearDown() {
 }

tests/threadpooltester.h

@@ -29,6 +29,7 @@ class ThreadPoolTester : public CppUnit::TestFixture {
     CPPUNIT_TEST_SUITE(ThreadPoolTester);
     CPPUNIT_TEST(testThreadPool);
     CPPUNIT_TEST(testExecutor);
+    CPPUNIT_TEST(testContext);
     CPPUNIT_TEST_SUITE_END();
 
  public:
@@ -43,6 +44,7 @@ class ThreadPoolTester : public CppUnit::TestFixture {
 
     void testThreadPool();
     void testExecutor();
+    void testContext();
 };
 
 }  // namespace test