Use context to keep track of extra storagelink instance used to

inject benchmark feed.

3 files changed, 46 insertions, 18 deletions

diff --git a/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.cpp b/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.cpp
index 6ff8f0a2159..6f1acd10fe4 100644
--- a/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.cpp
+++ b/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.cpp
@@ -89,19 +89,28 @@ BmStorageLink::onUp(const std::shared_ptr<storage::api::StorageMessage>& msg)
     return release(msg->getMsgId());
 }
 
+struct StorageApiChainBmFeedHandler::Context {
+    BmStorageLink* bm_link;
+    Context()
+        : bm_link(nullptr)
+    {
+    }
+    ~Context() = default;
+};
+
 class MyStorageChainBuilder : public storage::StorageChainBuilder
 {
     using Parent = storage::StorageChainBuilder;
+    std::shared_ptr<StorageApiChainBmFeedHandler::Context> _context;
 public:
-    MyStorageChainBuilder();
+    MyStorageChainBuilder(std::shared_ptr<StorageApiChainBmFeedHandler::Context> context);
     ~MyStorageChainBuilder() override;
     void add(std::unique_ptr<StorageLink> link) override;
 };
 
-BmStorageLink *bm_link = nullptr;
-
-MyStorageChainBuilder::MyStorageChainBuilder()
-    : storage::StorageChainBuilder()
+MyStorageChainBuilder::MyStorageChainBuilder(std::shared_ptr<StorageApiChainBmFeedHandler::Context> context)
+    : storage::StorageChainBuilder(),
+      _context(std::move(context))
 {
 }
 
@@ -114,13 +123,14 @@ MyStorageChainBuilder::add(std::unique_ptr<StorageLink> link)
     Parent::add(std::move(link));
     if (name == "Communication manager") {
         auto my_link = std::make_unique<BmStorageLink>();
-        bm_link = my_link.get();
+        _context->bm_link = my_link.get();
         Parent::add(std::move(my_link));
     }
 }
 
-StorageApiChainBmFeedHandler::StorageApiChainBmFeedHandler()
-    : IBmFeedHandler()
+StorageApiChainBmFeedHandler::StorageApiChainBmFeedHandler(std::shared_ptr<Context> context)
+    : IBmFeedHandler(),
+      _context(std::move(context))
 {
     auto cmd = make_set_cluster_state_cmd();
     PendingTracker tracker(1);
@@ -134,6 +144,7 @@ void
 StorageApiChainBmFeedHandler::send_msg(std::shared_ptr<storage::api::StorageCommand> cmd, PendingTracker& pending_tracker)
 {
     cmd->setSourceIndex(0);
+    auto bm_link = _context->bm_link;
     bm_link->retain(cmd->getMsgId(), pending_tracker);
     bm_link->sendDown(std::move(cmd));
 }
@@ -159,10 +170,16 @@ StorageApiChainBmFeedHandler::remove(const document::Bucket& bucket, const Docum
     send_msg(std::move(cmd), tracker);
 }
 
+std::shared_ptr<StorageApiChainBmFeedHandler::Context>
+StorageApiChainBmFeedHandler::get_context()
+{
+    return std::make_shared<Context>();
+}
+
 std::unique_ptr<storage::IStorageChainBuilder>
-StorageApiChainBmFeedHandler::get_storage_chain_builder()
+StorageApiChainBmFeedHandler::get_storage_chain_builder(std::shared_ptr<Context> context)
 {
-    return std::make_unique<MyStorageChainBuilder>();
+    return std::make_unique<MyStorageChainBuilder>(std::move(context));
 }
 
 }
