From 244f9279f5b07932aeac7aedf8b0460174c990c5 Mon Sep 17 00:00:00 2001
From: chrisxu333 <chrisxu333@gmail.com>
Date: Sun, 14 Jan 2024 17:29:23 -0700
Subject: [PATCH] persist cluster-enabled status in RocksDB

---
 src/server/server.cc                          | 13 +++++++++++
 src/storage/storage.h                         |  2 ++
 .../integration/cluster/cluster_test.go       |  8 +++++++
 tests/gocase/util/server.go                   | 22 +++++++++++++++++++
 4 files changed, 45 insertions(+)

diff --git a/src/server/server.cc b/src/server/server.cc
index efe721b27ba..dde4e807f1f 100644
--- a/src/server/server.cc
+++ b/src/server/server.cc
@@ -152,6 +152,19 @@ Status Server::Start() {
     }
   }
 
+  std::string value;
+  auto cf = storage->GetCFHandle(engine::kPropagateColumnFamilyName);
+  rocksdb::Status check_cluster_enabled = storage->Get(rocksdb::ReadOptions(), cf, rocksdb::Slice(engine::kClusterEnabledKey), &value);
+  if(check_cluster_enabled.IsNotFound()) {
+    s = storage->WriteToPropagateCF(engine::kClusterEnabledKey, std::to_string(config_->cluster_enabled));
+    if(!s.IsOK()) return s;
+  } else if(check_cluster_enabled.ok()) {
+    // use the one that's previously persisted.
+    config_->cluster_enabled = std::stoi(value);
+  } else {
+    return {Status::NotOK, "failed reading from storage"};
+  }
+
   if (config_->cluster_enabled) {
     if (config_->persist_cluster_nodes_enabled) {
       auto s = cluster->LoadClusterNodes(config_->NodesFilePath());
diff --git a/src/storage/storage.h b/src/storage/storage.h
index 7e37d82f7ea..1c8f08976d1 100644
--- a/src/storage/storage.h
+++ b/src/storage/storage.h
@@ -80,6 +80,8 @@ constexpr const char *kLuaFuncSHAPrefix = "lua_f_";
 constexpr const char *kLuaFuncLibPrefix = "lua_func_lib_";
 constexpr const char *kLuaLibCodePrefix = "lua_lib_code_";
 
+const std::string kClusterEnabledKey = "cluster_enabled";
+
 struct CompressionOption {
   rocksdb::CompressionType type;
   const std::string name;
diff --git a/tests/gocase/integration/cluster/cluster_test.go b/tests/gocase/integration/cluster/cluster_test.go
index 8bc42fdafdb..f76c537f745 100644
--- a/tests/gocase/integration/cluster/cluster_test.go
+++ b/tests/gocase/integration/cluster/cluster_test.go
@@ -93,6 +93,14 @@ func TestClusterNodes(t *testing.T) {
 		require.EqualValues(t, []redis.ClusterNode{{ID: nodeID, Addr: srv.HostPort()}}, slots[0].Nodes)
 	})
 
+	t.Run("enable/disable cluster-enabled option", func(t *testing.T) {
+		// force change cluster-enabled status in kvrocks.conf file
+		srv.ForceChangeClusterMode(false)
+		defer srv.ForceChangeClusterMode(true)
+		srv.Restart()
+		require.NoError(t, rdb.Do(ctx, "clusterx", "version").Err())
+	})
+
 	t.Run("enable/disable the persist cluster nodes", func(t *testing.T) {
 		require.NoError(t, rdb.ConfigSet(ctx, "persist-cluster-nodes-enabled", "yes").Err())
 		srv.Restart()
diff --git a/tests/gocase/util/server.go b/tests/gocase/util/server.go
index a3f4314e342..4bd28a896de 100644
--- a/tests/gocase/util/server.go
+++ b/tests/gocase/util/server.go
@@ -28,6 +28,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"regexp"
+	"strings"
 	"sync"
 	"syscall"
 	"testing"
@@ -134,6 +135,27 @@ func (s *KvrocksServer) close(keepDir bool) {
 	s.clean(keepDir)
 }
 
+func (s *KvrocksServer) ForceChangeClusterMode(enable bool) {
+	dir := s.configs["dir"]
+	f, err := os.OpenFile(filepath.Join(dir, "kvrocks.conf"), os.O_RDWR, 0666)
+	require.NoError(s.t, err)
+	defer func() { require.NoError(s.t, f.Close()) }()
+
+	// change the line containing cluster-enabled to no
+	data, err := os.ReadFile(filepath.Join(dir, "kvrocks.conf"))
+	require.NoError(s.t, err)
+
+	content := string(data)
+	var new_content string
+	if !enable {
+		new_content = strings.ReplaceAll(content, "cluster-enabled yes", "cluster-enabled no")
+	} else {
+		new_content = strings.ReplaceAll(content, "cluster-enabled no", "cluster-enabled yes")
+	}
+	err = os.WriteFile(filepath.Join(dir, "kvrocks.conf"), []byte(new_content), 0666)
+	require.NoError(s.t, err)
+}
+
 func (s *KvrocksServer) Restart() {
 	s.close(true)