MB-58402: Limit concurrency of bucket deletion

Add a new task type that both limits concurrency using a semaphore
and allows the function being run to return a sleep time after which
it wishes to run again.

Change-Id: Idf7c499cf5ae7d67290872fa3f594d51a0db28a7
Reviewed-on: https://review.couchbase.org/c/kv_engine/+/196661
Reviewed-by: Trond Norbye <[email protected]>
Tested-by: Build Bot <[email protected]>

Author: pavlosg

daemon/CMakeLists.txt

@@ -109,6 +109,7 @@ add_library(memcached_daemon STATIC
             tls_configuration.h
             tracing.cc
             tracing.h
+            yielding_limited_concurrency_task.cc
             yielding_task.cc
         )
 

daemon/protocol/mcbp/bucket_management_command_context.cc

@@ -16,7 +16,7 @@
 #include <daemon/one_shot_limited_concurrency_task.h>
 #include <daemon/session_cas.h>
 #include <daemon/settings.h>
-#include <daemon/yielding_task.h>
+#include <daemon/yielding_limited_concurrency_task.h>
 #include <executor/executorpool.h>
 #include <logger/logger.h>
 #include <memcached/config_parser.h>
@@ -157,35 +157,36 @@ cb::engine_errc BucketManagementCommandContext::remove() {
 
     std::string taskname{"Delete bucket [" + name + "]"};
     using namespace std::chrono_literals;
-    ExecutorPool::get()->schedule(std::make_shared<YieldingTask>(
-            TaskId::Core_DeleteBucketTask,
-            taskname,
-            [destroyer = std::move(*optDestroyer),
-             client = &cookie,
-             nm = std::move(name)]() mutable
+    auto task = [destroyer = std::move(*optDestroyer),
+                 client = &cookie,
+                 nm = std::move(name)]() mutable
             -> std::optional<std::chrono::duration<double>> {
-                auto& connection = client->getConnection();
-                cb::engine_errc status;
-                try {
-                    status = destroyer.drive();
-                } catch (const std::runtime_error& error) {
-                    LOG_WARNING(
-                            "{}: An error occurred while deleting bucket [{}]: "
-                            "{}",
-                            connection.getId(),
-                            nm,
-                            error.what());
-                    status = cb::engine_errc::failed;
-                }
-                if (status == cb::engine_errc::would_block) {
-                    // destroyer hasn't completed yet (waiting for connections
-                    // or operations), try again in 10ms
-                    return {10ms};
-                }
-                client->notifyIoComplete(status);
-                return {};
-            },
-            100ms));
+        auto& connection = client->getConnection();
+        cb::engine_errc status;
+        try {
+            status = destroyer.drive();
+        } catch (const std::runtime_error& error) {
+            LOG_WARNING("{}: An error occurred while deleting bucket [{}]: {}",
+                        connection.getId(),
+                        nm,
+                        error.what());
+            status = cb::engine_errc::failed;
+        }
+        if (status == cb::engine_errc::would_block) {
+            // destroyer hasn't completed yet (waiting for connections
+            // or operations), try again in 10ms
+            return {10ms};
+        }
+        client->notifyIoComplete(status);
+        return {};
+    };
+    ExecutorPool::get()->schedule(
+            std::make_shared<YieldingLimitedConcurrencyTask>(
+                    TaskId::Core_DeleteBucketTask,
+                    std::move(taskname),
+                    std::move(task),
+                    ConcurrencySemaphores::instance().bucket_management,
+                    100ms));
 
     state = State::Done;
     return cb::engine_errc::would_block;

daemon/tasks_test.cc

@@ -11,6 +11,7 @@
 #include <folly/portability/GMock.h>
 #include <folly/portability/GTest.h>
 
+#include <daemon/yielding_limited_concurrency_task.h>
 #include <daemon/yielding_task.h>
 #include <executor/executorpool.h>
 #include <executor/fake_executorpool.h>
@@ -37,19 +38,17 @@ class TasksTest : public ::testing::Test {
     }
 };
 
