fix: run every EloqStore in a process under one brpc module

docs/architecture/02-runtime-and-lifecycle.md

@@ -29,11 +29,18 @@ A comment in `eloq_store.h` records one non-obvious ordering constraint:
 `cloud_service_` is declared before `shards_` because shard teardown triggers
 async HTTP cleanup that calls back into `CloudStorageService`.
 
-There is also a process-global `inline EloqStore *eloq_store` and a
-thread-local `inline thread_local Shard *shard` (`include/tasks/task.h`).
-Task code reaches its shard, options, and io manager through these — which is
-why **one process hosts at most one running EloqStore** and why task code must
-run on a shard thread.
+Task code reaches its store, shard, options, and io manager through two
+thread-locals, `eloq_store` and `shard` (`include/tasks/task.h`), which is why
+task code must run on a shard thread. They are bound per thread, so one process
+may run several stores (a data store service runs one per data shard it
+leads): a shard's thread binds its shard and store when its loop starts; under
+`ELOQ_MODULE_ENABLED` one module serves every store in the process, and worker
+`thd_id` runs shard `thd_id` of each store in turn, so
+`EloqStoreModule::Process` binds the shard it runs each time. Only a thread
+doing a shard's work binds them. `Start` and `Stop` bind nothing: they run on
+the caller's thread (in an embedding, often a bthread, which can resume on
+another worker after a wait) and reach the store through `this`. `Start` adds
+the store to the module, and `Stop` removes it once its shards have stopped.
 
 ## Status state machine
 
@@ -66,7 +73,8 @@ the order is load-bearing; later steps assume earlier ones:
    object; a stale term fails the whole `Start`). Or `StandbyService::Start()`.
 8. `Shard::Start()` for each shard — spawns worker threads (normal build) or
    marks shards runnable (module build).
-9. Module build only: register `EloqStoreModule` and notify workers.
+9. Module build only: join the process's `EloqStoreModule` (`AddStore`; the
+   first store to join registers the module with brpc) and notify workers.
 10. Start `ArchiveCrond` if `data_append_mode && num_retained_archives > 0 &&
     archive_interval_secs > 0`.
 11. Start `PrewarmService` if `prewarm_cloud_cache` and mode is Cloud.
@@ -85,7 +93,9 @@ state:
 2. Stop `PrewarmService`, then `ArchiveCrond` (background request *sources*
    stop before request *executors*).
 3. Module build: flag shards `Stopping`, wait until each reports `Stopped`,
