From 05bdc4461103929835af59eec55fb16e23d91f41 Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Thu, 12 Oct 2023 17:05:53 +0200
Subject: [PATCH] Fix waiting for log writers on shutdown, improve persistence
 enabling

---
 metrics/api.go          |  5 -----
 metrics/metrics_host.go |  2 +-
 metrics/metrics_logs.go |  2 +-
 metrics/module.go       | 27 ++++++++++++++-------------
 metrics/persistence.go  | 10 +++++++++-
 5 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/metrics/api.go b/metrics/api.go
index 638107d..4597882 100644
--- a/metrics/api.go
+++ b/metrics/api.go
@@ -11,7 +11,6 @@ import (
 	"github.com/safing/portbase/api"
 	"github.com/safing/portbase/config"
 	"github.com/safing/portbase/log"
-	"github.com/safing/portbase/utils"
 )
 
 func registerAPI() error {
@@ -140,11 +139,7 @@ func writeMetricsTo(ctx context.Context, url string) error {
 	)
 }
 
-var metricsPusherDone = utils.NewBroadcastFlag()
-
 func metricsWriter(ctx context.Context) error {
-	defer metricsPusherDone.NotifyAndReset()
-
 	pushURL := pushOption()
 	ticker := module.NewSleepyTicker(1*time.Minute, 0)
 	defer ticker.Stop()
diff --git a/metrics/metrics_host.go b/metrics/metrics_host.go
index ecf2472..e582c77 100644
--- a/metrics/metrics_host.go
+++ b/metrics/metrics_host.go
@@ -16,7 +16,7 @@ import (
 
 const hostStatTTL = 1 * time.Second
 
-func registeHostMetrics() (err error) {
+func registerHostMetrics() (err error) {
 	// Register load average metrics.
 	_, err = NewGauge("host/load/avg/1", nil, getFloat64HostStat(LoadAvg1), &Options{Name: "Host Load Avg 1min", Permission: api.PermitUser})
 	if err != nil {
diff --git a/metrics/metrics_logs.go b/metrics/metrics_logs.go
index fbf570a..8ec9697 100644
--- a/metrics/metrics_logs.go
+++ b/metrics/metrics_logs.go
@@ -5,7 +5,7 @@ import (
 	"github.com/safing/portbase/log"
 )
 
-func registeLogMetrics() (err error) {
+func registerLogMetrics() (err error) {
 	_, err = NewFetchingCounter(
 		"logs/warning/total",
 		nil,
diff --git a/metrics/module.go b/metrics/module.go
index 4dbd9b2..e7c9497 100644
--- a/metrics/module.go
+++ b/metrics/module.go
@@ -5,8 +5,8 @@ import (
 	"fmt"
 	"sort"
 	"sync"
-	"time"
 
+	"github.com/safing/portbase/log"
 	"github.com/safing/portbase/modules"
 )
 
@@ -36,7 +36,7 @@ var (
 )
 
 func init() {
-	module = modules.Register("metrics", prep, start, stop, "config", "database", "api")
+	module = modules.Register("metrics", prep, start, stop, "config", "database", "api", "base")
 }
 
 func prep() error {
@@ -59,11 +59,11 @@ func start() error {
 		return err
 	}
 
-	if err := registeHostMetrics(); err != nil {
+	if err := registerHostMetrics(); err != nil {
 		return err
 	}
 
-	if err := registeLogMetrics(); err != nil {
+	if err := registerLogMetrics(); err != nil {
 		return err
 	}
 
@@ -71,6 +71,10 @@ func start() error {
 		return err
 	}
 
+	if err := loadPersistentMetrics(); err != nil {
+		log.Errorf("metrics: failed to load persistent metrics: %s", err)
+	}
+
 	if pushOption() != "" {
 		module.StartServiceWorker("metric pusher", 0, metricsWriter)
 	}
@@ -82,16 +86,13 @@ func stop() error {
 	// Wait until the metrics pusher is done, as it may have started reporting
 	// and may report a higher number than we store to disk. For persistent
 	// metrics it can then happen that the first report is lower than the
-	// previous report, making prometheus think that al that happened since the
+	// previous report, making prometheus think that all that happened since the
 	// last report, due to the automatic restart detection.
-	done := metricsPusherDone.NewFlag()
-	done.Refresh()
-	if !done.IsSet() {
-		select {
-		case <-done.Signal():
-		case <-time.After(10 * time.Second):
-		}
-	}
+
+	// The registry is read locked when writing metrics.
+	// Write lock the registry to make sure all writes are finished.
+	registryLock.Lock()
+	registryLock.Unlock() //nolint:staticcheck
 
 	storePersistentMetrics()
 
diff --git a/metrics/persistence.go b/metrics/persistence.go
index ad77bac..b0c1b5f 100644
--- a/metrics/persistence.go
+++ b/metrics/persistence.go
@@ -52,10 +52,18 @@ func EnableMetricPersistence(key string) error {
 
 	// Set storage key.
 	storageKey = key
+	return nil
+}
+
+func loadPersistentMetrics() error {
+	// Abort if storage is not enabled.
+	if storageInit.SetToIf(false, true) {
+		return nil
+	}
 
 	// Load metrics from storage.
 	var err error
-	storage, err = getMetricsStorage(key)
+	storage, err = getMetricsStorage(storageKey)
 	switch {
 	case err == nil:
 		// Continue.