Add storage::lib::ClusterStateBundle to storage::api::SetSystemStateCommand.

2 files changed, 19 insertions, 8 deletions

diff --git a/storageapi/src/vespa/storageapi/message/state.cpp b/storageapi/src/vespa/storageapi/message/state.cpp
index b128e8f6485..efa9a45764f 100644
--- a/storageapi/src/vespa/storageapi/message/state.cpp
+++ b/storageapi/src/vespa/storageapi/message/state.cpp
@@ -2,6 +2,7 @@
 
 #include "state.h"
 #include <vespa/storageapi/messageapi/storagemessage.h>
+#include <vespa/vdslib/state/clusterstate.h>
 #include <ostream>
 
 namespace storage {
@@ -61,6 +62,12 @@ GetNodeStateReply::print(std::ostream& out, bool verbose,
     }
 }
 
+SetSystemStateCommand::SetSystemStateCommand(const lib::ClusterStateBundle& state)
+    : StorageCommand(MessageType::SETSYSTEMSTATE),
+      _state(state)
+{
+}
+
 SetSystemStateCommand::SetSystemStateCommand(const lib::ClusterState& state)
     : StorageCommand(MessageType::SETSYSTEMSTATE),
       _state(state)
@@ -71,7 +78,7 @@ void
 SetSystemStateCommand::print(std::ostream& out, bool verbose,
                              const std::string& indent) const
 {
-    out << "SetSystemStateCommand(" << _state << ")";
+    out << "SetSystemStateCommand(" << *_state.getBaselineClusterState() << ")";
     if (verbose) {
         out << " : ";
         StorageCommand::print(out, verbose, indent);
@@ -80,7 +87,7 @@ SetSystemStateCommand::print(std::ostream& out, bool verbose,
 
 SetSystemStateReply::SetSystemStateReply(const SetSystemStateCommand& cmd)
     : StorageReply(cmd),
-      _state(cmd.getSystemState())
+      _state(cmd.getClusterStateBundle())
 {
 }
 
diff --git a/storageapi/src/vespa/storageapi/message/state.h b/storageapi/src/vespa/storageapi/message/state.h
index 746d92fce6b..4e5ad92b259 100644
--- a/storageapi/src/vespa/storageapi/message/state.h
+++ b/storageapi/src/vespa/storageapi/message/state.h
@@ -4,7 +4,8 @@
 
 #include <vespa/storageapi/messageapi/storagecommand.h>
 #include <vespa/storageapi/messageapi/storagereply.h>
-#include <vespa/vdslib/state/clusterstate.h>
+#include <vespa/vdslib/state/nodestate.h>
+#include <vespa/vdslib/state/cluster_state_bundle.h>
 
 namespace storage::api {
 
@@ -60,11 +61,13 @@ public:
  *  put/get/remove etx)
  */
 class SetSystemStateCommand : public StorageCommand {
-    lib::ClusterState _state;
+    lib::ClusterStateBundle _state;
 
 public:
-    explicit SetSystemStateCommand(const lib::ClusterState&);
-    const lib::ClusterState& getSystemState() const { return _state; }
+    explicit SetSystemStateCommand(const lib::ClusterStateBundle &state);
+    explicit SetSystemStateCommand(const lib::ClusterState &state);
+    const lib::ClusterState& getSystemState() const { return *_state.getBaselineClusterState(); }
+    const lib::ClusterStateBundle& getClusterStateBundle() const { return _state; }
     void print(std::ostream& out, bool verbose, const std::string& indent) const override;
 
     DECLARE_STORAGECOMMAND(SetSystemStateCommand, onSetSystemState)
@@ -77,13 +80,14 @@ public:
  * @brief Reply received after a SetSystemStateCommand.
  */
 class SetSystemStateReply : public StorageReply {
-    lib::ClusterState _state;
+    lib::ClusterStateBundle _state;
 
 public:
     explicit SetSystemStateReply(const SetSystemStateCommand& cmd);
 
     // Not serialized. Available locally
-    const lib::ClusterState& getSystemState() const { return _state; }
+    const lib::ClusterState& getSystemState() const { return *_state.getBaselineClusterState(); }
+    const lib::ClusterStateBundle& getClusterStateBundle() const { return _state; }
     void print(std::ostream& out, bool verbose, const std::string& indent) const override;
 
     DECLARE_STORAGEREPLY(SetSystemStateReply, onSetSystemStateReply)