internal/aggregate: make shutdown context-aware proper goroutine sync

Adds `WaitGroup` to track background goroutine and make Shutdown respect
context deadlines

Signed-off-by: NotAShelf <raf@notashelf.dev>
Change-Id: Ia7f074725717f037412dacb93e34105b6a6a6964

internal/aggregate/metrics_test.go

@@ -2,9 +2,11 @@ package aggregate
 
 import (
 	"context"
+	"errors"
 	"os"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/testutil"
@@ -278,3 +280,160 @@ func TestMetricsAggregator_Shutdown_DefaultStatePath(t *testing.T) {
 		t.Logf("Shutdown returned error (might be expected): %v", err)
 	}
 }
+
+func TestMetricsAggregator_Shutdown_RespectsContext(t *testing.T) {
+	registry := NewPathRegistry(100)
+	tmpDir := t.TempDir()
+
+	cfg := config.Config{
+		Site: config.SiteConfig{
+			SaltRotation: "daily",
+			Collect: config.CollectConfig{
+				Pageviews: true,
+			},
+		},
+		Server: config.ServerConfig{
+			StatePath: tmpDir + "/hll.state",
+		},
+	}
+
+	agg := NewMetricsAggregator(registry, NewCustomEventRegistry(100), &cfg)
+
+	// Create a context with very short timeout
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
+	defer cancel()
+
+	// Wait for context to expire
+	time.Sleep(10 * time.Millisecond)
+
+	// Shutdown should respect context timeout
+	err := agg.Shutdown(ctx)
+	if err == nil {
+		t.Error("expected context deadline exceeded error, got nil")
+	}
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Errorf(
+			"expected context.DeadlineExceeded, got %v",
+			err,
+		)
+	}
+}
+
+func TestMetricsAggregator_Shutdown_WaitsForGoroutine(t *testing.T) {
+	registry := NewPathRegistry(100)
+	tmpDir := t.TempDir()
+
+	cfg := config.Config{
+		Site: config.SiteConfig{
+			SaltRotation: "daily",
+			Collect: config.CollectConfig{
+				Pageviews: true,
+			},
+		},
+		Server: config.ServerConfig{
+			StatePath: tmpDir + "/hll.state",
+		},
+	}
+
+	agg := NewMetricsAggregator(registry, NewCustomEventRegistry(100), &cfg)
+
+	// Give the goroutine time to start
+	time.Sleep(10 * time.Millisecond)
+
+	// Track if goroutine is still running
+	done := make(chan struct{})
+	go func() {
+		ctx := context.Background()
+		agg.Shutdown(ctx)
+		close(done)
+	}()
+
+	// Shutdown should complete quickly (goroutine should stop)
+	select {
+	case <-done:
+		// Success - shutdown completed
+	case <-time.After(1 * time.Second):
+		t.Fatal("Shutdown did not complete within timeout - goroutine not stopping")
+	}
+}
+
+func TestMetricsAggregator_LoadState(t *testing.T) {
+	tmpDir := t.TempDir()
+	statePath := tmpDir + "/hll.state"
+
+	cfg := config.Config{
+		Site: config.SiteConfig{
+			SaltRotation: "daily",
+			Collect: config.CollectConfig{
+				Pageviews: true,
+			},
+		},
+		Server: config.ServerConfig{
+			StatePath: statePath,
+		},
+	}
+
+	// Create first aggregator and add some visitors
+	registry1 := NewPathRegistry(100)
+	agg1 := NewMetricsAggregator(registry1, NewCustomEventRegistry(100), &cfg)
+	agg1.AddUnique("192.168.1.1", "Mozilla/5.0")
+	agg1.AddUnique("192.168.1.2", "Mozilla/5.0")
+	agg1.AddUnique("192.168.1.3", "Mozilla/5.0")
+
+	// Get estimate before shutdown
+	estimate1 := agg1.estimator.Estimate()
+	if estimate1 == 0 {
+		t.Fatal("expected non-zero estimate before shutdown")
+	}
+
+	// Shutdown to save state
+	ctx := context.Background()
+	if err := agg1.Shutdown(ctx); err != nil {
+		t.Fatalf("Shutdown failed: %v", err)
+	}
+
+	// Verify state file was created
+	if _, err := os.Stat(statePath); os.IsNotExist(err) {
+		t.Fatal("state file was not created")
+	}
+
+	// Create second aggregator and load state
+	registry2 := NewPathRegistry(100)
+	agg2 := NewMetricsAggregator(registry2, NewCustomEventRegistry(100), &cfg)
+
+	// Load should restore the state
+	if err := agg2.LoadState(); err != nil {
+		t.Fatalf("LoadState failed: %v", err)
+	}
+
+	// Estimate should match (approximately - HLL is probabilistic)
+	estimate2 := agg2.estimator.Estimate()
+	if estimate2 != estimate1 {
+		t.Errorf("expected estimate %d after load, got %d", estimate1, estimate2)
+	}
+}
+
+func TestMetricsAggregator_LoadState_NoFile(t *testing.T) {
+	tmpDir := t.TempDir()
+	statePath := tmpDir + "/nonexistent.state"
+
+	cfg := config.Config{
+		Site: config.SiteConfig{
+			SaltRotation: "daily",
+			Collect: config.CollectConfig{
+				Pageviews: true,
+			},
+		},
+		Server: config.ServerConfig{
+			StatePath: statePath,
+		},
+	}
+
+	registry := NewPathRegistry(100)
+	agg := NewMetricsAggregator(registry, NewCustomEventRegistry(100), &cfg)
+
+	// LoadState should not error if file doesn't exist (first run)
+	if err := agg.LoadState(); err != nil {
+		t.Errorf("LoadState should not error on missing file, got: %v", err)
+	}
+}