diff --git a/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.h b/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.h
index 79c04c98de9..521deddd19e 100644
--- a/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.h
+++ b/searchcore/src/apps/vespa-feed-bm/storage_api_chain_bm_feed_handler.h
@@ -15,15 +15,20 @@ namespace feedbm {
  */
 class StorageApiChainBmFeedHandler : public IBmFeedHandler
 {
+public:
+    struct Context;
+private:
+    std::shared_ptr<Context> _context;
     void send_msg(std::shared_ptr<storage::api::StorageCommand> cmd, PendingTracker& tracker);
 public:
-    StorageApiChainBmFeedHandler();
+    StorageApiChainBmFeedHandler(std::shared_ptr<Context> context);
     ~StorageApiChainBmFeedHandler();
     void put(const document::Bucket& bucket, std::unique_ptr<document::Document> document, uint64_t timestamp, PendingTracker& tracker) override;
     void update(const document::Bucket& bucket, std::unique_ptr<document::DocumentUpdate> document_update, uint64_t timestamp, PendingTracker& tracker) override;
     void remove(const document::Bucket& bucket, const document::DocumentId& document_id,  uint64_t timestamp, PendingTracker& tracker) override;
 
-    static std::unique_ptr<storage::IStorageChainBuilder> get_storage_chain_builder();
+    static std::shared_ptr<Context> get_context();
+    static std::unique_ptr<storage::IStorageChainBuilder> get_storage_chain_builder(std::shared_ptr<Context> context);
 };
 
 }
diff --git a/searchcore/src/apps/vespa-feed-bm/vespa_feed_bm.cpp b/searchcore/src/apps/vespa-feed-bm/vespa_feed_bm.cpp
index 2c143c7ab23..73d741c3fee 100644
--- a/searchcore/src/apps/vespa-feed-bm/vespa_feed_bm.cpp
+++ b/searchcore/src/apps/vespa-feed-bm/vespa_feed_bm.cpp
@@ -282,7 +282,7 @@ class MyServiceLayerProcess : public storage::ServiceLayerProcess {
 public:
     MyServiceLayerProcess(const config::ConfigUri & configUri,
                           PersistenceProvider &provider,
-                          bool use_storage_chain);
+                          std::unique_ptr<storage::IStorageChainBuilder> chain_builder);
     ~MyServiceLayerProcess() override { shutdown(); }
 
     void shutdown() override;
@@ -292,12 +292,12 @@ public:
 
 MyServiceLayerProcess::MyServiceLayerProcess(const config::ConfigUri & configUri,
                                              PersistenceProvider &provider,
-                                             bool use_storage_chain)
+                                             std::unique_ptr<storage::IStorageChainBuilder> chain_builder)
     : ServiceLayerProcess(configUri),
       _provider(provider)
 {
-    if (use_storage_chain) {
-        set_storage_chain_builder(StorageApiChainBmFeedHandler::get_storage_chain_builder());
+    if (chain_builder) {
+        set_storage_chain_builder(std::move(chain_builder));
     }
 }
 
@@ -632,9 +632,15 @@ PersistenceProviderFixture::start_service_layer(bool use_storage_chain)
     _slobrok = std::make_unique<mbus::Slobrok>(_slobrok_port);
     LOG(info, "start service layer");
     config::ConfigUri config_uri("bm-servicelayer", _config_context);
+    std::unique_ptr<storage::IStorageChainBuilder> chain_builder;
+    std::shared_ptr<StorageApiChainBmFeedHandler::Context> context;
+    if (use_storage_chain) {
+        context = StorageApiChainBmFeedHandler::get_context();
+        chain_builder = StorageApiChainBmFeedHandler::get_storage_chain_builder(context);
+    }
     _service_layer = std::make_unique<MyServiceLayerProcess>(config_uri,
                                                              *_persistence_engine,
-                                                             use_storage_chain);
+                                                             std::move(chain_builder));
     _service_layer->setupConfig(100ms);
     _service_layer->createNode();
     _service_layer->getNode().waitUntilInitialized();
@@ -643,7 +649,7 @@ PersistenceProviderFixture::start_service_layer(bool use_storage_chain)
     _rpc_client_shared_rpc_resources = std::make_unique<SharedRpcResources>(client_config_uri, _rpc_client_port, 100);
     _rpc_client_shared_rpc_resources->start_server_and_register_slobrok("bm-rpc-client");
     if (use_storage_chain) {
-        _feed_handler = std::make_unique<StorageApiChainBmFeedHandler>();
+        _feed_handler = std::make_unique<StorageApiChainBmFeedHandler>(std::move(context));
     } else {
         _feed_handler = std::make_unique<StorageApiRpcBmFeedHandler>(*_rpc_client_shared_rpc_resources, _repo);
     }