Add config option for using B-tree bucket DB in distributor

Still disabled by default; this will be swapped later.
Expose read guard generation for easier debugging.
Add some explicit tests for read guard snapshot semantics.

18 files changed, 106 insertions, 30 deletions

diff --git a/storage/src/tests/distributor/btree_bucket_database_test.cpp b/storage/src/tests/distributor/btree_bucket_database_test.cpp
index c253a758f98..43d74ca2fb5 100644
--- a/storage/src/tests/distributor/btree_bucket_database_test.cpp
+++ b/storage/src/tests/distributor/btree_bucket_database_test.cpp
@@ -2,11 +2,56 @@
 
 #include <vespa/storage/bucketdb/btree_bucket_database.h>
 #include <tests/distributor/bucketdatabasetest.h>
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
+
+using namespace ::testing;
 
 namespace storage::distributor {
 
 INSTANTIATE_TEST_CASE_P(BTreeDatabase, BucketDatabaseTest,
                         ::testing::Values(std::make_shared<BTreeBucketDatabase>()));
 
+using document::BucketId;
+
+namespace {
+
+BucketCopy BC(uint32_t node_idx, uint32_t state) {
+    api::BucketInfo info(0x123, state, state);
+    return BucketCopy(0, node_idx, info);
+}
+
+
+BucketInfo BI(uint32_t node_idx, uint32_t state) {
+    BucketInfo bi;
+    bi.addNode(BC(node_idx, state), toVector<uint16_t>(0));
+    return bi;
+}
+
+}
+
+struct BTreeReadGuardTest : Test {
+    BTreeBucketDatabase _db;
+};
+
+TEST_F(BTreeReadGuardTest, guard_does_not_observe_new_entries) {
+    auto guard = _db.acquire_read_guard();
+    _db.update(BucketDatabase::Entry(BucketId(16, 16), BI(1, 1234)));
+    std::vector<BucketDatabase::Entry> entries;
+    guard->find_parents_and_self(BucketId(16, 16), entries);
+    EXPECT_EQ(entries.size(), 0U);
+}
+
+TEST_F(BTreeReadGuardTest, guard_observes_entries_alive_at_acquire_time) {
+    BucketId bucket(16, 16);
+    _db.update(BucketDatabase::Entry(bucket, BI(1, 1234)));
+    auto guard = _db.acquire_read_guard();
+    _db.remove(bucket);
+    std::vector<BucketDatabase::Entry> entries;
+    guard->find_parents_and_self(bucket, entries);
+    ASSERT_EQ(entries.size(), 1U);
+    EXPECT_EQ(entries[0].getBucketInfo(), BI(1, 1234));
+}
+
 }
 
diff --git a/storage/src/tests/distributor/distributortestutil.cpp b/storage/src/tests/distributor/distributortestutil.cpp
index 91af37e0f30..4a9ef147741 100644
--- a/storage/src/tests/distributor/distributortestutil.cpp
+++ b/storage/src/tests/distributor/distributortestutil.cpp
@@ -29,6 +29,7 @@ DistributorTestUtil::createLinks()
             *_threadPool,
             *this,
             true,
+            false, // TODO swap default
             _hostInfo,
             &_messageSender));
     _component.reset(new storage::DistributorComponent(_node->getComponentRegister(), "distrtestutil"));
diff --git a/storage/src/tests/distributor/simplemaintenancescannertest.cpp b/storage/src/tests/distributor/simplemaintenancescannertest.cpp
index b21a10c319e..abf061dd990 100644
--- a/storage/src/tests/distributor/simplemaintenancescannertest.cpp
+++ b/storage/src/tests/distributor/simplemaintenancescannertest.cpp
@@ -36,7 +36,7 @@ void
 SimpleMaintenanceScannerTest::SetUp()
 {
     _priorityGenerator = std::make_unique<MockMaintenancePriorityGenerator>();
-    _bucketSpaceRepo = std::make_unique<DistributorBucketSpaceRepo>();
+    _bucketSpaceRepo = std::make_unique<DistributorBucketSpaceRepo>(false);
     _priorityDb = std::make_unique<SimpleBucketPriorityDatabase>();
     _scanner = std::make_unique<SimpleMaintenanceScanner>(*_priorityDb, *_priorityGenerator, *_bucketSpaceRepo);
 }
