From 1dc92ed407b8229ae35c260da6ff91d926a4c440 Mon Sep 17 00:00:00 2001
From: Kane York <rikingcoding@gmail.com>
Date: Fri, 8 Jul 2016 13:08:36 -0700
Subject: [PATCH] Add tests

---
 socketserver/server/handlecore.go     |  5 +-
 socketserver/server/publisher.go      | 55 ++++++++++++++++------
 socketserver/server/publisher_test.go | 66 +++++++++++++++++++++++++++
 3 files changed, 109 insertions(+), 17 deletions(-)
 create mode 100644 socketserver/server/publisher_test.go

diff --git a/socketserver/server/handlecore.go b/socketserver/server/handlecore.go
index 1ed4fde4..5b426c63 100644
--- a/socketserver/server/handlecore.go
+++ b/socketserver/server/handlecore.go
@@ -131,10 +131,11 @@ func startJanitors() {
 	loadUniqueUsers()
 
 	go authorizationJanitor()
-	go bunchCacheJanitor()
-	go pubsubJanitor()
 	go aggregateDataSender()
+	go bunchCacheJanitor()
+	go cachedMessageJanitor()
 	go commandCounter()
+	go pubsubJanitor()
 
 	go ircConnection()
 	go shutdownHandler()
diff --git a/socketserver/server/publisher.go b/socketserver/server/publisher.go
index 0aa8e5f2..d8aff0a3 100644
--- a/socketserver/server/publisher.go
+++ b/socketserver/server/publisher.go
@@ -10,7 +10,7 @@ import (
 )
 
 type LastSavedMessage struct {
-	Timestamp time.Time
+	Expires time.Time
 	Data      string
 }
 
@@ -21,6 +21,31 @@ type LastSavedMessage struct {
 var CachedLastMessages = make(map[Command]map[string]LastSavedMessage)
 var CachedLSMLock sync.RWMutex
 
+func cachedMessageJanitor() {
+	for {
+		time.Sleep(1*time.Hour)
+		cachedMessageJanitor_do()
+	}
+}
+
+func cachedMessageJanitor_do() {
+	CachedLSMLock.Lock()
+	defer CachedLSMLock.Unlock()
+
+	now := time.Now()
+
+	for cmd, chanMap := range CachedLastMessages {
+		for channel, msg := range chanMap {
+			if !msg.Expires.IsZero() && msg.Expires.Before(now) {
+				delete(chanMap, channel)
+			}
+		}
+		if len(chanMap) == 0 {
+			delete(CachedLastMessages, cmd)
+		}
+	}
+}
+
 // DumpBacklogData drops all /cached_pub data.
 func DumpBacklogData() {
 	CachedLSMLock.Lock()
@@ -74,10 +99,8 @@ type timestampArray interface {
 	GetTime(int) time.Time
 }
 
-func SaveLastMessage(which map[Command]map[string]LastSavedMessage, locker sync.Locker, cmd Command, channel string, timestamp time.Time, data string, deleting bool) {
-	locker.Lock()
-	defer locker.Unlock()
-
+// the CachedLSMLock must be held when calling this
+func saveLastMessage(cmd Command, channel string, expires time.Time, data string, deleting bool) {
 	chanMap, ok := CachedLastMessages[cmd]
 	if !ok {
 		if deleting {
@@ -90,7 +113,7 @@ func SaveLastMessage(which map[Command]map[string]LastSavedMessage, locker sync.
 	if deleting {
 		delete(chanMap, channel)
 	} else {
-		chanMap[channel] = LastSavedMessage{Timestamp: timestamp, Data: data}
+		chanMap[channel] = LastSavedMessage{Expires: expires, Data: data}
 	}
 }
 
@@ -126,24 +149,26 @@ func HTTPBackendCachedPublish(w http.ResponseWriter, r *http.Request) {
 	json := formData.Get("args")
 	channel := formData.Get("channel")
 	deleteMode := formData.Get("delete") != ""
-	timeStr := formData.Get("time")
-	timeNum, err := strconv.ParseInt(timeStr, 10, 64)
-	if err != nil {
-		w.WriteHeader(422)
-		fmt.Fprintf(w, "error parsing time: %v", err)
-		return
+	timeStr := formData.Get("expires")
+	var expires time.Time
+	if timeStr != "" {
+		timeNum, err := strconv.ParseInt(timeStr, 10, 64)
+		if err != nil {
+			w.WriteHeader(422)
+			fmt.Fprintf(w, "error parsing time: %v", err)
+			return
+		}
+		expires = time.Unix(timeNum, 0)
 	}
-	timestamp := time.Unix(timeNum, 0)
 
 	var count int
 	msg := ClientMessage{MessageID: -1, Command: cmd, origArguments: json}
 	msg.parseOrigArguments()
 
 	channels := strings.Split(channel, ",")
-	var dummyLock sync.Mutex
 	CachedLSMLock.Lock()
 	for _, channel := range channels {
-		SaveLastMessage(CachedLastMessages, &dummyLock, cmd, channel, timestamp, json, deleteMode)
+		saveLastMessage(cmd, channel, expires, json, deleteMode)
 	}
 	CachedLSMLock.Unlock()
 	count = PublishToMultiple(channels, msg)
diff --git a/socketserver/server/publisher_test.go b/socketserver/server/publisher_test.go
new file mode 100644
index 00000000..a8f62b37
--- /dev/null
+++ b/socketserver/server/publisher_test.go
@@ -0,0 +1,66 @@
+package server
+
+import (
+	"testing"
+	"time"
+)
+
+func TestExpiredCleanup(t *testing.T) {
+	const cmd = "test_command"
+	const channel = "trihex"
+	const channel2 = "twitch"
+	const channel3 = "360chrism"
+	const channel4 = "qa_partner"
+
+	DumpBacklogData()
+	defer DumpBacklogData()
+
+	var zeroTime time.Time
+	hourAgo := time.Now().Add(-1*time.Hour)
+	now := time.Now()
+	hourFromNow := time.Now().Add(1*time.Hour)
+
+	saveLastMessage(cmd, channel, hourAgo, "1", false)
+	saveLastMessage(cmd, channel2, now, "2", false)
+
+	if len(CachedLastMessages) != 1 {
+		t.Error("messages not saved")
+	}
+	if len(CachedLastMessages[cmd]) != 2{
+		t.Error("messages not saved")
+	}
+
+	time.Sleep(2*time.Millisecond)
+
+	cachedMessageJanitor_do()
+
+	if len(CachedLastMessages) != 0 {
+		t.Error("messages still present")
+	}
+
+	saveLastMessage(cmd, channel, hourAgo, "1", false)
+	saveLastMessage(cmd, channel2, now, "2", false)
+	saveLastMessage(cmd, channel3, hourFromNow, "3", false)
+	saveLastMessage(cmd, channel4, zeroTime, "4", false)
+
+	if len(CachedLastMessages[cmd]) != 4 {
+		t.Error("messages not saved")
+	}
+
+	time.Sleep(2*time.Millisecond)
+
+	cachedMessageJanitor_do()
+
+	if len(CachedLastMessages) != 1 {
+		t.Error("messages not saved")
+	}
+	if len(CachedLastMessages[cmd]) != 2 {
+		t.Error("messages not saved")
+	}
+	if CachedLastMessages[cmd][channel3].Data != "3" {
+		t.Error("saved wrong message")
+	}
+	if CachedLastMessages[cmd][channel4].Data != "4" {
+		t.Error("saved wrong message")
+	}
+}
\ No newline at end of file