-TEST_F(TasksTest, YieldingTaskCalledAgain) {
-    // test that a YieldingTask is called again if it returns a snooze time,
-    // and is not called again when returning a nullopt to indicate it is
-    // "done"
+using MockTaskFunc = testing::StrictMock<
+        testing::MockFunction<std::optional<std::chrono::duration<double>>()>>;
+
+/**
+ * Test that a task is called again if it returns a snooze time,
+ * and is not called again when returning a nullopt to indicate it is "done"
+ */
+static void testTaskCalledAgain(ExTask task, MockTaskFunc& mockTaskFunc) {
     using namespace testing;
-    StrictMock<MockFunction<std::optional<std::chrono::duration<double>>()>>
-            mockTaskFunc;
 
-    ExecutorPool::get()->schedule(
-            std::make_shared<YieldingTask>(TaskId::Core_DeleteBucketTask,
-                                           "foobar",
-                                           mockTaskFunc.AsStdFunction(),
-                                           std::chrono::seconds(30)));
+    ExecutorPool::get()->schedule(task);
 
     InSequence s;
 
@@ -67,17 +66,34 @@ TEST_F(TasksTest, YieldingTaskCalledAgain) {
                  std::logic_error);
 }
 
-TEST_F(TasksTest, YieldingTaskSnoozes) {
-    // test that a YieldingTask is correctly snooze()'ed
+TEST_F(TasksTest, YieldingTaskCalledAgain) {
+    MockTaskFunc func;
+    testTaskCalledAgain(
+            std::make_shared<YieldingTask>(TaskId::Core_DeleteBucketTask,
+                                           "foobar",
+                                           func.AsStdFunction(),
+                                           std::chrono::seconds(30)),
+            func);
+}
 
+TEST_F(TasksTest, YieldingLimitedConcurrencyTaskCalledAgain) {
+    MockTaskFunc func;
+    cb::AwaitableSemaphore semaphore{4};
+    testTaskCalledAgain(std::make_shared<YieldingLimitedConcurrencyTask>(
+                                TaskId::Core_DeleteBucketTask,
+                                "foobar",
+                                func.AsStdFunction(),
+                                semaphore,
+                                std::chrono::seconds(30)),
+                        func);
+}
+
+/**
+ * Test that a task is correctly snooze()'ed
+ */
+static void testTaskSnoozes(ExTask task, MockTaskFunc& mockTaskFunc) {
     using namespace testing;
-    StrictMock<MockFunction<std::optional<std::chrono::duration<double>>()>>
-            mockTaskFunc;
 
-    auto task = std::make_shared<YieldingTask>(TaskId::Core_DeleteBucketTask,
-                                               "foobar",
-                                               mockTaskFunc.AsStdFunction(),
-                                               std::chrono::seconds(30));
     ExecutorPool::get()->schedule(task);
 
     using namespace std::chrono;
@@ -95,4 +111,26 @@ TEST_F(TasksTest, YieldingTaskSnoozes) {
 
     EXPECT_EQ(TASK_SNOOZED, task->getState());
     EXPECT_GE(task->getWaketime(), beforeTime + 10min);
-}
+}
+
+TEST_F(TasksTest, YieldingTaskSnoozes) {
+    MockTaskFunc func;
+    testTaskSnoozes(
+            std::make_shared<YieldingTask>(TaskId::Core_DeleteBucketTask,
+                                           "foobar",
+                                           func.AsStdFunction(),
+                                           std::chrono::seconds(30)),
+            func);
+}
+
+TEST_F(TasksTest, YieldingLimitedConcurrencyTaskSnoozes) {
+    MockTaskFunc func;
+    cb::AwaitableSemaphore semaphore{4};
+    testTaskSnoozes(std::make_shared<YieldingLimitedConcurrencyTask>(
+                            TaskId::Core_DeleteBucketTask,
+                            "foobar",
+                            func.AsStdFunction(),
+                            semaphore,
+                            std::chrono::seconds(30)),
+                    func);
+}

daemon/yielding_limited_concurrency_task.cc

@@ -0,0 +1,53 @@
+/*
+ *     Copyright 2023-Present Couchbase, Inc.
+ *
+ *   Use of this software is governed by the Business Source License included
+ *   in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
+ *   in that file, in accordance with the Business Source License, use of this
+ *   software will be governed by the Apache License, Version 2.0, included in
+ *   the file licenses/APL2.txt.
+ */
+
+#include "yielding_limited_concurrency_task.h"
+#include <logger/logger.h>
+
+YieldingLimitedConcurrencyTask::YieldingLimitedConcurrencyTask(
+        TaskId id,
+        std::string name,
+        YieldingFunc function,
+        cb::AwaitableSemaphore& semaphore,
+        std::chrono::microseconds expectedRuntime)
+    : LimitedConcurrencyTask(NoBucketTaskable::instance(), id, semaphore, true),
+      name(std::move(name)),
+      function(std::move(function)),
+      expectedRuntime(std::move(expectedRuntime)) {
+}
+
+bool YieldingLimitedConcurrencyTask::runInner() {
+    auto guard = acquireOrWait();
+    if (!guard) {
+        // Could not acquire a token, queued for notification.
+        // Already snooze()-ed forever, just return true to
+        // reschedule.
+        return true;
+    }
+
+    // Do concurrency-limited work
+    try {
+        auto optSleepTime = function();
+        if (optSleepTime) {
+            snooze(optSleepTime->count());
+            // Task _does_ wish to run again
+            return true;
+        }
+    } catch (const std::exception& e) {
+        LOG_CRITICAL(
+                "YieldingLimitedConcurrencyTask::runInner(\"{}\"):"
+                " received exception: {}",
+                name,
+                e.what());
+    }
+
+    // Not running again
+    return false;
+}

daemon/yielding_limited_concurrency_task.h

@@ -0,0 +1,62 @@
+/*
+ *     Copyright 2023-Present Couchbase, Inc.
+ *
+ *   Use of this software is governed by the Business Source License included
+ *   in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
+ *   in that file, in accordance with the Business Source License, use of this
+ *   software will be governed by the Apache License, Version 2.0, included in
+ *   the file licenses/APL2.txt.
+ */
+#pragma once
+
+#include "nobucket_taskable.h"
+#include <executor/limited_concurrency_task.h>
+#include <functional>
+#include <optional>
+
+/**
+ * YieldingLimitedConcurrencyTask accepts a "generic" task which runs the
+ * provided function until completed to avoid having to duplicate all of the
+ * boilerplate code for a simple task.
+ *
+ * A semaphore is used to limit the number of tasks run concurrently.
+ * The function may return a sleep time after which it wishes to run again
+ * (0s may be used to yield the thread and run again asap), or a
+ * nullopt to indicate it is complete.
+ */
+class YieldingLimitedConcurrencyTask : public LimitedConcurrencyTask {
+public:
+    using YieldingFunc =
+            std::function<std::optional<std::chrono::duration<double>>()>;
+
+    /**
+     * Create a new instance of the YieldingLimitedConcurrencyTask
+     *
+     * @param id The task identifier for the task
+     * @param name The name for the task
+     * @param function The function to run on the executor
+     * @param semaphore The semaphore to use for concurrency limiting
+     * @param expectedRuntime Maximum duration `function` is expected to run
+     */
+    YieldingLimitedConcurrencyTask(TaskId id,
+                                   std::string name,
+                                   YieldingFunc function,
+                                   cb::AwaitableSemaphore& semaphore,
+                                   std::chrono::microseconds expectedRuntime =
+                                           std::chrono::milliseconds(100));
+
+    std::string getDescription() const override {
+        return name;
+    }
+
+    std::chrono::microseconds maxExpectedDuration() const override {
+        return expectedRuntime;
+    }
+
+protected:
+    bool runInner() override;
+
+    const std::string name;
+    const YieldingFunc function;
+    const std::chrono::microseconds expectedRuntime;
+};