@@ -79,7 +79,7 @@ TEST_F(SimpleMaintenanceScannerTest, prioritize_single_bucket) {
 
 TEST_F(SimpleMaintenanceScannerTest, prioritize_single_bucket_alt_bucket_space) {
     document::BucketSpace bucketSpace(4);
-    _bucketSpaceRepo->add(bucketSpace, std::make_unique<DistributorBucketSpace>());
+    _bucketSpaceRepo->add(bucketSpace, std::make_unique<DistributorBucketSpace>(false));
     _scanner->reset();
     addBucketToDb(bucketSpace, 1);
     std::string expected("PrioritizedBucket(Bucket(BucketSpace(0x0000000000000004), BucketId(0x4000000000000001)), pri VERY_HIGH)\n");
diff --git a/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp b/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp
index a901bbdd96b..c3ade3c2877 100644
--- a/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp
+++ b/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp
@@ -537,4 +537,8 @@ void BTreeBucketDatabase::ReadGuardImpl::find_parents_and_self(const document::B
     _db->find_parents_and_self_internal(_frozen_view, bucket, entries);
 }
 
+uint64_t BTreeBucketDatabase::ReadGuardImpl::generation() const noexcept {
+    return _guard.getGeneration();
+}
+
 }
diff --git a/storage/src/vespa/storage/bucketdb/btree_bucket_database.h b/storage/src/vespa/storage/bucketdb/btree_bucket_database.h
index 5c69b57956b..1f2b25814a8 100644
--- a/storage/src/vespa/storage/bucketdb/btree_bucket_database.h
+++ b/storage/src/vespa/storage/bucketdb/btree_bucket_database.h
@@ -81,16 +81,16 @@ private:
 
         void find_parents_and_self(const document::BucketId& bucket,
                                    std::vector<Entry>& entries) const override;
+        uint64_t generation() const noexcept override;
     };
 
     friend class ReadGuardImpl;
-
+    friend struct BTreeBuilderMerger;
+    friend struct BTreeTrailingInserter;
+public:
     std::unique_ptr<ReadGuard> acquire_read_guard() const override {
         return std::make_unique<ReadGuardImpl>(*this);
     }
-
-    friend struct BTreeBuilderMerger;
-    friend struct BTreeTrailingInserter;
 };
 
 }
diff --git a/storage/src/vespa/storage/bucketdb/bucketdatabase.h b/storage/src/vespa/storage/bucketdb/bucketdatabase.h
index 0659d02cc15..46aaaa997d9 100644
--- a/storage/src/vespa/storage/bucketdb/bucketdatabase.h
+++ b/storage/src/vespa/storage/bucketdb/bucketdatabase.h
@@ -244,6 +244,7 @@ public:
 
         virtual void find_parents_and_self(const document::BucketId& bucket,
                                            std::vector<Entry>& entries) const = 0;
+        virtual uint64_t generation() const noexcept = 0;
     };
 
     virtual std::unique_ptr<ReadGuard> acquire_read_guard() const {
diff --git a/storage/src/vespa/storage/bucketdb/mapbucketdatabase.h b/storage/src/vespa/storage/bucketdb/mapbucketdatabase.h
index b295be588a6..9fe5e2d7740 100644
--- a/storage/src/vespa/storage/bucketdb/mapbucketdatabase.h
+++ b/storage/src/vespa/storage/bucketdb/mapbucketdatabase.h
@@ -73,6 +73,7 @@ private:
 
         void find_parents_and_self(const document::BucketId& bucket,
                                    std::vector<Entry>& entries) const override;
+        uint64_t generation() const noexcept override { return 0; }
     };
 
     uint32_t allocate();
diff --git a/storage/src/vespa/storage/config/stor-distributormanager.def b/storage/src/vespa/storage/config/stor-distributormanager.def
index 182aa2008c5..e789f03bb14 100644
--- a/storage/src/vespa/storage/config/stor-distributormanager.def
+++ b/storage/src/vespa/storage/config/stor-distributormanager.def
@@ -199,3 +199,9 @@ allow_stale_reads_during_cluster_state_transitions bool default=false
 ## Setting any of these values only makes sense for testing!
 simulated_db_pruning_latency_msec int default=0
 simulated_db_merging_latency_msec int default=0
