From 15d289c33d95d573487ac2b4fe39382156b19d60 Mon Sep 17 00:00:00 2001
From: Kelvin Clement Mwinuka <kelvinmwinuka@hotmail.co.uk>
Date: Thu, 1 Feb 2024 04:42:44 +0800
Subject: [PATCH] Added config flags to choose between aof restore and snapshot
 restore in standalone mode.

---
 src/server/server.go | 14 ++++++++++----
 src/utils/config.go  |  6 ++++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/server/server.go b/src/server/server.go
index fd9b364..5f25172 100644
--- a/src/server/server.go
+++ b/src/server/server.go
@@ -256,8 +256,11 @@ func (server *Server) Start(ctx context.Context) {
 			StartRewriteAOF:  server.StartRewriteAOF,
 			FinishRewriteAOF: server.FinishRewriteAOF,
 		})
-		if err := server.AOFEngine.Restore(ctx); err != nil {
-			log.Println(err)
+		if conf.RestoreAOF && !conf.RestoreSnapshot {
+			err := server.AOFEngine.Restore(ctx)
+			if err != nil {
+				log.Println(err)
+			}
 		}
 		server.AOFEngine.Start(ctx)
 		// Initialize and start standalone snapshot engine
@@ -272,8 +275,11 @@ func (server *Server) Start(ctx context.Context) {
 			KeyUnlock:                     server.KeyUnlock,
 			SetValue:                      server.SetValue,
 		})
-		if err := server.SnapshotEngine.Restore(ctx); err != nil {
-			log.Println(err)
+		if conf.RestoreSnapshot {
+			err := server.SnapshotEngine.Restore(ctx)
+			if err != nil {
+				log.Println(err)
+			}
 		}
 		server.SnapshotEngine.Start(ctx)
 	}
diff --git a/src/utils/config.go b/src/utils/config.go
index b20e3a2..1eaca45 100644
--- a/src/utils/config.go
+++ b/src/utils/config.go
@@ -30,6 +30,8 @@ type Config struct {
 	Password           string `json:"password" yaml:"password"`
 	SnapShotThreshold  uint64 `json:"snapshotThreshold" yaml:"snapshotThreshold"`
 	SnapshotInterval   uint   `json:"snapshotInterval" yaml:"snapshotInterval"`
+	RestoreSnapshot    bool   `json:"restoreSnapshot" yaml:"restoreSnapshot"`
+	RestoreAOF         bool   `json:"restoreAOF" yaml:"restoreAOF"`
 }
 
 func GetConfig() (Config, error) {
@@ -49,6 +51,8 @@ func GetConfig() (Config, error) {
 	aclConfig := flag.String("aclConfig", "", "ACL config file path.")
 	snapshotThreshold := flag.Uint64("snapshotThreshold", 1000, "The number of entries that trigger a snapshot. Default is 1000.")
 	snapshotInterval := flag.Uint("snapshotInterval", 60*5, "The time interval between snapshots (in seconds). Default is 5 minutes.")
+	restoreSnapshot := flag.Bool("restoreSnapshot", false, "This flag prompts the server to restore state from snapshot when set to true. Only works in standalone mode. Higher priority than restoreAOF.")
+	restoreAOF := flag.Bool("restoreAOF", false, "This flag prompts the server to restore state from append-only logs. Only works in standalone mode. Lower priority than restoreSnapshot.")
 	forwardCommand := flag.Bool(
 		"forwardCommand",
 		false,
@@ -93,6 +97,8 @@ It is a plain text value by default but you can provide a SHA256 hash by adding
 		Password:           *password,
 		SnapShotThreshold:  *snapshotThreshold,
 		SnapshotInterval:   *snapshotInterval,
+		RestoreSnapshot:    *restoreSnapshot,
+		RestoreAOF:         *restoreAOF,
 	}
 
 	if len(*config) > 0 {