2 files changed, 56 insertions, 1 deletions

diff --git a/vdslib/src/vespa/vdslib/state/cluster_state_bundle.cpp b/vdslib/src/vespa/vdslib/state/cluster_state_bundle.cpp
index a2bbba5e52c..cf9c0ba6dff 100644
--- a/vdslib/src/vespa/vdslib/state/cluster_state_bundle.cpp
+++ b/vdslib/src/vespa/vdslib/state/cluster_state_bundle.cpp
@@ -1,7 +1,9 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
+#include <vespa/document/bucket/fixed_bucket_spaces.h>
 #include "cluster_state_bundle.h"
 #include "clusterstate.h"
+#include <iostream>
 
 namespace storage::lib {
 
@@ -10,6 +12,13 @@ ClusterStateBundle::ClusterStateBundle(const ClusterState &baselineClusterState)
 {
 }
 
+ClusterStateBundle::ClusterStateBundle(const ClusterState& baselineClusterState,
+                                       BucketSpaceStateMapping derivedBucketSpaceStates)
+    : _baselineClusterState(std::make_shared<const ClusterState>(baselineClusterState)),
+      _derivedBucketSpaceStates(std::move(derivedBucketSpaceStates))
+{
+}
+
 ClusterStateBundle::~ClusterStateBundle() = default;
 
 const std::shared_ptr<const lib::ClusterState> &
@@ -22,6 +31,7 @@ const std::shared_ptr<const lib::ClusterState> &
 ClusterStateBundle::getDerivedClusterState(document::BucketSpace) const
 {
     // For now, just return the baseline cluster state.
+    // TODO use _derivedBucketSpaceStates
     return _baselineClusterState;
 }
 
@@ -34,7 +44,36 @@ ClusterStateBundle::getVersion() const
 bool
 ClusterStateBundle::operator==(const ClusterStateBundle &rhs) const
 {
-    return *_baselineClusterState == *rhs._baselineClusterState;
+    if (!(*_baselineClusterState == *rhs._baselineClusterState)) {
+        return false;
+    }
+    if (_derivedBucketSpaceStates.size() != rhs._derivedBucketSpaceStates.size()) {
+        return false;
+    }
+    // Can't do a regular operator== comparison since we must check equality
+    // of cluster state _values_, not their _pointers_.
+    for (auto& lhs_ds : _derivedBucketSpaceStates) {
+        auto rhs_iter = rhs._derivedBucketSpaceStates.find(lhs_ds.first);
+        if ((rhs_iter == rhs._derivedBucketSpaceStates.end())
+            || !(*lhs_ds.second == *rhs_iter->second)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+std::ostream& operator<<(std::ostream& os, const ClusterStateBundle& bundle) {
+    os << "ClusterStateBundle('" << *bundle.getBaselineClusterState();
+    if (!bundle.getDerivedClusterStates().empty()) {
+        // Output ordering is undefined for of per-space states.
+        for (auto& ds : bundle.getDerivedClusterStates()) {
+            os << "', ";
+            os << document::FixedBucketSpaces::to_string(ds.first);
+            os << " '" << *ds.second;
+        }
+    }
+    os << "')";
+    return os;
 }
 
 }
diff --git a/vdslib/src/vespa/vdslib/state/cluster_state_bundle.h b/vdslib/src/vespa/vdslib/state/cluster_state_bundle.h
index 77d26092f4e..a8ee0b54713 100644
--- a/vdslib/src/vespa/vdslib/state/cluster_state_bundle.h
+++ b/vdslib/src/vespa/vdslib/state/cluster_state_bundle.h
@@ -3,6 +3,8 @@
 #pragma once
 
 #include <vespa/document/bucket/bucketspace.h>
+#include <unordered_map>
+#include <iosfwd>
 
 namespace storage::lib {
 
@@ -14,14 +16,28 @@ class ClusterState;
  */
 class ClusterStateBundle
 {
+public:
+    using BucketSpaceStateMapping = std::unordered_map<
+        document::BucketSpace,
+        std::shared_ptr<const ClusterState>,
+        document::BucketSpace::hash
+    >;
     std::shared_ptr<const ClusterState> _baselineClusterState;
+    BucketSpaceStateMapping _derivedBucketSpaceStates;
 public:
     explicit ClusterStateBundle(const ClusterState &baselineClusterState);
+    ClusterStateBundle(const ClusterState& baselineClusterState,
+                       BucketSpaceStateMapping derivedBucketSpaceStates);
     ~ClusterStateBundle();
     const std::shared_ptr<const ClusterState> &getBaselineClusterState() const;
     const std::shared_ptr<const ClusterState> &getDerivedClusterState(document::BucketSpace bucketSpace) const;
+    const BucketSpaceStateMapping& getDerivedClusterStates() const noexcept {
+        return _derivedBucketSpaceStates;
+    }
     uint32_t getVersion() const;
     bool operator==(const ClusterStateBundle &rhs) const;
 };
 
+std::ostream& operator<<(std::ostream&, const ClusterStateBundle&);
+
 }