Add bucketSpace argument to storage::spi::PersistenceProvider::setClusterState() method.

6 files changed, 21 insertions, 19 deletions

diff --git a/searchcore/src/apps/proton/downpersistence.cpp b/searchcore/src/apps/proton/downpersistence.cpp
index 4ee9a0635b0..33ad4bc5024 100644
--- a/searchcore/src/apps/proton/downpersistence.cpp
+++ b/searchcore/src/apps/proton/downpersistence.cpp
@@ -45,7 +45,7 @@ DownPersistence::listBuckets(BucketSpace, PartitionId) const
 }
 
 Result
-DownPersistence:: setClusterState(const ClusterState&)
+DownPersistence:: setClusterState(BucketSpace, const ClusterState&)
 {
     return Result();
 }
diff --git a/searchcore/src/apps/proton/downpersistence.h b/searchcore/src/apps/proton/downpersistence.h
index 9e64b89f065..0a602c4467e 100644
--- a/searchcore/src/apps/proton/downpersistence.h
+++ b/searchcore/src/apps/proton/downpersistence.h
@@ -29,7 +29,7 @@ public:
     Result initialize() override;
     PartitionStateListResult getPartitionStates() const override;
     BucketIdListResult listBuckets(BucketSpace bucketSpace, PartitionId) const override;
-    Result setClusterState(const ClusterState&) override;
+    Result setClusterState(BucketSpace, const ClusterState&) override;
     Result setActiveState(const Bucket&, BucketInfo::ActiveState) override;
     BucketInfoResult getBucketInfo(const Bucket&) const override;
     Result put(const Bucket&, Timestamp, const DocumentSP&, Context&) override;
diff --git a/searchcore/src/tests/proton/persistenceengine/persistenceengine_test.cpp b/searchcore/src/tests/proton/persistenceengine/persistenceengine_test.cpp
index aa7c5a4eda9..4a195514db1 100644
--- a/searchcore/src/tests/proton/persistenceengine/persistenceengine_test.cpp
+++ b/searchcore/src/tests/proton/persistenceengine/persistenceengine_test.cpp
@@ -611,7 +611,7 @@ TEST_F("require that setClusterState() is routed to handlers", SimpleFixture)
 {
     ClusterState state(createClusterState());
 
-    f.engine.setClusterState(state);
+    f.engine.setClusterState(makeBucketSpace(), state);
     EXPECT_EQUAL(&state, f.hset.handler1.lastCalc);
     EXPECT_EQUAL(&state, f.hset.handler2.lastCalc);
 }
