From b2516eda9fdfc6838383685a6f5da557fbefff1a Mon Sep 17 00:00:00 2001
From: Pijus Kamandulis <pkpjuklas@gmail.com>
Date: Wed, 12 Mar 2025 22:00:30 +0200
Subject: [PATCH] Stability improvements

---
 api/tests/config_test.go                         |  1 +
 cmd/server/server.go                             |  3 +++
 .../badger_datastore/badger_datastore.go         |  1 +
 sharedlibrary/shared.go                          | 16 ++++++++--------
 sharedlibrary/sharedlibrary.go                   |  3 ++-
 5 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/api/tests/config_test.go b/api/tests/config_test.go
index cbc172b..c9c5628 100644
--- a/api/tests/config_test.go
+++ b/api/tests/config_test.go
@@ -90,6 +90,7 @@ func runTestsWithPreset(t *testing.T, name string, testPreset testPreset, f test
 
 	ts := runTestServerCustomConfig(serverConfig)
 	defer ts.Server.Close()
+	defer ts.DataStore.Close()
 
 	client, err := azcosmos.NewClientFromConnectionString(
 		fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
diff --git a/cmd/server/server.go b/cmd/server/server.go
index 6d6c9ab..ed49b25 100644
--- a/cmd/server/server.go
+++ b/cmd/server/server.go
@@ -10,6 +10,7 @@ import (
 	"github.com/pikami/cosmium/internal/datastore"
 	badgerdatastore "github.com/pikami/cosmium/internal/datastore/badger_datastore"
 	mapdatastore "github.com/pikami/cosmium/internal/datastore/map_datastore"
+	"github.com/pikami/cosmium/internal/logger"
 )
 
 func main() {
@@ -19,11 +20,13 @@ func main() {
 	switch configuration.DataStore {
 	case config.DataStoreBadger:
 		dataStore = badgerdatastore.NewBadgerDataStore()
+		logger.InfoLn("Using Badger data store")
 	default:
 		dataStore = mapdatastore.NewMapDataStore(mapdatastore.MapDataStoreOptions{
 			InitialDataFilePath: configuration.InitialDataFilePath,
 			PersistDataFilePath: configuration.PersistDataFilePath,
 		})
+		logger.InfoLn("Using in-memory data store")
 	}
 
 	server := api.NewApiServer(dataStore, &configuration)
diff --git a/internal/datastore/badger_datastore/badger_datastore.go b/internal/datastore/badger_datastore/badger_datastore.go
index 17292e4..690c6db 100644
--- a/internal/datastore/badger_datastore/badger_datastore.go
+++ b/internal/datastore/badger_datastore/badger_datastore.go
@@ -28,6 +28,7 @@ func NewBadgerDataStore() *BadgerDataStore {
 
 func (r *BadgerDataStore) Close() {
 	r.db.Close()
+	r.db = nil
 }
 
 func (r *BadgerDataStore) DumpToJson() (string, error) {
diff --git a/sharedlibrary/shared.go b/sharedlibrary/shared.go
index 8d08ab2..8302eba 100644
--- a/sharedlibrary/shared.go
+++ b/sharedlibrary/shared.go
@@ -20,14 +20,14 @@ var (
 const (
 	ResponseSuccess = 0
 
-	ResponseUnknown                                   = 100
-	ResponseFailedToParseConfiguration                = 101
-	ResponseFailedToLoadState                         = 102
-	ResponseFailedToParseRequest                      = 103
-	ResponseServerInstanceAlreadyExists               = 104
-	ResponseServerInstanceNotFound                    = 105
-	ResponseFailedToStartServer                       = 106
-	ResponseCurentDataStoreDoesNotSupportStateLoading = 107
+	ResponseUnknown                                    = 100
+	ResponseFailedToParseConfiguration                 = 101
+	ResponseFailedToLoadState                          = 102
+	ResponseFailedToParseRequest                       = 103
+	ResponseServerInstanceAlreadyExists                = 104
+	ResponseServerInstanceNotFound                     = 105
+	ResponseFailedToStartServer                        = 106
+	ResponseCurrentDataStoreDoesNotSupportStateLoading = 107
 
 	ResponseDataStoreNotFound   = 200
 	ResponseDataStoreConflict   = 201
diff --git a/sharedlibrary/sharedlibrary.go b/sharedlibrary/sharedlibrary.go
index d3edd62..5d69387 100644
--- a/sharedlibrary/sharedlibrary.go
+++ b/sharedlibrary/sharedlibrary.go
@@ -65,6 +65,7 @@ func StopServerInstance(serverName *C.char) int {
 
 	if serverInstance, ok := getInstance(serverNameStr); ok {
 		serverInstance.server.Stop()
+		serverInstance.dataStore.Close()
 		removeInstance(serverNameStr)
 		return ResponseSuccess
 	}
@@ -100,7 +101,7 @@ func LoadServerInstanceState(serverName *C.char, stateJSON *C.char) int {
 			}
 			return ResponseSuccess
 		}
-		return ResponseCurentDataStoreDoesNotSupportStateLoading
+		return ResponseCurrentDataStoreDoesNotSupportStateLoading
 	}
 
 	return ResponseServerInstanceNotFound