-   unregister the module.
+   then leave the module (`RemoveStore`, which waits until no worker is still
+   running one of the store's shards). The module stays registered for the
+   process's other stores.
 4. `Shard::Stop()` for each shard — joins the worker thread; the thread's exit
    path runs `TaskManager::Shutdown()` then `PageManager::Shutdown()` then
    `io_mgr_->Stop()` *on the shard thread* (these structures assume

docs/architecture/10-bindings-and-embedding.md

@@ -57,7 +57,11 @@ only when the embedding runtime (EloqKV's tx service, with bthread) defines
 
 - `EloqStoreModule : eloq::EloqModule` exposes `ExtThdStart/ExtThdEnd`
   (bind/unbind the thread-local `shard`), `Process(thd_id)` (run
-  `Shard::WorkOneRound`), and `HasTask(thd_id)`.
+  `Shard::WorkOneRound`), and `HasTask(thd_id)`. One module serves every
+  store in the process, registered with brpc once: worker `thd_id` runs shard
+  `thd_id` of each running store in turn, binding that shard and its store
+  first. Stores join at `Start` (`AddStore`) and leave after their shards
+  stop (`RemoveStore`, which waits for workers still inside one of them).
 - Shards own no threads; the host's workers drive shard progress. Shard
   count must match the host's worker topology (`thd_id` ↔ `shard_id`).
 - Request sync waits switch from `std::atomic` wait to bthread

include/eloq_store.h

@@ -1050,12 +1050,14 @@ class EloqStore
     std::unique_ptr<PrewarmService> prewarm_service_{nullptr};
     std::unique_ptr<StandbyService> standby_service_{nullptr};
 #ifdef ELOQ_MODULE_ENABLED
-    std::unique_ptr<EloqStoreModule> module_{nullptr};
+    /** @brief Whether the process's EloqStoreModule runs this store. */
+    bool in_module_{false};
 #endif
 
     bool enable_eloqstore_metrics_{false};
     std::atomic<StoreMode> store_mode_{StoreMode::Local};
 
+    friend class EloqStoreModule;
     friend class Shard;
     friend class KvRequest;
     friend class AsyncIoManager;

include/eloqstore_module.h

@@ -2,22 +2,37 @@
 #ifdef ELOQ_MODULE_ENABLED
 #include <bthread/eloq_module.h>
 
-#include <vector>
+#include <array>
+#include <atomic>
+#include <cstddef>
+#include <mutex>
 
 #include "storage/shard.h"
 
 namespace eloqstore
 {
+class EloqStore;
 
+/**
+ * @brief The brpc module that runs the shards of every EloqStore in the
+ * process on brpc workers.
+ *
+ * One module serves all running stores: a process may run several (a data
+ * store service runs one per data shard it leads), and brpc registers few
+ * modules, each registration or removal waiting for every worker. Worker
+ * thd_id runs shard thd_id of each store in turn. A store joins at Start
+ * (AddStore) and leaves once its shards have stopped (RemoveStore), which
+ * waits until no worker is still running one of them.
+ */
 class EloqStoreModule : public eloq::EloqModule
 {
 public:
-    EloqStoreModule() = default;
-    explicit EloqStoreModule(std::vector<std::unique_ptr<Shard>> *shards)
-        : shards_(shards)
-    {
-    }
-    ~EloqStoreModule() = default;
+    /**
+     * @brief The process's module, registered with brpc when the first store
+     * joins.
+     * @return The module.
+     */
+    static EloqStoreModule &Instance();
 
     eloq::ModuleType Type() const override
     {
@@ -29,7 +44,45 @@ class EloqStoreModule : public eloq::EloqModule
     void Process(int thd_id) override;
     bool HasTask(int thd_id) const override;
 
-    std::vector<std::unique_ptr<Shard>> *shards_;
+    /**
+     * @brief Adds a store whose shards the workers start running.
+     * @param store The store, its shards started.
+     */
+    void AddStore(EloqStore *store);
+
+    /**
+     * @brief Removes a store, and waits until no worker is running one of its
+     * shards. Called once the shards have stopped: the workers run them until
+     * then.
+     * @param store The store.
+     */
+    void RemoveStore(EloqStore *store);
+
+    /** @brief How many stores may run in one process at once. */
+    static constexpr size_t kMaxStores = 8;
+
+private:
+    EloqStoreModule() = default;
+
+    /** @brief A store the workers run, and how many are using it now. */
+    struct Slot
+    {
+        std::atomic<EloqStore *> store{nullptr};  ///< nullptr: free.
+        mutable std::atomic<int> users{0};  ///< Workers inside the store now.
+    };
+
+    /**
+     * @brief Calls fn on shard thd_id of each store, until fn returns true.
+     * @param thd_id The worker, and the shard index.
+     * @param fn Called as fn(Shard *); true stops the walk.
+     * @return true if fn returned true.
+     */
+    template <typename Fn>
+    bool ForEachShard(int thd_id, Fn &&fn) const;
+
+    std::array<Slot, kMaxStores> slots_;
+    std::mutex mux_;          ///< Serializes AddStore and RemoveStore.
+    bool registered_{false};  ///< Registered with brpc; guarded by mux_.
 };
 
 }  // namespace eloqstore

include/storage/shard.h

@@ -128,9 +128,14 @@ class Shard
                task_mgr_.NumActive() == 0 && io_mgr_->IsIdle() &&
                !io_mgr_->NeedPrewarm() && delayed_requests_.empty();
     }
+    /**
+     * @brief Binds this shard and its store to the calling thread (shard,
+     * eloq_store), for a brpc worker that runs it.
+     */
     void BindExtThd()
     {
         shard = this;
+        eloq_store = store_;
     }
 #endif
 

include/tasks/batch_write_task.h

@@ -15,7 +15,13 @@ namespace eloqstore
 class BatchWriteTask : public WriteTask
 {
 public:
-    BatchWriteTask();
+    /**
+     * @brief Builds an idle task.
+     * @param opts The store's options. Given rather than read through the
+     * thread's store binding: a shard builds its pooled tasks when it is
+     * constructed, on EloqStore::Start's caller, which binds no store.
+     */
+    explicit BatchWriteTask(const KvOptions *opts);
     TaskType Type() const override
     {
         return TaskType::BatchWrite;

include/tasks/task.h

@@ -26,8 +26,22 @@ class MappingSnapshot;
 class Shard;
 class EloqStore;
 
-inline EloqStore *eloq_store;
+/**
+ * @brief The store whose work this thread is doing.
+ *
+ * A process may run several stores (a data store service runs one per data
+ * shard it leads), so this is bound per thread, as shard is, and only by a
+ * thread doing a shard's work: a shard's own thread, and a brpc worker each
+ * time a module runs one of its shards (EloqStoreModule::Process).
+ *
+ * @note Code that runs on any other thread, such as EloqStore::Start and
+ * EloqStore::Stop, reaches the store explicitly and must not read this: the
+ * caller may be a bthread, which can resume on another pthread after a wait.
+ */
+inline thread_local EloqStore *eloq_store;
+/** @brief The shard whose work this thread is doing (bound with eloq_store). */
 inline thread_local Shard *shard;
+
 KvTask *ThdTask();
 AsyncIoManager *IoMgr();
 const KvOptions *Options();

include/tasks/task_manager.h

@@ -2,8 +2,10 @@
 #pragma once
 
 #include <cassert>
+#include <deque>
 #include <limits>
 #include <memory>
+#include <type_traits>
 #include <vector>
 
 #include "tasks/background_write.h"
@@ -52,16 +54,28 @@ class TaskManager
     class TaskPool
     {
     public:
-        TaskPool(uint32_t size, bool allow_growth = true)
-            : init_pool_size_(size), allow_growth_(allow_growth)
+        /**
+         * @brief Builds the pool's initial tasks.
+         * @param size Number of tasks built now.
+         * @param allow_growth Whether GetTask builds another task when none is
+         * free.
+         * @param opts The store's options, given to the task types built with
+         * them (see MakeTask).
+         */
+        TaskPool(uint32_t size, bool allow_growth, const KvOptions *opts)
+            : opts_(opts), allow_growth_(allow_growth)
         {
-            if (size > 0)
+            for (uint32_t i = 0; i < size; i++)
             {
-                init_pool_ = std::make_unique<T[]>(size);
-                for (uint32_t i = 0; i < size; i++)
+                if constexpr (std::is_constructible_v<T, const KvOptions *>)
                 {
-                    FreeTask(&init_pool_[i]);
+                    init_pool_.emplace_back(opts_);
                 }
+                else
+                {
+                    init_pool_.emplace_back();
+                }
+                FreeTask(&init_pool_.back());
             }
         };
 
@@ -80,7 +94,7 @@ class TaskManager
             {
                 return nullptr;
             }
-            auto &task = ext_pool_.emplace_back(std::make_unique<T>());
+            auto &task = ext_pool_.emplace_back(MakeTask());
             return task.get();
         }
 
@@ -94,9 +108,9 @@ class TaskManager
         template <typename F>
         void ForEachTask(F &&visitor)
         {
-            for (uint32_t i = 0; i < init_pool_size_; ++i)
+            for (T &task : init_pool_)
             {
-                visitor(&init_pool_[i]);
+                visitor(&task);
             }
             for (auto &task : ext_pool_)
             {
@@ -108,14 +122,33 @@ class TaskManager
         {
             free_head_ = nullptr;
             ext_pool_.clear();
-            init_pool_.reset();
+            init_pool_.clear();
         }
 
     private:
-        std::unique_ptr<T[]> init_pool_{nullptr};
+        /**
+         * @brief Builds one task, with the store's options if its type takes
+         * them.
+         * @return The new task.
+         */
+        std::unique_ptr<T> MakeTask() const
+        {
+            if constexpr (std::is_constructible_v<T, const KvOptions *>)
+            {
+                return std::make_unique<T>(opts_);
+            }
+            else
+            {
+                return std::make_unique<T>();
+            }
+        }
+
+        const KvOptions *opts_;  ///< The store's options, for MakeTask.
+        /// The tasks built up front; a deque, so building one never moves
+        /// those already handed out.
+        std::deque<T> init_pool_;
         std::vector<std::unique_ptr<T>> ext_pool_;
         T *free_head_{nullptr};
-        uint32_t init_pool_size_{0};
         bool allow_growth_{true};
     };
 

src/eloq_store.cpp

@@ -583,7 +583,7 @@ void EloqStore::CleanupRuntime(size_t started_shards)
         archive_crond_->Stop();
     }
 #ifdef ELOQ_MODULE_ENABLED
-    if (module_ != nullptr)
+    if (in_module_)
     {
         for (size_t i = 0; i < shard_count; ++i)
         {
@@ -614,7 +614,9 @@ void EloqStore::CleanupRuntime(size_t started_shards)
             }
             bthread_usleep(1000);
         }
-        eloq::unregister_module(module_.get());
+        // Stopped: the workers no longer need to run the shards.
+        EloqStoreModule::Instance().RemoveStore(this);
+        in_module_ = false;
     }
 #endif
 
@@ -649,10 +651,6 @@ void EloqStore::CleanupRuntime(size_t started_shards)
     {
         cloud_service_->Stop();
     }
-    if (eloq_store == this)
-    {
-        eloq_store = nullptr;
-    }
 }
 
 KvError EloqStore::Start(std::string_view branch,
@@ -750,7 +748,6 @@ KvError EloqStore::Start(std::string_view branch,
         standby_service_.reset();
     }
 
-    eloq_store = this;
     // Initialize
     if (!options_.store_path.empty())
     {
@@ -868,8 +865,8 @@ KvError EloqStore::Start(std::string_view branch,
     }
 
 #ifdef ELOQ_MODULE_ENABLED
-    module_ = std::make_unique<EloqStoreModule>(&shards_);
-    eloq::register_module(module_.get());
+    EloqStoreModule::Instance().AddStore(this);
+    in_module_ = true;
     for (auto &shard : shards_)
     {
         eloq::EloqModule::NotifyWorker(static_cast<int>(shard->shard_id_));