+
+## Whether to use a B-tree data structure for the distributor bucket database instead
+## of the legacy database. Setting this option may trigger alternate code paths for
+## read only operations, as the B-tree database is thread safe for concurrent reads.
+use_btree_database bool default=false restart
+
diff --git a/storage/src/vespa/storage/distributor/distributor.cpp b/storage/src/vespa/storage/distributor/distributor.cpp
index 7cb7d687446..69b64ac8dc1 100644
--- a/storage/src/vespa/storage/distributor/distributor.cpp
+++ b/storage/src/vespa/storage/distributor/distributor.cpp
@@ -58,6 +58,7 @@ Distributor::Distributor(DistributorComponentRegister& compReg,
                          framework::TickingThreadPool& threadPool,
                          DoneInitializeHandler& doneInitHandler,
                          bool manageActiveBucketCopies,
+                         bool use_btree_database,
                          HostInfo& hostInfoReporterRegistrar,
                          ChainedMessageSender* messageSender)
     : StorageLink("distributor"),
@@ -66,8 +67,8 @@ Distributor::Distributor(DistributorComponentRegister& compReg,
       _clusterStateBundle(lib::ClusterState()),
       _compReg(compReg),
       _component(compReg, "distributor"),
-      _bucketSpaceRepo(std::make_unique<DistributorBucketSpaceRepo>()),
-      _readOnlyBucketSpaceRepo(std::make_unique<DistributorBucketSpaceRepo>()),
+      _bucketSpaceRepo(std::make_unique<DistributorBucketSpaceRepo>(use_btree_database)),
+      _readOnlyBucketSpaceRepo(std::make_unique<DistributorBucketSpaceRepo>(use_btree_database)),
       _metrics(new DistributorMetricSet(_component.getLoadTypes()->getMetricLoadTypes())),
       _operationOwner(*this, _component.getClock()),
       _maintenanceOperationOwner(*this, _component.getClock()),
@@ -103,6 +104,10 @@ Distributor::Distributor(DistributorComponentRegister& compReg,
                 std::chrono::seconds(0))), // Set by config later
       _must_send_updated_host_info(false)
 {
+    if (use_btree_database) {
+        LOG(info, "Using new B-tree bucket database implementation instead of legacy implementation"); // TODO remove this once default is swapped
+    }
+
     _component.registerMetric(*_metrics);
     _component.registerMetricUpdateHook(_metricUpdateHook,
                                         framework::SecondTime(0));
diff --git a/storage/src/vespa/storage/distributor/distributor.h b/storage/src/vespa/storage/distributor/distributor.h
index 424ac0e7a78..638704adf24 100644
--- a/storage/src/vespa/storage/distributor/distributor.h
+++ b/storage/src/vespa/storage/distributor/distributor.h
@@ -51,6 +51,7 @@ public:
                 framework::TickingThreadPool&,
                 DoneInitializeHandler&,
                 bool manageActiveBucketCopies,
+                bool use_btree_database,
                 HostInfo& hostInfoReporterRegistrar,
                 ChainedMessageSender* = nullptr);
 
