MB-70482: Add topkeys controller

And return an error if topkeys is already running including
a UUID

Change-Id: I311ff239eb8c854b540a3e9fc63919e680bd8566
Reviewed-on: https://review.couchbase.org/c/kv_engine/+/239905
Tested-by: Build Bot <build@couchbase.com>
Reviewed-by: Jim Walker <jim@couchbase.com>

Author: trondn

daemon/CMakeLists.txt

@@ -118,6 +118,8 @@ add_library(memcached_daemon STATIC
             token_auth_data.h
             top_keys.cc
             top_keys.h
+            top_keys_controller.cc
+            top_keys_controller.h
             tracing.cc
             tracing.h
             yielding_limited_concurrency_task.cc

daemon/ioctl.cc

@@ -17,6 +17,7 @@
 #include "external_auth_manager_thread.h"
 #include "front_end_thread.h"
 #include "settings.h"
+#include "top_keys_controller.h"
 #include "tracing.h"
 #include "utilities/string_utilities.h"
 #include <logger/logger.h>
@@ -118,20 +119,6 @@ static cb::engine_errc ioctlGetTopkeysStop(Cookie& cookie,
                                            const StrToStrMap& args,
                                            std::string& value,
                                            cb::mcbp::Datatype& datatype) {
-    std::shared_ptr<cb::trace::topkeys::Collector> collector;
-    FrontEndThread::forEach(
-            [&collector](auto& t) {
-                if (!t.keyTrace) {
-                    return;
-                }
-
-                if (!collector) {
-                    collector = t.keyTrace;
-                }
-                t.keyTrace.reset();
-            },
-            true);
-
     std::size_t limit = 100;
     if (args.contains("limit")) {
         try {
@@ -145,25 +132,39 @@ static cb::engine_errc ioctlGetTopkeysStop(Cookie& cookie,
         }
     }
 
-    if (!collector) {
-        cookie.setErrorContext(
-                "Failed to collect topkeys - tracing not active");
-        LOG_WARNING_CTX("Failed to collect topkeys - tracing not active",
-                        {"conn_id", cookie.getConnection().getId()});
-        return cb::engine_errc::failed;
+    cb::uuid::uuid_t uuid;
+    if (args.contains("uuid")) {
+        try {
+            uuid = cb::uuid::from_string(args.find("uuid")->second);
+        } catch (const std::exception& exception) {
+            LOG_ERROR_CTX("Failed to parse uuid",
+                          {"conn_id", cookie.getConnection().getId()},
+                          {"error", exception.what()});
+            return cb::engine_errc::invalid_arguments;
+        }
     }
 
-    try {
-        nlohmann::json json = collector->getResults(limit);
-        value = json.dump();
-    } catch (const std::exception& exception) {
-        LOG_ERROR_CTX("Failed to get trace data",
-                      {"conn_id", cookie.getConnection().getId()},
-                      {"error", exception.what()});
-        return cb::engine_errc::failed;
+    auto [status, json] =
+            cb::trace::topkeys::Controller::instance().stop(uuid, limit);
+    if (status == cb::engine_errc::success) {
+        try {
+            value = json.dump();
+        } catch (const std::exception& exception) {
+            LOG_ERROR_CTX(
+                    "Failed to get trace data. Trace data will be discarded",
+                    {"conn_id", cookie.getConnection().getId()},
+                    {"error", exception.what()});
+            return cb::engine_errc::failed;
+        }
+        datatype = cb::mcbp::Datatype::JSON;
+        return cb::engine_errc::success;
     }
-    datatype = cb::mcbp::Datatype::JSON;
-    return cb::engine_errc::success;
+
+    if (json.is_object()) {
+        cookie.setErrorJsonExtras(json);
+    }
+    cookie.setErrorContext("Failed to start topkeys collection");
+    return status;
 }
 
 cb::engine_errc ioctl_get_property(Cookie& cookie,
@@ -325,7 +326,9 @@ static cb::engine_errc ioctlSetServerlessUnitSize(Cookie& cookie,
 
 static cb::engine_errc ioctlSetTopkeysStart(Cookie& cookie,
                                             const StrToStrMap& args,
-                                            const std::string&) {
+                                            const std::string&,
+                                            std::string& result,
+                                            cb::mcbp::Datatype& datatype) {
     std::size_t limit = 10000;
     if (args.contains("limit")) {
         try {
@@ -402,27 +405,28 @@ static cb::engine_errc ioctlSetTopkeysStart(Cookie& cookie,
         }
     }
 
-    auto expiry_time = cb::time::steady_clock::now() +
-                       std::chrono::seconds(expected_duration);
-
-    auto collector = cb::trace::topkeys::Collector::create(
-            limit, shards, expiry_time, bucket_filter);
-    if (!collector) {
-        cookie.setErrorContext("Failed to create topkeys collector");
-        LOG_WARNING_CTX("Failed to create topkeys collector",
-                        {"conn_id", cookie.getConnection().getId()});
-        return cb::engine_errc::failed;
+    const auto [status, uuid] =
+            cb::trace::topkeys::Controller::instance().create(
+                    limit,
+                    shards,
+                    std::chrono::seconds(expected_duration),
+                    bucket_filter);
+    nlohmann::json json = {{"uuid", to_string(uuid)}};
+    if (status == cb::engine_errc::success) {
+        result = json.dump();
+        datatype = cb::mcbp::Datatype::JSON;
+    } else {
+        cookie.setErrorJsonExtras(json);
+        cookie.setErrorContext("Failed to start topkeys collection");
     }
-
-    FrontEndThread::forEach([&collector](auto& t) { t.keyTrace = collector; },
-                            true);
-
-    return cb::engine_errc::success;
+    return status;
 }
 
 cb::engine_errc ioctl_set_property(Cookie& cookie,
                                    const std::string& key,
-                                   const std::string& value) {
+                                   const std::string& value,
+                                   std::string& result,
+                                   cb::mcbp::Datatype& datatype) {
     std::pair<std::string, StrToStrMap> request;
 
     try {
@@ -431,6 +435,9 @@ cb::engine_errc ioctl_set_property(Cookie& cookie,
         return cb::engine_errc::invalid_arguments;
     }
 
+    result.clear();
+    datatype = cb::mcbp::Datatype::Raw;
+
     auto& manager = cb::ioctl::Manager::getInstance();
     auto* id = manager.lookup(request.first);
     if (id) {
@@ -471,7 +478,8 @@ cb::engine_errc ioctl_set_property(Cookie& cookie,
             }
             return cb::engine_errc::not_supported;
         case cb::ioctl::Id::TopkeysStart:
-            return ioctlSetTopkeysStart(cookie, request.second, value);
+            return ioctlSetTopkeysStart(
+                    cookie, request.second, value, result, datatype);
 
         case cb::ioctl::Id::TraceDumpBegin: // may only be used with Get
         case cb::ioctl::Id::TraceDumpGet: // may only be used with Get

daemon/ioctl.h

@@ -32,9 +32,18 @@ cb::engine_errc ioctl_get_property(Cookie& cookie,
 
 /**
  * Attempts to set property `key` to the value `value`.
- * If the property could be written, return cb::engine_errc::success.
- * Otherwise returns a status code indicating why the write failed.
+ *
+ * @param cookie represents the command
+ * @param key is the name of the control
+ * @param value is the value to set for the key
+ * @param result is an out parameter to send back to the client
+ * @param datatype is another out parameter describing the value in result
+ *
+ * @return cb::engine_errc::success if the property could be written,
+ *         otherwise a status code indicating why the write failed.
  */
 cb::engine_errc ioctl_set_property(Cookie& cookie,
                                    const std::string& key,
-                                   const std::string& value);
+                                   const std::string& value,
+                                   std::string& result,
+                                   cb::mcbp::Datatype& datatype);

daemon/one_shot_task.h

@@ -21,13 +21,18 @@ class OneShotTask : public GlobalTask {
      * @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 expectedRuntime_ The expected runtime of the task
+     * @param initialSleepTime The initial sleep time before the task is first
+     *                         run
      */
     OneShotTask(TaskId id,
                 std::string name_,
                 std::function<void()> function_,
                 std::chrono::microseconds expectedRuntime_ =
-                        std::chrono::milliseconds(100))
-        : GlobalTask(NoBucketTaskable::instance(), id, 0, true),
+                        std::chrono::milliseconds(100),
+                std::chrono::nanoseconds initialSleepTime =
+                        std::chrono::nanoseconds(0))
+        : GlobalTask(NoBucketTaskable::instance(), id, initialSleepTime, true),
           name(std::move(name_)),
           function(std::move(function_)),
           expectedRuntime(std::move(expectedRuntime_)) {

daemon/protocol/mcbp/ioctl_command_context.cc

@@ -28,7 +28,7 @@ cb::engine_errc IoctlCommandContext::execute() {
     const auto key = std::string{cookie.getRequest().getKeyString()};
     if (req.getClientOpcode() == cb::mcbp::ClientOpcode::IoctlSet) {
         const auto value = std::string{cookie.getRequest().getValueString()};
-        return ioctl_set_property(cookie, key, value);
+        return ioctl_set_property(cookie, key, value, response, datatype);
     }
     return ioctl_get_property(cookie, key, response, datatype);
 }

daemon/top_keys.cc

@@ -11,7 +11,11 @@
 #include "top_keys.h"
 
 #include "bucket_manager.h"
+#include "nobucket_taskable.h"
+#include "one_shot_task.h"
+#include "top_keys_controller.h"
 
+#include <executor/executorpool.h>
 #include <memcached/dockey_view.h>
 #include <memcached/limits.h>
 #include <memcached/storeddockey.h>
@@ -73,12 +77,21 @@ void to_json(nlohmann::json& json, const Result& result) {
     }
 }
 
+Collector::~Collector() {
+    if (expiry_remover) {
+        expiry_remover->cancel();
+    }
+}
+
 class CountingCollector : public Collector {
 public:
     CountingCollector(std::size_t max,
                       std::size_t shards,
-                      cb::time::steady_clock::time_point expiry_time)
-        : Collector(expiry_time), limit(max), shardmaps(shards) {
+                      std::chrono::seconds expiry_time,
+                      bool install_cleanup_task)
+        : Collector(expiry_time, install_cleanup_task),
+          limit(max),
+          shardmaps(shards) {
     }
 
     void access(const size_t bucket,
@@ -211,9 +224,10 @@ class FilteredCountingCollector : public CountingCollector {
 public:
     FilteredCountingCollector(std::size_t max,
                               std::size_t shards,
-                              cb::time::steady_clock::time_point expiry_time,
-                              std::vector<std::size_t> buckets)
-        : CountingCollector(max, shards, expiry_time),
+                              std::chrono::seconds expiry_time,
+                              std::vector<std::size_t> buckets,
+                              bool install_cleanup_task)
+        : CountingCollector(max, shards, expiry_time, install_cleanup_task),
           bucketfilter(std::move(buckets)) {
     }
 
@@ -229,20 +243,45 @@ class FilteredCountingCollector : public CountingCollector {
     std::vector<std::size_t> bucketfilter;
 };
 
-std::shared_ptr<Collector> Collector::create(
-        std::size_t num_keys,
-        std::size_t shards,
-        cb::time::steady_clock::time_point expiry_time,
-        std::vector<std::size_t> buckets) {
+std::shared_ptr<Collector> Collector::create(std::size_t num_keys,
+                                             std::size_t shards,
+                                             std::chrono::seconds expiry_time,
+                                             std::vector<std::size_t> buckets,
+                                             bool install_cleanup_task) {
     if (num_keys) {
         if (buckets.empty()) {
             return std::make_unique<CountingCollector>(
-                    num_keys, shards, expiry_time);
+                    num_keys, shards, expiry_time, install_cleanup_task);
         }
         return std::make_unique<FilteredCountingCollector>(
-                num_keys, shards, expiry_time, std::move(buckets));
+                num_keys,
+                shards,
+                expiry_time,
+                std::move(buckets),
+                install_cleanup_task);
     }
     return {};
 }
 
+Collector::Collector(std::chrono::seconds exp, bool install_cleanup_task)
+    : expiry_time(cb::time::steady_clock::now() + exp) {
+    if (install_cleanup_task) {
+        // Install the cleanup handler to run 1 second *after* it should
+        // expire as that would allow the front end threads to potentially
+        // disconnect from the trace *before* we remove it (which means
+        // we don't need to inject a task in each front end thread)
+        // and can release the memory immediately in this thread context
+        expiry_remover = std::make_shared<OneShotTask>(
+                TaskId::Core_ExpiredTopKeysRemover,
+                fmt::format("Expired TopKeys remover: {}", ::to_string(uuid)),
+                [collector_uuid = uuid]() {
+                    Controller::instance().onExpiry(collector_uuid);
+                },
+                std::chrono::milliseconds(20),
+                exp + std::chrono::seconds(1));
+
+        ExecutorPool::get()->schedule(expiry_remover);
+    }
+}
+
 } // namespace cb::trace::topkeys

daemon/top_keys.h

@@ -10,10 +10,11 @@
 
 #pragma once
 
-#include "memcached/dockey_view.h"
-
+#include <executor/globaltask.h>
+#include <memcached/dockey_view.h>
 #include <nlohmann/json_fwd.hpp>
 #include <platform/cb_time.h>
+#include <platform/uuid.h>
 #include <memory>
 #include <unordered_map>
 #include <vector>
@@ -45,7 +46,7 @@ void to_json(nlohmann::json& json, const Result& result);
 
 class Collector {
 public:
-    virtual ~Collector() = default;
+    virtual ~Collector();
 
     /**
      * Create a Collector instance used to track key access.
@@ -59,13 +60,23 @@ class Collector {
      *                bucket. For space efficiencty we use the bucket *id*
      *                and not the bucket name (as buckets typically don't
      *                come and go all the time)
+     * @param install_cleanup_task If true, the collector will install a task
+     *                             to remove itself after the expiry time
+     *                             has passed. This parameter is set to true
+     *                             in the normal case where the collector is
+     *                             created by the front end thread, but we
+     *                             would like to avoid that when running
+     *                             isolated unit tests testing other
+     *                             functionality of the collector.
+     *
+     * @return a shared pointer to the created Collector instance
      */
     static std::shared_ptr<Collector> create(
             std::size_t num_keys,
             std::size_t shards,
-            cb::time::steady_clock::time_point expiry_time =
-                    cb::time::steady_clock::now() + std::chrono::minutes(1),
-            std::vector<std::size_t> buckets = {});
+            std::chrono::seconds expiry_time = std::chrono::minutes(1),
+            std::vector<std::size_t> buckets = {},
+            bool install_cleanup_task = false);
 
     /// Is this collector expired or not (e.g. should we discard the collected
     /// data). The time is passed in as an argument to avoid having to fetch the
@@ -75,6 +86,12 @@ class Collector {
         return now >= expiry_time;
     }
 
+    /// Get the uuid of this collector (used by the clients to identify the
+    /// collector when they want to stop it and retrieve the collected data)
+    auto get_uuid() const {
+        return uuid;
+    }
+
     /**
      * Register access for a key in a given bucket
      *
@@ -94,8 +111,10 @@ class Collector {
     virtual Result getResults(size_t limit) const = 0;
 
 protected:
-    Collector(cb::time::steady_clock::time_point exp) : expiry_time(exp) {};
+    Collector(std::chrono::seconds exp, bool install_cleanup_task);
     const cb::time::steady_clock::time_point expiry_time;
+    const cb::uuid::uuid_t uuid = cb::uuid::random();
+    ExTask expiry_remover;
 };
 
 } // namespace cb::trace::topkeys