Merge pull request #4144 from vespa-engine/toregge/scan-all-bucket-dbs-in-distributor-simple-maintenance-scanner

Scan all bucket dbs in storage::distributor::SimpleMaintenanceScanner.

7 files changed, 87 insertions, 28 deletions

diff --git a/storage/src/tests/distributor/simplemaintenancescannertest.cpp b/storage/src/tests/distributor/simplemaintenancescannertest.cpp
index a46419b71a4..66a2d3efa6c 100644
--- a/storage/src/tests/distributor/simplemaintenancescannertest.cpp
+++ b/storage/src/tests/distributor/simplemaintenancescannertest.cpp
@@ -1,6 +1,9 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
+#include <vespa/document/test/make_bucket_space.h>
 #include <vespa/vdstestlib/cppunit/macros.h>
+#include <vespa/storage/distributor/distributor_bucket_space_repo.h>
+#include <vespa/storage/distributor/distributor_bucket_space.h>
 #include <vespa/storage/distributor/maintenance/simplemaintenancescanner.h>
 #include <vespa/storage/distributor/maintenance/simplebucketprioritydatabase.h>
 #include <vespa/storage/bucketdb/mapbucketdatabase.h>
@@ -11,11 +14,13 @@
 namespace storage::distributor {
 
 using document::BucketId;
+using document::test::makeBucketSpace;
 typedef MaintenancePriority Priority;
 
 class SimpleMaintenanceScannerTest : public CppUnit::TestFixture {
     CPPUNIT_TEST_SUITE(SimpleMaintenanceScannerTest);
     CPPUNIT_TEST(testPrioritizeSingleBucket);
+    CPPUNIT_TEST(testPrioritizeSingleBucketAltBucketSpace);
     CPPUNIT_TEST(testPrioritizeMultipleBuckets);
     CPPUNIT_TEST(testPendingMaintenanceOperationStatistics);
     CPPUNIT_TEST(perNodeMaintenanceStatsAreTracked);
@@ -27,10 +32,11 @@ class SimpleMaintenanceScannerTest : public CppUnit::TestFixture {
     std::string dumpPriorityDbToString(const BucketPriorityDatabase&) const;
 
     std::unique_ptr<MockMaintenancePriorityGenerator> _priorityGenerator;
-    std::unique_ptr<MapBucketDatabase> _bucketDb;
+    std::unique_ptr<DistributorBucketSpaceRepo> _bucketSpaceRepo;
     std::unique_ptr<SimpleBucketPriorityDatabase> _priorityDb;
     std::unique_ptr<SimpleMaintenanceScanner> _scanner;
 
+    void addBucketToDb(document::BucketSpace bucketSpace, int bucketNum);
     void addBucketToDb(int bucketNum);
 
     bool scanEntireDatabase(int expected);
@@ -39,6 +45,7 @@ class SimpleMaintenanceScannerTest : public CppUnit::TestFixture {
 
 public:
     void testPrioritizeSingleBucket();
+    void testPrioritizeSingleBucketAltBucketSpace();
     void testPrioritizeMultipleBuckets();
     void testPendingMaintenanceOperationStatistics();
     void perNodeMaintenanceStatsAreTracked();
@@ -53,16 +60,23 @@ void
 SimpleMaintenanceScannerTest::setUp()
 {
     _priorityGenerator.reset(new MockMaintenancePriorityGenerator());
-    _bucketDb.reset(new MapBucketDatabase());
+    _bucketSpaceRepo = std::make_unique<DistributorBucketSpaceRepo>();
     _priorityDb.reset(new SimpleBucketPriorityDatabase());
-    _scanner.reset(new SimpleMaintenanceScanner(*_priorityDb, *_priorityGenerator, *_bucketDb));
+    _scanner.reset(new SimpleMaintenanceScanner(*_priorityDb, *_priorityGenerator, *_bucketSpaceRepo));
 }
 
 void
-SimpleMaintenanceScannerTest::addBucketToDb(int bucketNum)
+SimpleMaintenanceScannerTest::addBucketToDb(document::BucketSpace bucketSpace, int bucketNum)
 {
     BucketDatabase::Entry entry(BucketId(16, bucketNum), BucketInfo());
-    _bucketDb->update(entry);
+    auto &bucketDb(_bucketSpaceRepo->get(bucketSpace).getBucketDatabase());
+    bucketDb.update(entry);
+}
+
+void
+SimpleMaintenanceScannerTest::addBucketToDb(int bucketNum)
+{
+    addBucketToDb(makeBucketSpace(), bucketNum);
 }
 
 std::string
@@ -80,7 +94,27 @@ SimpleMaintenanceScannerTest::testPrioritizeSingleBucket()
     addBucketToDb(1);
     std::string expected("PrioritizedBucket(Bucket(BucketSpace(0x0000000000000000), BucketId(0x4000000000000001)), pri VERY_HIGH)\n");
 
-    CPPUNIT_ASSERT(!_scanner->scanNext().isDone());
+    auto scanResult = _scanner->scanNext();
+    CPPUNIT_ASSERT(!scanResult.isDone());
+    CPPUNIT_ASSERT_EQUAL(makeBucketSpace().getId(), scanResult.getBucketSpace().getId());
+    CPPUNIT_ASSERT_EQUAL(expected, _priorityDb->toString());
+
+    CPPUNIT_ASSERT(_scanner->scanNext().isDone());
+    CPPUNIT_ASSERT_EQUAL(expected, _priorityDb->toString());
+}
+
+void
+SimpleMaintenanceScannerTest::testPrioritizeSingleBucketAltBucketSpace()
+{
+    document::BucketSpace bucketSpace(4);
+    _bucketSpaceRepo->add(bucketSpace, std::make_unique<DistributorBucketSpace>());
+    _scanner->reset();
+    addBucketToDb(bucketSpace, 1);
+    std::string expected("PrioritizedBucket(Bucket(BucketSpace(0x0000000000000004), BucketId(0x4000000000000001)), pri VERY_HIGH)\n");
+
+    auto scanResult = _scanner->scanNext();
+    CPPUNIT_ASSERT(!scanResult.isDone());
+    CPPUNIT_ASSERT_EQUAL(bucketSpace.getId(), scanResult.getBucketSpace().getId());
     CPPUNIT_ASSERT_EQUAL(expected, _priorityDb->toString());
 
     CPPUNIT_ASSERT(_scanner->scanNext().isDone());
diff --git a/storage/src/vespa/storage/distributor/distributor.cpp b/storage/src/vespa/storage/distributor/distributor.cpp
index c9930b7299c..665f2edd569 100644
--- a/storage/src/vespa/storage/distributor/distributor.cpp
+++ b/storage/src/vespa/storage/distributor/distributor.cpp
@@ -87,7 +87,7 @@ Distributor::Distributor(DistributorComponentRegister& compReg,
       _bucketPriorityDb(new SimpleBucketPriorityDatabase()),
       _scanner(new SimpleMaintenanceScanner(
             *_bucketPriorityDb, _idealStateManager,
-            getDefaultBucketSpace().getBucketDatabase())),
+            *_bucketSpaceRepo)),
       _throttlingStarter(new ThrottlingOperationStarter(
             _maintenanceOperationOwner)),
       _blockingStarter(new BlockingOperationStarter(_pendingMessageTracker,
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 27f2c344e9f..ddf28b3d95d 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.cpp
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.cpp
@@ -16,12 +16,18 @@ namespace distributor {
 DistributorBucketSpaceRepo::DistributorBucketSpaceRepo()
     : _map()
 {
-    _map.emplace(BucketSpace::placeHolder(), std::make_unique<DistributorBucketSpace>());
+    add(BucketSpace::placeHolder(), std::make_unique<DistributorBucketSpace>());
 }
 
 DistributorBucketSpaceRepo::~DistributorBucketSpaceRepo() {
 }
 
+void
+DistributorBucketSpaceRepo::add(document::BucketSpace bucketSpace, std::unique_ptr<DistributorBucketSpace> distributorBucketSpace)
+{
+    _map.emplace(bucketSpace, std::move(distributorBucketSpace));
+}
+
 void DistributorBucketSpaceRepo::setDefaultDistribution(
         std::shared_ptr<lib::Distribution> distr)
 {
@@ -33,7 +39,6 @@ void DistributorBucketSpaceRepo::setDefaultDistribution(
 DistributorBucketSpace &
 DistributorBucketSpaceRepo::get(BucketSpace bucketSpace)
 {
-    assert(bucketSpace == BucketSpace::placeHolder());
     auto itr = _map.find(bucketSpace);
     assert(itr != _map.end());
     return *itr->second;
@@ -42,7 +47,6 @@ DistributorBucketSpaceRepo::get(BucketSpace bucketSpace)
 const DistributorBucketSpace &
 DistributorBucketSpaceRepo::get(BucketSpace bucketSpace) const
 {
-    assert(bucketSpace == BucketSpace::placeHolder());
     auto itr = _map.find(bucketSpace);
     assert(itr != _map.end());
     return *itr->second;
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 a9892a1039f..b0f367e8be5 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.h
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space_repo.h
@@ -37,6 +37,7 @@ public:
     void setDefaultDistribution(std::shared_ptr<lib::Distribution> distr);
     BucketSpaceMap::const_iterator begin() const { return _map.begin(); }
     BucketSpaceMap::const_iterator end() const { return _map.end(); }
+    void add(document::BucketSpace bucketSpace, std::unique_ptr<DistributorBucketSpace> distributorBucketSpace);
 };
 
 }
diff --git a/storage/src/vespa/storage/distributor/maintenance/maintenancescanner.h b/storage/src/vespa/storage/distributor/maintenance/maintenancescanner.h
index 783e8e1e5ba..c1d76b57c7c 100644
--- a/storage/src/vespa/storage/distributor/maintenance/maintenancescanner.h
+++ b/storage/src/vespa/storage/distributor/maintenance/maintenancescanner.h
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 #pragma once
 
+#include <vespa/document/bucket/bucketspace.h>
 #include <vespa/storage/bucketdb/bucketdatabase.h>
 
 namespace storage {
@@ -13,20 +14,22 @@ public:
 
     class ScanResult {
         bool _done;
+        document::BucketSpace _bucketSpace;
         BucketDatabase::Entry _entry;
 
     public:
         bool isDone() const { return _done; }
+        document::BucketSpace getBucketSpace() const { return _bucketSpace; }
         const BucketDatabase::Entry& getEntry() const { return _entry; }
 
         static ScanResult createDone() { return ScanResult(true); }
-        static ScanResult createNotDone(BucketDatabase::Entry entry) {
-            return ScanResult(entry);
+        static ScanResult createNotDone(document::BucketSpace bucketSpace, BucketDatabase::Entry entry) {
+            return ScanResult(bucketSpace, entry);
         }
 
     private:
-        ScanResult(bool done) : _done(done), _entry() {}
-        ScanResult(const BucketDatabase::Entry& e) : _done(false), _entry(e) {}
+        ScanResult(bool done) : _done(done), _bucketSpace(document::BucketSpace::placeHolder()), _entry() {}
+        ScanResult(document::BucketSpace bucketSpace, const BucketDatabase::Entry& e) : _done(false), _bucketSpace(bucketSpace), _entry(e) {}
     };
 
     virtual ScanResult scanNext() = 0;
diff --git a/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.cpp b/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.cpp
index 2bdef7ed320..870dcc25a4b 100644
--- a/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.cpp
+++ b/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.cpp
@@ -1,8 +1,20 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 #include "simplemaintenancescanner.h"
+#include <vespa/storage/distributor/distributor_bucket_space.h>
 
 namespace storage::distributor {
 
+SimpleMaintenanceScanner::SimpleMaintenanceScanner(BucketPriorityDatabase& bucketPriorityDb,
+                                                   const MaintenancePriorityGenerator& priorityGenerator,
+                                                   const DistributorBucketSpaceRepo& bucketSpaceRepo)
+    : _bucketPriorityDb(bucketPriorityDb),
+      _priorityGenerator(priorityGenerator),
+      _bucketSpaceRepo(bucketSpaceRepo),
+      _bucketSpaceItr(_bucketSpaceRepo.begin()),
+      _bucketCursor()
+{
+}
+
 SimpleMaintenanceScanner::~SimpleMaintenanceScanner() {}
 
 SimpleMaintenanceScanner::PendingMaintenanceStats::PendingMaintenanceStats() {}
@@ -14,19 +26,28 @@ SimpleMaintenanceScanner::PendingMaintenanceStats::operator = (const PendingMain
 MaintenanceScanner::ScanResult
 SimpleMaintenanceScanner::scanNext()
 {
-    BucketDatabase::Entry entry(_bucketDb.getNext(_bucketCursor));
-    if (!entry.valid()) {
-        return ScanResult::createDone();
+    for (;;) {
+        if (_bucketSpaceItr == _bucketSpaceRepo.end()) {
+            return ScanResult::createDone();
+        }
+        const auto &bucketDb(_bucketSpaceItr->second->getBucketDatabase());
+        BucketDatabase::Entry entry(bucketDb.getNext(_bucketCursor));
+        if (!entry.valid()) {
+            ++_bucketSpaceItr;
+            _bucketCursor = document::BucketId();
+            continue;
+        }
+        prioritizeBucket(document::Bucket(_bucketSpaceItr->first, entry.getBucketId()));
+        _bucketCursor = entry.getBucketId();
+        return ScanResult::createNotDone(_bucketSpaceItr->first, entry);
     }
-    prioritizeBucket(document::Bucket(document::BucketSpace::placeHolder(), entry.getBucketId()));
-    _bucketCursor = entry.getBucketId();
-    return ScanResult::createNotDone(entry);
 }
 
 void
 SimpleMaintenanceScanner::reset()
 {
     _bucketCursor = document::BucketId();
+    _bucketSpaceItr = _bucketSpaceRepo.begin();
     _pendingMaintenance = PendingMaintenanceStats();
 }
 
diff --git a/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.h b/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.h
index 05de7674d6a..f4ad53957e9 100644
--- a/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.h
+++ b/storage/src/vespa/storage/distributor/maintenance/simplemaintenancescanner.h
@@ -5,7 +5,7 @@
 #include "bucketprioritydatabase.h"
 #include "maintenanceprioritygenerator.h"
 #include "node_maintenance_stats_tracker.h"
-#include <vespa/storage/bucketdb/bucketdatabase.h>
+#include <vespa/storage/distributor/distributor_bucket_space_repo.h>
 
 namespace storage {
 namespace distributor {
@@ -31,18 +31,14 @@ public:
 private:
     BucketPriorityDatabase& _bucketPriorityDb;
     const MaintenancePriorityGenerator& _priorityGenerator;
-    const BucketDatabase& _bucketDb;
+    const DistributorBucketSpaceRepo &_bucketSpaceRepo;
+    DistributorBucketSpaceRepo::BucketSpaceMap::const_iterator _bucketSpaceItr;
     document::BucketId _bucketCursor;
     PendingMaintenanceStats _pendingMaintenance;
 public:
     SimpleMaintenanceScanner(BucketPriorityDatabase& bucketPriorityDb,
                              const MaintenancePriorityGenerator& priorityGenerator,
-                             const BucketDatabase& bucketDb)
-        : _bucketPriorityDb(bucketPriorityDb),
-          _priorityGenerator(priorityGenerator),
-          _bucketDb(bucketDb),
-          _bucketCursor()
-    {}
+                             const DistributorBucketSpaceRepo& bucketSpaceRepo);
     SimpleMaintenanceScanner(const SimpleMaintenanceScanner&) = delete;
     SimpleMaintenanceScanner& operator=(const SimpleMaintenanceScanner&) = delete;
     ~SimpleMaintenanceScanner();