Track extra modified buckets (due to document db removal) per bucket space.

3 files changed, 12 insertions, 6 deletions

diff --git a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp
index 72f2ea8c34d..69a3a902af8 100644
--- a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp
+++ b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.cpp
@@ -585,7 +585,7 @@ PersistenceEngine::getModifiedBuckets(BucketSpace bucketSpace) const
     MBV extraModifiedBuckets;
     {
         LockGuard guard(_lock);
-        extraModifiedBuckets.swap(_extraModifiedBuckets);
+        extraModifiedBuckets.swap(_extraModifiedBuckets[bucketSpace]);
     }
     HandlerSnapshot::UP snap = getHandlerSnapshot(bucketSpace);
     SynchronizedBucketIdListResultHandler resultHandler(snap->size() + extraModifiedBuckets.size());
@@ -699,13 +699,13 @@ PersistenceEngine::propagateSavedClusterState(IPersistenceHandler &handler)
 }
 
 void
-PersistenceEngine::grabExtraModifiedBuckets(IPersistenceHandler &handler)
+PersistenceEngine::grabExtraModifiedBuckets(BucketSpace bucketSpace, IPersistenceHandler &handler)
 {
     BucketIdListResultHandler resultHandler;
     handler.handleListBuckets(resultHandler);
     auto result = std::make_shared<BucketIdListResult>(resultHandler.getResult());
     LockGuard guard(_lock);
-    _extraModifiedBuckets.push_back(result);
+    _extraModifiedBuckets[bucketSpace].push_back(result);
 }
 
 
diff --git a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h
index d9a27f9e460..dec336b05e4 100644
--- a/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h
+++ b/searchcore/src/vespa/searchcore/proton/persistenceengine/persistenceengine.h
@@ -59,8 +59,14 @@ private:
               in_use(false),
               bucket_guards() {}
     };
+    struct BucketSpaceHash {
+        std::size_t operator() (const document::BucketSpace &bucketSpace) const { return bucketSpace.getId(); }
+    };
+
     typedef std::map<IteratorId, IteratorEntry *> Iterators;
     typedef std::vector<std::shared_ptr<BucketIdListResult> > BucketIdListResultV;
+    using ExtraModifiedBuckets = std::unordered_map<BucketSpace, BucketIdListResultV, BucketSpaceHash>;
+
 
     const ssize_t                           _defaultSerializedSize;
     const bool                              _ignoreMaxBytes;
@@ -71,7 +77,7 @@ private:
     IPersistenceEngineOwner                &_owner;
     const IResourceWriteFilter             &_writeFilter;
     ClusterState::SP                        _clusterState;
-    mutable BucketIdListResultV             _extraModifiedBuckets;
+    mutable ExtraModifiedBuckets            _extraModifiedBuckets;
     mutable std::shared_timed_mutex         _rwMutex;
 
     IPersistenceHandler::SP getHandler(document::BucketSpace bucketSpace,
@@ -124,7 +130,7 @@ public:
 
     void destroyIterators();
     void propagateSavedClusterState(IPersistenceHandler &handler);
-    void grabExtraModifiedBuckets(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 34cfaaeddb2..7caed5198bd 100644
--- a/searchcore/src/vespa/searchcore/proton/server/proton.cpp
+++ b/searchcore/src/vespa/searchcore/proton/server/proton.cpp
@@ -623,7 +623,7 @@ Proton::removeDocumentDB(const DocTypeName &docTypeName)
             oldHandler = _persistenceEngine->removeHandler(old->getBucketSpace(), docTypeName);
             if (_initComplete && oldHandler) {
                 // TODO: Fix race with bucket db modifying ops.
-                _persistenceEngine->grabExtraModifiedBuckets(*oldHandler);
+                _persistenceEngine->grabExtraModifiedBuckets(old->getBucketSpace(), *oldHandler);
             }
         }
         _persistenceEngine->destroyIterators();