Feature: refactor storage to save/load messages by topic

Author: ThomasBoom89

internal/storage.go

@@ -5,54 +5,103 @@ import (
 	"os"
 )
 
+const DataDir = "./data/"
+const FileExcluded = ".gitkeep"
+
 type Storage struct {
-	queue Queue[string]
+	topicManager *TopicManager
 }
 
-func NewStorage(queue Queue[string]) *Storage {
+func NewStorage(topicManager *TopicManager) *Storage {
 	return &Storage{
-		queue: queue,
+		topicManager: topicManager,
+	}
+}
+
+func (S *Storage) createStoreDirectory() {
+	_, err := os.ReadDir(DataDir)
+	if err != nil {
+		if os.IsNotExist(err) {
+			err := os.Mkdir(DataDir, 0755)
+			if err != nil {
+				panic(err)
+			}
+		} else if !os.IsExist(err) {
+			panic(err)
+		}
 	}
 }
 
 func (S *Storage) clear() {
-	err := os.WriteFile("./data", []byte(""), 0644)
+	dirEntries, err := os.ReadDir(DataDir)
 	if err != nil {
 		panic(err)
 	}
+	for _, entry := range dirEntries {
+		if entry.Name() == FileExcluded {
+			continue
+		}
+		os.Remove(DataDir + entry.Name())
+	}
 }
 
 func (S *Storage) Save() {
 	S.clear()
+	S.createStoreDirectory()
+
+	for _, topic := range S.topicManager.GetTopics() {
+		S.saveTopic(topic)
+	}
+}
 
-	file, err := os.OpenFile("./data", os.O_RDWR|os.O_CREATE, 0644)
+func (S *Storage) saveTopic(topic Topic) {
+	file, err := os.OpenFile(string(DataDir+topic), os.O_RDWR|os.O_CREATE, 0644)
 	if err != nil {
 		panic(err)
 	}
 	defer file.Close()
-	for !S.queue.IsEmpty() {
-		value, _ := S.queue.Dequeue()
-		//err := os.OpenFile("./data", []byte(value+"\n"), 0644)
-		_, err = file.WriteString(value + "\n")
+	for {
+		message, err := S.topicManager.GetNextMessage(topic)
+		if err != nil {
+			return
+		}
+		_, err = file.Write(message)
 		if err != nil {
 			panic(err)
 		}
+		// Line feed
+		_, _ = file.Write([]byte{10})
 	}
 }
 
 func (S *Storage) Load() {
-	file, err := os.OpenFile("./data", os.O_RDONLY, 0644)
+	dirEntries, err := os.ReadDir(DataDir)
+	if err != nil {
+		panic(err)
+	}
+
+	for _, dirEntry := range dirEntries {
+		if dirEntry.Name() == FileExcluded {
+			continue
+		}
+		S.loadTopic(dirEntry)
+	}
+}
+
+func (S *Storage) loadTopic(dirEntry os.DirEntry) {
+	topic := dirEntry.Name()
+	file, err := os.OpenFile(DataDir+topic, os.O_RDONLY, 0644)
 	if err != nil {
-		return
+		panic(err)
 	}
 	defer file.Close()
 
 	scanner := bufio.NewScanner(file)
 	for scanner.Scan() {
-		text := scanner.Text()
-		if text == "" {
+		text := scanner.Bytes()
+		if len(text) <= 0 {
 			continue
 		}
-		S.queue.Enqueue(text)
+		S.topicManager.AddMessage(Topic(topic), text)
 	}
 }

internal/system.go

@@ -4,11 +4,12 @@ import (
 	"github.com/gofiber/fiber/v2/log"
 	"os"
 	"os/signal"
+	"syscall"
 )
 
 func HandleOsSignal(storage *Storage) {
 	sigchan := make(chan os.Signal, 1)
-	signals := []os.Signal{os.Kill, os.Interrupt}
+	signals := []os.Signal{os.Kill, os.Interrupt, syscall.SIGTERM}
 	signal.Notify(sigchan, signals...)
 	<-sigchan
 

internal/topic.go

@@ -9,6 +9,7 @@ type TopicManager struct {
 	messageQueue    map[Topic]Queue[Message]
 	connectionQueue map[Topic]Queue[*websocket.Conn]
 	connections     map[Topic]map[*websocket.Conn]bool
+	topics          []Topic
 }
 
 type Topic string
@@ -38,6 +39,7 @@ func (T *TopicManager) GetNextMessage(topic Topic) (Message, error) {
 func (T *TopicManager) AddMessage(topic Topic, message Message) {
 	if _, ok := T.messageQueue[topic]; !ok {
 		T.messageQueue[topic] = NewLinkedListQueue[Message]()
+		T.topics = append(T.topics, topic)
 	}
 	T.messageQueue[topic].Enqueue(message)
 }
@@ -58,12 +60,7 @@ func (T *TopicManager) RemoveConnection(topic Topic, conn *websocket.Conn) {
 }
 
 func (T *TopicManager) GetTopics() []Topic {
-	topics := make([]Topic, 0)
-	for topic := range T.messageQueue {
-		topics = append(topics, topic)
-	}
-
-	return topics
+	return T.topics
 }
 
 func (T *TopicManager) GetNextConnection(topic Topic) (*websocket.Conn, error) {

src/main.go

@@ -10,8 +10,7 @@ import (
 
 func main() {
 	topicManager := internal.NewTopicManager()
-	queue := internal.NewLinkedListQueue[string]()
-	storage := internal.NewStorage(queue)
+	storage := internal.NewStorage(topicManager)
 	storage.Load()
 	defer internal.SaveOnPanic(storage)
 	go internal.HandleOsSignal(storage)