From 9fa0f7d61f029bc3617eeb997896e4fc270151e8 Mon Sep 17 00:00:00 2001
From: Fred Boniface <fred@fjla.uk>
Date: Fri, 21 Jul 2023 10:02:55 +0100
Subject: [PATCH] Organize cleanup code

---
 src/dbAccess/client.go | 12 ++++++++++++
 src/main.go            |  9 ++++++++-
 src/vstp/subscribe.go  |  4 ++--
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/dbAccess/client.go b/src/dbAccess/client.go
index 7ae33d1..c7223e4 100644
--- a/src/dbAccess/client.go
+++ b/src/dbAccess/client.go
@@ -54,3 +54,15 @@ func initDataAccess() *mongo.Client {
 	}
 	return client
 }
+
+func CloseMongoClient() {
+	if MongoClient != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		if err := MongoClient.Disconnect(ctx); err != nil {
+			log.Msg.Warn("Error disconnecting MongoDB client: " + err.Error())
+		} else {
+			log.Msg.Info("MongoDB client disconnected.")
+		}
+	}
+}
diff --git a/src/main.go b/src/main.go
index daaf0a8..026aec9 100644
--- a/src/main.go
+++ b/src/main.go
@@ -10,7 +10,14 @@ import (
 
 func main() {
 	log.Msg.Info("Initialised OwlBoard MQ Client " + helpers.Version)
+	defer cleanup()
 	dbAccess.PrintFromDbPackage()
-	defer messaging.Disconnect(messaging.Client)
 	vstp.Subscribe()
 }
+
+func cleanup() {
+	log.Msg.Debug("Cleaning up open connections")
+	defer messaging.Disconnect(messaging.Client)
+	defer dbAccess.CloseMongoClient()
+	log.Msg.Info("Program ready to exit")
+}
diff --git a/src/vstp/subscribe.go b/src/vstp/subscribe.go
index 7f68826..e53fcbd 100644
--- a/src/vstp/subscribe.go
+++ b/src/vstp/subscribe.go
@@ -19,9 +19,9 @@ func Subscribe() {
 		for {
 			msg := <-sub.C
 			if msg.Err != nil {
-				log.Msg.Error("\nSTOMP Message Error: " + msg.Err.Error())
+				log.Msg.Error("STOMP Message Error: " + msg.Err.Error())
 			} else {
-				log.Msg.Debug("\nSTOMP Message Received")
+				log.Msg.Debug("STOMP Message Received")
 				saveToFile(string(msg.Body), "msgBody")
 				handle(msg)
 			}