diff --git a/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp b/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp
index f013ce43048..3f7dbda62d9 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp
@@ -10,15 +10,17 @@ namespace storage::distributor {
 
 namespace {
 
-std::unique_ptr<BucketDatabase> make_default_db_impl() {
-    //return std::make_unique<BTreeBucketDatabase>();
+std::unique_ptr<BucketDatabase> make_default_db_impl(bool use_btree_db) {
+    if (use_btree_db) {
+        return std::make_unique<BTreeBucketDatabase>();
+    }
     return std::make_unique<MapBucketDatabase>();
 }
 
 }
 
-DistributorBucketSpace::DistributorBucketSpace()
-    : _bucketDatabase(make_default_db_impl()),
+DistributorBucketSpace::DistributorBucketSpace(bool use_btree_db)
+    : _bucketDatabase(make_default_db_impl(use_btree_db)),
       _clusterState(),
       _distribution()
 {
diff --git a/storage/src/vespa/storage/distributor/distributor_bucket_space.h b/storage/src/vespa/storage/distributor/distributor_bucket_space.h
index effb0dc3e17..26a0ee9098c 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space.h
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space.h
@@ -30,7 +30,7 @@ class DistributorBucketSpace {
     std::shared_ptr<const lib::ClusterState> _clusterState;
     std::shared_ptr<const lib::Distribution> _distribution;
 public:
-    DistributorBucketSpace();
+    explicit DistributorBucketSpace(bool use_btree_db);
     ~DistributorBucketSpace();
 
     DistributorBucketSpace(const DistributorBucketSpace&) = delete;
diff --git a/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.cpp b/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.cpp
index 744c54676ae..54287d32666 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.cpp
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.cpp
@@ -13,11 +13,11 @@ using document::BucketSpace;
 
 namespace storage::distributor {
 
-DistributorBucketSpaceRepo::DistributorBucketSpaceRepo()
+DistributorBucketSpaceRepo::DistributorBucketSpaceRepo(bool use_btree_db)
     : _map()
 {
-    add(document::FixedBucketSpaces::default_space(), std::make_unique<DistributorBucketSpace>());
-    add(document::FixedBucketSpaces::global_space(), std::make_unique<DistributorBucketSpace>());
+    add(document::FixedBucketSpaces::default_space(), std::make_unique<DistributorBucketSpace>(use_btree_db));
+    add(document::FixedBucketSpaces::global_space(), std::make_unique<DistributorBucketSpace>(use_btree_db));
 }
 
 DistributorBucketSpaceRepo::~DistributorBucketSpaceRepo() = default;
diff --git a/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.h b/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.h
index ee36842969a..bc42ae8bb3a 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.h
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.h
@@ -19,7 +19,7 @@ private:
     BucketSpaceMap _map;
 
 public:
-    DistributorBucketSpaceRepo();
+    explicit DistributorBucketSpaceRepo(bool use_btree_db); // TODO remove param once B-tree is default
     ~DistributorBucketSpaceRepo();
 
     DistributorBucketSpaceRepo(const DistributorBucketSpaceRepo&&) = delete;
diff --git a/storage/src/vespa/storage/storageserver/distributornode.cpp b/storage/src/vespa/storage/storageserver/distributornode.cpp
index 1cd1477e769..20b5dd641f7 100644
--- a/storage/src/vespa/storage/storageserver/distributornode.cpp
+++ b/storage/src/vespa/storage/storageserver/distributornode.cpp
@@ -19,6 +19,7 @@ DistributorNode::DistributorNode(
         DistributorNodeContext& context,
         ApplicationGenerationFetcher& generationFetcher,
         NeedActiveState activeState,
+        bool use_btree_database,
         StorageLink::UP communicationManager)
     : StorageNode(configUri, context, generationFetcher,
             std::unique_ptr<HostInfo>(new HostInfo()),
@@ -29,6 +30,7 @@ DistributorNode::DistributorNode(
       _lastUniqueTimestampRequested(0),
       _uniqueTimestampCounter(0),
       _manageActiveBucketCopies(activeState == NEED_ACTIVE_BUCKET_STATES_SET),
+      _use_btree_database(use_btree_database),
       _retrievedCommunicationManager(std::move(communicationManager))
 {
     try{
@@ -108,6 +110,7 @@ DistributorNode::createChain()
             new storage::distributor::Distributor(
                 dcr, *_threadPool, getDoneInitializeHandler(),
                 _manageActiveBucketCopies,
+                _use_btree_database,
                 stateManager->getHostInfo())));
 
     chain->push_back(StorageLink::UP(stateManager.release()));
diff --git a/storage/src/vespa/storage/storageserver/distributornode.h b/storage/src/vespa/storage/storageserver/distributornode.h
index 34f2dbb42a7..3a81e31fc56 100644
--- a/storage/src/vespa/storage/storageserver/distributornode.h
+++ b/storage/src/vespa/storage/storageserver/distributornode.h
@@ -24,6 +24,7 @@ class DistributorNode
     uint64_t _lastUniqueTimestampRequested;
     uint32_t _uniqueTimestampCounter;
     bool _manageActiveBucketCopies;
+    bool _use_btree_database;
     std::unique_ptr<StorageLink> _retrievedCommunicationManager;
 
 public:
@@ -38,8 +39,9 @@ public:
                     DistributorNodeContext&,
                     ApplicationGenerationFetcher& generationFetcher,
                     NeedActiveState,
+                    bool use_btree_database,
                     std::unique_ptr<StorageLink> communicationManager);
-    ~DistributorNode();
+    ~DistributorNode() override;
 
     const lib::NodeType& getNodeType() const override { return lib::NodeType::DISTRIBUTOR; }
     ResumeGuard pause() override;
diff --git a/storageserver/src/vespa/storageserver/app/distributorprocess.cpp b/storageserver/src/vespa/storageserver/app/distributorprocess.cpp
index 57fdcdeb248..ff4b2e98cca 100644
--- a/storageserver/src/vespa/storageserver/app/distributorprocess.cpp
+++ b/storageserver/src/vespa/storageserver/app/distributorprocess.cpp
@@ -11,7 +11,8 @@ namespace storage {
 
 DistributorProcess::DistributorProcess(const config::ConfigUri & configUri)
     : Process(configUri),
-      _activeFlag(DistributorNode::NO_NEED_FOR_ACTIVE_STATES)
+      _activeFlag(DistributorNode::NO_NEED_FOR_ACTIVE_STATES),
+      _use_btree_database(false)
 {
 }
 
@@ -29,19 +30,22 @@ DistributorProcess::shutdown()
 void
 DistributorProcess::setupConfig(uint64_t subscribeTimeout)
 {
-    std::unique_ptr<vespa::config::content::core::StorServerConfig> config =
-        config::ConfigGetter<vespa::config::content::core::StorServerConfig>::getConfig(_configUri.getConfigId(), _configUri.getContext(), subscribeTimeout);
-    if (config->persistenceProvider.type
-        != vespa::config::content::core::StorServerConfig::PersistenceProvider::Type::STORAGE)
-    {
+    using vespa::config::content::core::StorServerConfig;
+    using vespa::config::content::core::StorDistributormanagerConfig;
+    using vespa::config::content::core::StorVisitordispatcherConfig;
+
+    auto stor_config = config::ConfigGetter<StorServerConfig>::getConfig(
+            _configUri.getConfigId(), _configUri.getContext(), subscribeTimeout);
+    if (stor_config->persistenceProvider.type != StorServerConfig::PersistenceProvider::Type::STORAGE) {
         _activeFlag = DistributorNode::NEED_ACTIVE_BUCKET_STATES_SET;
     }
+    auto dist_config = config::ConfigGetter<StorDistributormanagerConfig>::getConfig(
+            _configUri.getConfigId(), _configUri.getContext(), subscribeTimeout);
+    _use_btree_database = dist_config->useBtreeDatabase;
     _distributorConfigHandler
-            = _configSubscriber.subscribe<vespa::config::content::core::StorDistributormanagerConfig>(
-                    _configUri.getConfigId(), subscribeTimeout);
+            = _configSubscriber.subscribe<StorDistributormanagerConfig>(_configUri.getConfigId(), subscribeTimeout);
     _visitDispatcherConfigHandler
-            = _configSubscriber.subscribe<vespa::config::content::core::StorVisitordispatcherConfig>(
-                    _configUri.getConfigId(), subscribeTimeout);
+            = _configSubscriber.subscribe<StorVisitordispatcherConfig>(_configUri.getConfigId(), subscribeTimeout);
     Process::setupConfig(subscribeTimeout);
 }
 
@@ -75,7 +79,7 @@ DistributorProcess::configUpdated()
 void
 DistributorProcess::createNode()
 {
-    _node.reset(new DistributorNode(_configUri, _context, *this, _activeFlag, StorageLink::UP()));
+    _node.reset(new DistributorNode(_configUri, _context, *this, _activeFlag, _use_btree_database, StorageLink::UP()));
     _node->handleConfigChange(*_distributorConfigHandler->getConfig());
     _node->handleConfigChange(*_visitDispatcherConfigHandler->getConfig());
 }
diff --git a/storageserver/src/vespa/storageserver/app/distributorprocess.h b/storageserver/src/vespa/storageserver/app/distributorprocess.h
index 0e8d0e2d599..57193f77e42 100644
--- a/storageserver/src/vespa/storageserver/app/distributorprocess.h
+++ b/storageserver/src/vespa/storageserver/app/distributorprocess.h
@@ -15,6 +15,7 @@ namespace storage {
 class DistributorProcess : public Process {
     DistributorNodeContext _context;
     DistributorNode::NeedActiveState _activeFlag;
+    bool _use_btree_database;
     DistributorNode::UP _node;
     config::ConfigHandle<vespa::config::content::core::StorDistributormanagerConfig>::UP
             _distributorConfigHandler;
@@ -23,7 +24,7 @@ class DistributorProcess : public Process {
 
 public:
     DistributorProcess(const config::ConfigUri & configUri);
-    ~DistributorProcess();
+    ~DistributorProcess() override;
 
     void shutdown() override;
     void setupConfig(uint64_t subscribeTimeout) override;