From 0508f378969b504530890b4997330091ab54af29 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thorben=20G=C3=BCnther?= <admin@xenrox.net>
Date: Wed, 8 Feb 2023 14:46:14 +0100
Subject: [PATCH] cache: Add config options for duration and cleanup interval

---
 README.md      |  9 +++++++++
 cache.go       | 12 +++++++-----
 config.go      | 44 ++++++++++++++++++++++++++++++++++++++++++++
 config_test.go | 16 +++++++++++++++-
 main.go        |  8 +++-----
 5 files changed, 78 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 50fb575..78524ce 100644
--- a/README.md
+++ b/README.md
@@ -60,6 +60,15 @@ ntfy {
     user user
     password pass
 }
+
+# When the alert-mode is set to single, ntfy-alertmanager will cache each single alert
+# to avoid sending recurrences.
+cache {
+    # How long messages stay in the cache for
+    duration 24h
+    # Interval in which the cache is cleaned up
+    cleanup-interval 1h
+}
 ```
 
 ### Alertmanager config
diff --git a/cache.go b/cache.go
index 72f0f6a..17bead6 100644
--- a/cache.go
+++ b/cache.go
@@ -12,16 +12,18 @@ type cachedAlert struct {
 }
 
 type cache struct {
-	mu     sync.Mutex
-	alerts map[fingerprint]*cachedAlert
+	mu       sync.Mutex
+	duration time.Duration
+	alerts   map[fingerprint]*cachedAlert
 }
 
 func (a *cachedAlert) expired() bool {
 	return a.expires.Before(time.Now())
 }
 
-func newCache() *cache {
+func newCache(d time.Duration) *cache {
 	c := new(cache)
+	c.duration = d
 	c.alerts = make(map[fingerprint]*cachedAlert)
 
 	return c
@@ -37,11 +39,11 @@ func (c *cache) cleanup() {
 	}
 }
 
-func (c *cache) set(f fingerprint, d time.Duration) {
+func (c *cache) set(f fingerprint) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	alert := new(cachedAlert)
-	alert.expires = time.Now().Add(d)
+	alert.expires = time.Now().Add(c.duration)
 
 	c.alerts[f] = alert
 }
diff --git a/config.go b/config.go
index 7fb0653..a6c27a2 100644
--- a/config.go
+++ b/config.go
@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"git.sr.ht/~emersion/go-scfg"
 )
@@ -22,6 +23,7 @@ type config struct {
 	Password    string
 	ntfy        ntfyConfig
 	labels      labels
+	cache       cacheConfig
 }
 
 type ntfyConfig struct {
@@ -40,6 +42,11 @@ type labelConfig struct {
 	Tags     []string
 }
 
+type cacheConfig struct {
+	CleanupInterval time.Duration
+	Duration        time.Duration
+}
+
 func readConfig(path string) (*config, error) {
 	cfg, err := scfg.Load(path)
 	if err != nil {
@@ -52,6 +59,9 @@ func readConfig(path string) (*config, error) {
 	config.LogLevel = "info"
 	config.alertMode = single
 
+	config.cache.CleanupInterval = time.Hour
+	config.cache.Duration = time.Hour * 24
+
 	d := cfg.Get("log-level")
 	if d != nil {
 		if err := d.ParseParams(&config.LogLevel); err != nil {
@@ -172,5 +182,39 @@ func readConfig(path string) (*config, error) {
 		}
 	}
 
+	cacheDir := cfg.Get("cache")
+
+	if cacheDir != nil {
+		var durationString string
+		d = cacheDir.Children.Get("duration")
+		if d != nil {
+			if err := d.ParseParams(&durationString); err != nil {
+				return nil, err
+			}
+
+			duration, err := time.ParseDuration(durationString)
+			if err != nil {
+				return nil, err
+			}
+
+			config.cache.Duration = duration
+		}
+
+		var cleanupIntervalString string
+		d = cacheDir.Children.Get("cleanup-interval")
+		if d != nil {
+			if err := d.ParseParams(&cleanupIntervalString); err != nil {
+				return nil, err
+			}
+
+			interval, err := time.ParseDuration(cleanupIntervalString)
+			if err != nil {
+				return nil, err
+			}
+
+			config.cache.CleanupInterval = interval
+		}
+	}
+
 	return config, nil
 }
diff --git a/config_test.go b/config_test.go
index 8e3b4e8..a64fc60 100644
--- a/config_test.go
+++ b/config_test.go
@@ -5,6 +5,7 @@ import (
 	"path/filepath"
 	"reflect"
 	"testing"
+	"time"
 )
 
 func TestReadConfig(t *testing.T) {
@@ -44,6 +45,15 @@ ntfy {
     user user
     password pass
 }
+
+# When the alert-mode is set to single, ntfy-alertmanager will cache each single alert
+# to avoid sending recurrences.
+cache {
+    # How long messages stay in the cache for
+    duration 48h
+    # Interval in which the cache is cleaned up
+    # cleanup-interval 1h
+}
 `
 
 	expectedCfg := &config{
@@ -53,7 +63,11 @@ ntfy {
 			Label: map[string]labelConfig{
 				"severity:critical":    {Priority: "5", Tags: []string{"rotating_light"}},
 				"severity:info":        {Priority: "1"},
-				"instance:example.com": {Tags: []string{"computer", "example"}}}}}
+				"instance:example.com": {Tags: []string{"computer", "example"}},
+			},
+		},
+		cache: cacheConfig{CleanupInterval: time.Hour, Duration: 48 * time.Hour},
+	}
 
 	configPath := filepath.Join(t.TempDir(), "config")
 	err := os.WriteFile(configPath, []byte(configContent), 0600)
diff --git a/main.go b/main.go
index 9e00fac..00490b4 100644
--- a/main.go
+++ b/main.go
@@ -58,8 +58,7 @@ func (rcv *receiver) singleAlertNotifications(p *payload) []*notification {
 			rcv.logger.Debugf("Alert %s skipped: Still in cache", alert.Fingerprint)
 			continue
 		}
-		// TODO: Make configurable
-		rcv.cache.set(alert.Fingerprint, time.Hour*24)
+		rcv.cache.set(alert.Fingerprint)
 
 		n := new(notification)
 
@@ -286,8 +285,7 @@ func (rcv *receiver) basicAuthMiddleware(handler http.HandlerFunc) http.HandlerF
 
 func (rcv *receiver) runCleanup() {
 	for {
-		// TODO: Make configurable
-		time.Sleep(time.Hour)
+		time.Sleep(rcv.cfg.cache.CleanupInterval)
 		rcv.logger.Info("Pruning cache")
 		rcv.cache.cleanup()
 	}
@@ -316,7 +314,7 @@ func main() {
 		logger.Errorf("config: %v", err)
 	}
 
-	receiver := &receiver{cfg: cfg, logger: logger, cache: newCache()}
+	receiver := &receiver{cfg: cfg, logger: logger, cache: newCache(cfg.cache.Duration)}
 
 	logger.Infof("Listening on %s, ntfy-alertmanager %s", cfg.HTTPAddress, version)