diff --git a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp
index 6f2fbabfb1c..c7b01d209ee 100644
--- a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp
+++ b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp
@@ -189,7 +189,7 @@ PersistenceEngine::PersistenceEngine(IPersistenceEngineOwner &owner, const IReso
       _iterators_lock(),
       _owner(owner),
       _writeFilter(writeFilter),
-      _clusterState(),
+      _clusterStates(),
       _extraModifiedBuckets(),
       _rwMutex()
 {
@@ -272,11 +272,11 @@ PersistenceEngine::listBuckets(BucketSpace bucketSpace, PartitionId id) const
 
 
 Result
-PersistenceEngine::setClusterState(const ClusterState &calc)
+PersistenceEngine::setClusterState(BucketSpace bucketSpace, const ClusterState &calc)
 {
     std::shared_lock<std::shared_timed_mutex> rguard(_rwMutex);
-    saveClusterState(calc);
-    HandlerSnapshot::UP snap = getHandlerSnapshot();
+    saveClusterState(bucketSpace, calc);
+    HandlerSnapshot::UP snap = getHandlerSnapshot(bucketSpace);
     GenericResultHandler resultHandler(snap->size());
     for (; snap->handlers().valid(); snap->handlers().next()) {
         IPersistenceHandler *handler = snap->handlers().get();
@@ -630,26 +630,27 @@ PersistenceEngine::destroyIterators()
 
 
 void
-PersistenceEngine::saveClusterState(const ClusterState &calc)
+PersistenceEngine::saveClusterState(BucketSpace bucketSpace, const ClusterState &calc)
 {
     auto clusterState = std::make_shared<ClusterState>(calc);
     {
         std::lock_guard<std::mutex> guard(_lock);
-        clusterState.swap(_clusterState);
+        clusterState.swap(_clusterStates[bucketSpace]);
     }
 }
 
 PersistenceEngine::ClusterState::SP
-PersistenceEngine::savedClusterState() const
+PersistenceEngine::savedClusterState(BucketSpace bucketSpace) const
 {
     std::lock_guard<std::mutex> guard(_lock);
-    return _clusterState;
+    auto itr(_clusterStates.find(bucketSpace));
+    return ((itr != _clusterStates.end()) ? itr->second : ClusterState::SP());
 }
 
 void
-PersistenceEngine::propagateSavedClusterState(IPersistenceHandler &handler)
+PersistenceEngine::propagateSavedClusterState(BucketSpace bucketSpace, IPersistenceHandler &handler)
 {
-    ClusterState::SP clusterState(savedClusterState());
+    ClusterState::SP clusterState(savedClusterState(bucketSpace));
     if (!clusterState)
         return;
     // Propagate saved cluster state.
diff --git a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h
index b3e70e9f433..a9c15e6f2b9 100644
--- a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h
+++ b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h
@@ -10,6 +10,7 @@
 #include <vespa/searchcore/proton/persistenceengine/ipersistencehandler.h>
 #include <mutex>
 #include <shared_mutex>
+#include <unordered_map>
 
 namespace proton {
 
@@ -75,7 +76,7 @@ private:
     mutable std::mutex                      _iterators_lock;
     IPersistenceEngineOwner                &_owner;
     const IResourceWriteFilter             &_writeFilter;
-    ClusterState::SP                        _clusterState;
+    std::unordered_map<BucketSpace, ClusterState::SP, BucketSpace::hash> _clusterStates;
     mutable ExtraModifiedBuckets            _extraModifiedBuckets;
     mutable std::shared_timed_mutex         _rwMutex;
 
@@ -86,8 +87,8 @@ private:
     HandlerSnapshot::UP getHandlerSnapshot(document::BucketSpace bucketSpace,
                                            const document::DocumentId &docId) const;
 
-    void saveClusterState(const ClusterState &calc);
-    ClusterState::SP savedClusterState() const;
+    void saveClusterState(BucketSpace bucketSpace, const ClusterState &calc);
+    ClusterState::SP savedClusterState(BucketSpace bucketSpace) const;
 
 public:
     typedef std::unique_ptr<PersistenceEngine> UP;
@@ -107,7 +108,7 @@ public:
     virtual Result initialize() override;
     virtual PartitionStateListResult getPartitionStates() const override;
     virtual BucketIdListResult listBuckets(BucketSpace bucketSpace, PartitionId) const override;
-    virtual Result setClusterState(const ClusterState& calc) override;
+    virtual Result setClusterState(BucketSpace bucketSpace, const ClusterState& calc) override;
     virtual Result setActiveState(const Bucket& bucket, BucketInfo::ActiveState newState) override;
     virtual BucketInfoResult getBucketInfo(const Bucket&) const override;
     virtual Result put(const Bucket&, Timestamp, const document::Document::SP&, Context&) override;
@@ -128,7 +129,7 @@ public:
     virtual Result maintain(const Bucket&, MaintenanceLevel) override;
 
     void destroyIterators();
-    void propagateSavedClusterState(IPersistenceHandler &handler);
+    void propagateSavedClusterState(BucketSpace bucketSpace, IPersistenceHandler &handler);
     void grabExtraModifiedBuckets(BucketSpace bucketSpace, IPersistenceHandler &handler);
     void populateInitialBucketDB(BucketSpace bucketSpace, IPersistenceHandler &targetHandler);
     std::unique_lock<std::shared_timed_mutex> getWLock() const;
diff --git a/searchcore/src/vespa/searchcore/proton/server/proton.cpp b/searchcore/src/vespa/searchcore/proton/server/proton.cpp
index 2794619273c..6de5410470f 100644
--- a/searchcore/src/vespa/searchcore/proton/server/proton.cpp
+++ b/searchcore/src/vespa/searchcore/proton/server/proton.cpp
@@ -568,7 +568,7 @@ Proton::addDocumentDB(const document::DocumentType &docType,
         std::unique_lock<std::shared_timed_mutex> persistenceWGuard(_persistenceEngine->getWLock());
         auto persistenceHandler = std::make_shared<PersistenceHandlerProxy>(ret);
         if (!_isInitializing) {
-            _persistenceEngine->propagateSavedClusterState(*persistenceHandler);
+            _persistenceEngine->propagateSavedClusterState(bucketSpace, *persistenceHandler);
             _persistenceEngine->populateInitialBucketDB(bucketSpace, *persistenceHandler);
         }
         // TODO: Fix race with new cluster state setting.