memory storage message limit

Author: Ian Bishop

memory.go

@@ -10,20 +10,48 @@ import (
 type InMemory struct {
 	MessageIDIndex map[string]int
 	Messages       []*data.Message
+
+	sliceIndexMap map[int]string
+	messageLimit  int
+	writeIndex    int
 }
 
 // CreateInMemory creates a new in memory storage backend
 func CreateInMemory() *InMemory {
 	return &InMemory{
 		MessageIDIndex: make(map[string]int),
 		Messages:       make([]*data.Message, 0),
+		sliceIndexMap:  make(map[int]string),
 	}
 }
 
+// Set a limit on the number of messages that will be stored.
+// Once this limit is reached, the oldest message will be overwritten
+// by every new message.
+func (memory *InMemory) SetMessageLimit(limit int) error {
+	memory.messageLimit = limit
+	return nil
+}
+
 // Store stores a message and returns its storage ID
 func (memory *InMemory) Store(m *data.Message) (string, error) {
-	memory.Messages = append(memory.Messages, m)
-	memory.MessageIDIndex[string(m.ID)] = len(memory.Messages) - 1
+	if memory.messageLimit > 0 && len(memory.Messages) == memory.messageLimit {
+		// overwrite oldest message
+		memory.Messages[memory.writeIndex] = m
+		messageID := memory.sliceIndexMap[memory.writeIndex]
+		delete(memory.MessageIDIndex, messageID)
+		memory.MessageIDIndex[string(m.ID)] = memory.writeIndex
+		memory.sliceIndexMap[memory.writeIndex] = string(m.ID)
+		memory.writeIndex++
+		// if we reach the end of the slice then start over
+		if memory.writeIndex == len(memory.Messages) {
+			memory.writeIndex = 0
+		}
+	} else {
+		memory.Messages = append(memory.Messages, m)
+		memory.MessageIDIndex[string(m.ID)] = len(memory.Messages) - 1
+		memory.sliceIndexMap[len(memory.Messages)-1] = string(m.ID)
+	}
 	return string(m.ID), nil
 }