fix(hub): stop launching system updaters after manager shutdown (#1950) (#1951)

The staggered Initialize goroutine kept calling AddSystem after the hub
had been torn down, which spawned StartUpdater goroutines that later
dereferenced a nil DB inside PocketBase. Track a manager-level context
that RemoveAllSystems cancels, abort the staggered loop on shutdown,
and refuse AddSystem once the manager is stopped.

agent/connection_manager.go

@@ -4,11 +4,15 @@ import (
 	"context"
 	"errors"
 	"log/slog"
+	"net"
+	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 	"time"
 
 	"github.com/henrygd/beszel/agent/health"
+	"github.com/henrygd/beszel/agent/utils"
 	"github.com/henrygd/beszel/internal/entities/system"
 )
 
@@ -111,13 +115,36 @@ func (c *ConnectionManager) Start(serverOptions ServerOptions) error {
 			_ = health.Update()
 		case <-sigCtx.Done():
 			slog.Info("Shutting down", "cause", context.Cause(sigCtx))
-			_ = c.agent.StopServer()
-			c.closeWebSocket()
-			return health.CleanUp()
+			return c.stop()
 		}
 	}
 }
 
+// stop does not stop the connection manager itself, just any active connections. The manager will attempt to reconnect after stopping, so this should only be called immediately before shutting down the entire agent.
+//
+// If we need or want to expose a graceful Stop method in the future, do something like this to actually stop the manager:
+//
+//	func (c *ConnectionManager) Start(serverOptions ServerOptions) error {
+//		ctx, cancel := context.WithCancel(context.Background())
+//		c.cancel = cancel
+//
+//		for {
+//			select {
+//			case <-ctx.Done():
+//				return c.stop()
+//			}
+//		}
+//	}
+//
+//	func (c *ConnectionManager) Stop() {
+//		c.cancel()
+//	}
+func (c *ConnectionManager) stop() error {
+	_ = c.agent.StopServer()
+	c.closeWebSocket()
+	return health.CleanUp()
+}
+
 // handleEvent processes connection events and updates the connection state accordingly.
 func (c *ConnectionManager) handleEvent(event ConnectionEvent) {
 	switch event {
@@ -185,9 +212,16 @@ func (c *ConnectionManager) connect() {
 
 	// Try WebSocket first, if it fails, start SSH server
 	err := c.startWebSocketConnection()
-	if err != nil && c.State == Disconnected {
-		c.startSSHServer()
-		c.startWsTicker()
+	if err != nil {
+		if shouldExitOnErr(err) {
+			time.Sleep(2 * time.Second) // prevent tight restart loop
+			_ = c.stop()
+			os.Exit(1)
+		}
+		if c.State == Disconnected {
+			c.startSSHServer()
+			c.startWsTicker()
+		}
 	}
 }
 
@@ -224,3 +258,14 @@ func (c *ConnectionManager) closeWebSocket() {
 		c.wsClient.Close()
 	}
 }
+
+// shouldExitOnErr checks if the error is a DNS resolution failure and if the
+// EXIT_ON_DNS_ERROR env var is set. https://github.com/henrygd/beszel/issues/1924.
+func shouldExitOnErr(err error) bool {
+	if val, _ := utils.GetEnv("EXIT_ON_DNS_ERROR"); val == "true" {
+		if opErr, ok := errors.AsType[*net.OpError](err); ok {
+			return strings.Contains(opErr.Err.Error(), "lookup")
+		}
+	}
+	return false
+}

agent/connection_manager_test.go

@@ -4,6 +4,7 @@ package agent
 
 import (
 	"crypto/ed25519"
+	"errors"
 	"fmt"
 	"net"
 	"net/url"
@@ -298,3 +299,65 @@ func TestConnectionManager_ConnectFlow(t *testing.T) {
 		cm.connect()
 	}, "Connect should not panic without WebSocket client")
 }
+
+func TestShouldExitOnErr(t *testing.T) {
+	createDialErr := func(msg string) error {
+		return &net.OpError{
+			Op:  "dial",
+			Net: "tcp",
+			Err: errors.New(msg),
+		}
+	}
+
+	tests := []struct {
+		name     string
+		err      error
+		envValue string
+		expected bool
+	}{
+		{
+			name:     "no env var",
+			err:      createDialErr("lookup lkahsdfasdf: no such host"),
+			envValue: "",
+			expected: false,
+		},
+		{
+			name:     "env var false",
+			err:      createDialErr("lookup lkahsdfasdf: no such host"),
+			envValue: "false",
+			expected: false,
+		},
+		{
+			name:     "env var true, matching error",
+			err:      createDialErr("lookup lkahsdfasdf: no such host"),
+			envValue: "true",
+			expected: true,
+		},
+		{
+			name:     "env var true, matching error with extra context",
+			err:      createDialErr("lookup beszel.server.lan on [::1]:53: read udp [::1]:44557->[::1]:53: read: connection refused"),
+			envValue: "true",
+			expected: true,
+		},
+		{
+			name:     "env var true, non-matching error",
+			err:      errors.New("connection refused"),
+			envValue: "true",
+			expected: false,
+		},
+		{
+			name:     "env var true, dial but not lookup",
+			err:      createDialErr("connection timeout"),
+			envValue: "true",
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Setenv("EXIT_ON_DNS_ERROR", tt.envValue)
+			result := shouldExitOnErr(tt.err)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}

internal/cmd/agent/agent.go

@@ -195,6 +195,6 @@ func main() {
 	}
 
 	if err := a.Start(serverConfig); err != nil {
-		log.Fatal("Failed to start server: ", err)
+		log.Fatal("Failed to start: ", err)
 	}
 }

internal/hub/systems/system_manager.go

@@ -1,6 +1,7 @@
 package systems
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"time"
@@ -45,6 +46,8 @@ type SystemManager struct {
 	systems       *store.Store[string, *System]         // Thread-safe store of active systems
 	sshConfig     *ssh.ClientConfig                     // SSH client configuration for system connections
 	smartFetchMap *expirymap.ExpiryMap[smartFetchState] // Stores last SMART fetch time/result; TTL is only for cleanup
+	ctx           context.Context                       // Cancelled when the manager is shutting down
+	cancel        context.CancelFunc                    // Cancels ctx
 }
 
 // hubLike defines the interface requirements for the hub dependency.
@@ -60,10 +63,13 @@ type hubLike interface {
 // NewSystemManager creates a new SystemManager instance with the provided hub.
 // The hub must implement the hubLike interface to provide database and alert functionality.
 func NewSystemManager(hub hubLike) *SystemManager {
+	ctx, cancel := context.WithCancel(context.Background())
 	return &SystemManager{
 		systems:       store.New(map[string]*System{}),
 		hub:           hub,
 		smartFetchMap: expirymap.New[smartFetchState](time.Hour),
+		ctx:           ctx,
+		cancel:        cancel,
 	}
 }
 
@@ -103,7 +109,13 @@ func (sm *SystemManager) Initialize() error {
 		sleepTime := time.Duration(delta) * time.Millisecond
 
 		for _, system := range systems {
-			time.Sleep(sleepTime)
+			select {
+			case <-sm.ctx.Done():
+				// Abort if the manager has been shut down (e.g. test cleanup)
+				// to avoid starting updater goroutines against a torn-down hub. See #1950.
+				return
+			case <-time.After(sleepTime):
+			}
 			_ = sm.AddSystem(system)
 		}
 	}()
@@ -238,6 +250,11 @@ func (sm *SystemManager) onRecordAfterDeleteSuccess(e *core.RecordEvent) error {
 // It validates required fields, initializes the system context, and starts the update goroutine.
 // Returns error if a system with the same ID already exists.
 func (sm *SystemManager) AddSystem(sys *System) error {
+	if sm.ctx.Err() != nil {
+		// Manager is shutting down; do not start new updater goroutines
+		// against a hub that may be torn down. See #1950.
+		return sm.ctx.Err()
+	}
 	if sm.systems.Has(sys.Id) {
 		return errSystemExists
 	}

internal/hub/systems/system_test.go

@@ -3,9 +3,11 @@
 package systems
 
 import (
+	"context"
 	"testing"
 
 	"github.com/henrygd/beszel/internal/entities/system"
+	"github.com/pocketbase/pocketbase/tools/store"
 )
 
 func TestCombinedData_MigrateDeprecatedFields(t *testing.T) {
@@ -157,3 +159,20 @@ func TestCombinedData_MigrateDeprecatedFields(t *testing.T) {
 		}
 	})
 }
+
+// TestAddSystemRefusedAfterShutdown verifies that AddSystem returns an error
+// once the manager has been shut down, so the staggered Initialize starter
+// cannot spawn updater goroutines against a torn-down hub. See #1950.
+func TestAddSystemRefusedAfterShutdown(t *testing.T) {
+	sm := &SystemManager{systems: store.New(map[string]*System{})}
+	sm.ctx, sm.cancel = context.WithCancel(context.Background())
+	sm.cancel()
+
+	err := sm.AddSystem(&System{Id: "sys", Host: "127.0.0.1"})
+	if err == nil {
+		t.Fatalf("AddSystem returned nil error after shutdown")
+	}
+	if sm.systems.Has("sys") {
+		t.Fatalf("system was added to store after shutdown")
+	}
+}

internal/hub/systems/systems_test_helpers.go

@@ -111,6 +111,11 @@ func (sm *SystemManager) SetSystemStatusInDB(systemID string, status string) boo
 
 // TESTING ONLY: RemoveAllSystems removes all systems from the store
 func (sm *SystemManager) RemoveAllSystems() {
+	// Signal shutdown first so any in-flight Initialize staggered-start goroutine
+	// stops adding new systems against a hub that is about to be torn down. See #1950.
+	if sm.cancel != nil {
+		sm.cancel()
+	}
 	for _, system := range sm.systems.GetAll() {
 		sm.RemoveSystem(system.Id)
 	}

internal/site/src/components/charts/line-chart.tsx

@@ -104,7 +104,7 @@ export default function LineChartDefault({
 					isAnimationActive={false}
 					// stackId={dataPoint.stackId}
 					order={dataPoint.order || i}
-					// activeDot={dataPoint.activeDot ?? true}
+					activeDot={dataPoint.activeDot ?? true}
 				/>
 			)
 		})

internal/site/src/components/login/auth-form.tsx

@@ -17,7 +17,7 @@ import { toast } from "../ui/use-toast"
 import { OtpInputForm } from "./otp-forms"
 
 const honeypot = v.literal("")
-const emailSchema = v.pipe(v.string(), v.email(t`Invalid email address.`))
+const emailSchema = v.pipe(v.string(), v.rfcEmail(t`Invalid email address.`))
 const passwordSchema = v.pipe(
 	v.string(),
 	v.minLength(8, t`Password must be at least 8 characters.`),

internal/site/src/components/routes/settings/notifications.tsx

@@ -24,7 +24,7 @@ interface ShoutrrrUrlCardProps {
 }
 
 const NotificationSchema = v.object({
-	emails: v.array(v.pipe(v.string(), v.email())),
+	emails: v.array(v.pipe(v.string(), v.rfcEmail())),
 	webhooks: v.array(v.pipe(v.string(), v.url())),
 })
 

internal/site/src/components/routes/system/charts/sensor-charts.tsx

@@ -120,7 +120,8 @@ export function TemperatureChart({
 				label: key,
 				dataKey: dataKeys[key],
 				color: colorMap[key],
-				opacity: strokeOpacity,
+				strokeOpacity,
+				activeDot: !filtered,
 			}
 		})
 	}, [sortedKeys, filter, dataKeys, colorMap])
@@ -134,7 +135,7 @@ export function TemperatureChart({
 	// 			label: `Test ${++i}`,
 	// 			dataKey: () => 0,
 	// 			color: "red",
-	// 			opacity: 1,
+	// 			strokeOpacity: 1,
 	// 		})
 	// 	}
 	// }
@@ -202,6 +203,7 @@ export function TemperatureChart({
 						return `${decimalString(value)} ${unit}`
 					}}
 					dataPoints={dataPoints}
+					filter={filter}
 				></LineChartDefault>
 			</ChartCard>
 		</div>