From 5f383c0eb18633b0a7962a2a566b5a8db81984fe Mon Sep 17 00:00:00 2001 From: henrygd Date: Mon, 17 Aug 2026 12:08:21 -0400 Subject: [PATCH] fix(tests): stop system updaters on app termination --- internal/hub/systems/system.go | 11 ++++--- internal/hub/systems/system_manager.go | 32 ++++++++++++++++++-- internal/hub/systems/systems_test.go | 3 ++ internal/hub/systems/systems_test_helpers.go | 5 +++ 4 files changed, 44 insertions(+), 7 deletions(-) diff --git a/internal/hub/systems/system.go b/internal/hub/systems/system.go index ceec0ad6..c04c2b78 100644 --- a/internal/hub/systems/system.go +++ b/internal/hub/systems/system.go @@ -56,7 +56,7 @@ func (sm *SystemManager) NewSystem(systemId string) *System { Id: systemId, data: &system.CombinedData{}, } - system.ctx, system.cancel = system.getContext() + system.ctx, system.cancel = system.getContext(sm.ctx) return system } @@ -79,7 +79,10 @@ func (sys *System) StartUpdater() { } else { // if the system does not have a websocket connection, wait before updating // to allow the agent to connect via websocket (makes sure fingerprint is set). - time.Sleep(11 * time.Second) + if !waitForContext(sys.ctx, 11*time.Second) { + return + } + } // update immediately if system is not paused (only for ws connections) @@ -402,9 +405,9 @@ func (sys *System) setDown(originalError error) error { return sys.manager.hub.SaveNoValidate(record) } -func (sys *System) getContext() (context.Context, context.CancelFunc) { +func (sys *System) getContext(ctx context.Context) (context.Context, context.CancelFunc) { if sys.ctx == nil { - sys.ctx, sys.cancel = context.WithCancel(context.Background()) + sys.ctx, sys.cancel = context.WithCancel(ctx) } return sys.ctx, sys.cancel } diff --git a/internal/hub/systems/system_manager.go b/internal/hub/systems/system_manager.go index 32112f2c..42f9562d 100644 --- a/internal/hub/systems/system_manager.go +++ b/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 app terminates + cancel context.CancelFunc // Cancels ctx and all child system contexts } // hubLike defines the interface requirements for the hub dependency. @@ -60,11 +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 { - return &SystemManager{ + sm := &SystemManager{ systems: store.New(map[string]*System{}), hub: hub, smartFetchMap: expirymap.New[smartFetchState](time.Hour), } + sm.ctx, sm.cancel = context.WithCancel(context.Background()) + return sm } // GetSystem returns a system by ID from the store @@ -103,7 +108,9 @@ func (sm *SystemManager) Initialize() error { sleepTime := time.Duration(delta) * time.Millisecond for _, system := range systems { - time.Sleep(sleepTime) + if !waitForContext(sm.ctx, sleepTime) { + return + } _ = sm.AddSystem(system) } }() @@ -121,6 +128,13 @@ func (sm *SystemManager) bindEventHooks() { sm.hub.OnRecordAfterUpdateSuccess("fingerprints").BindFunc(sm.onTokenRotated) sm.hub.OnRealtimeSubscribeRequest().BindFunc(sm.onRealtimeSubscribeRequest) sm.hub.OnRealtimeConnectRequest().BindFunc(sm.onRealtimeConnectRequest) + sm.hub.OnTerminate().BindFunc(sm.onTerminate) +} + +// onTerminate cancels SystemManager context on app shutdown +func (sm *SystemManager) onTerminate(e *core.TerminateEvent) error { + sm.cancel() + return e.Next() } // onTokenRotated handles fingerprint token rotation events. @@ -247,7 +261,7 @@ func (sm *SystemManager) AddSystem(sys *System) error { // Initialize system for monitoring sys.manager = sm - sys.ctx, sys.cancel = sys.getContext() + sys.ctx, sys.cancel = sys.getContext(sm.ctx) sys.data = &system.CombinedData{} sm.systems.Set(sys.Id, sys) @@ -372,3 +386,15 @@ func deactivateAlerts(app core.App, systemID string) error { } return nil } + +// waitForContext waits for delay or returns early when ctx is cancelled. +func waitForContext(ctx context.Context, delay time.Duration) bool { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} diff --git a/internal/hub/systems/systems_test.go b/internal/hub/systems/systems_test.go index e9632c6c..cf4997f8 100644 --- a/internal/hub/systems/systems_test.go +++ b/internal/hub/systems/systems_test.go @@ -30,6 +30,7 @@ func TestSystemManagerNew(t *testing.T) { require.NoError(t, err) synctest.Test(t, func(t *testing.T) { + sm.ResetContextForTesting() sm.Initialize() record, err := tests.CreateRecord(hub, "systems", map[string]any{ @@ -112,6 +113,8 @@ func TestSystemManagerNew(t *testing.T) { assert.False(t, sm.HasSystem(record.Id), "System should not exist in the store after deletion") }) + // The following subtests run outside the synctest bubble. + sm.ResetContextForTesting() testOld(t, hub) synctest.Test(t, func(t *testing.T) { diff --git a/internal/hub/systems/systems_test_helpers.go b/internal/hub/systems/systems_test_helpers.go index 5f598d7d..b6ae4782 100644 --- a/internal/hub/systems/systems_test_helpers.go +++ b/internal/hub/systems/systems_test_helpers.go @@ -117,6 +117,11 @@ func (sm *SystemManager) RemoveAllSystems() { sm.smartFetchMap.StopCleaner() } +// ResetContextForTesting replaces the manager context for a new synctest bubble. +func (sm *SystemManager) ResetContextForTesting() { + sm.ctx, sm.cancel = context.WithCancel(context.Background()) +} + func (s *System) StopUpdater() { s.cancel() }