mirror of
https://github.com/henrygd/beszel.git
synced 2026-08-19 08:47:46 +02:00
fix(tests): stop system updaters on app termination
This commit is contained in:
@@ -56,7 +56,7 @@ func (sm *SystemManager) NewSystem(systemId string) *System {
|
|||||||
Id: systemId,
|
Id: systemId,
|
||||||
data: &system.CombinedData{},
|
data: &system.CombinedData{},
|
||||||
}
|
}
|
||||||
system.ctx, system.cancel = system.getContext()
|
system.ctx, system.cancel = system.getContext(sm.ctx)
|
||||||
return system
|
return system
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -79,7 +79,10 @@ func (sys *System) StartUpdater() {
|
|||||||
} else {
|
} else {
|
||||||
// if the system does not have a websocket connection, wait before updating
|
// 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).
|
// 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)
|
// 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)
|
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 {
|
if sys.ctx == nil {
|
||||||
sys.ctx, sys.cancel = context.WithCancel(context.Background())
|
sys.ctx, sys.cancel = context.WithCancel(ctx)
|
||||||
}
|
}
|
||||||
return sys.ctx, sys.cancel
|
return sys.ctx, sys.cancel
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package systems
|
package systems
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"time"
|
"time"
|
||||||
@@ -45,6 +46,8 @@ type SystemManager struct {
|
|||||||
systems *store.Store[string, *System] // Thread-safe store of active systems
|
systems *store.Store[string, *System] // Thread-safe store of active systems
|
||||||
sshConfig *ssh.ClientConfig // SSH client configuration for system connections
|
sshConfig *ssh.ClientConfig // SSH client configuration for system connections
|
||||||
smartFetchMap *expirymap.ExpiryMap[smartFetchState] // Stores last SMART fetch time/result; TTL is only for cleanup
|
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.
|
// 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.
|
// NewSystemManager creates a new SystemManager instance with the provided hub.
|
||||||
// The hub must implement the hubLike interface to provide database and alert functionality.
|
// The hub must implement the hubLike interface to provide database and alert functionality.
|
||||||
func NewSystemManager(hub hubLike) *SystemManager {
|
func NewSystemManager(hub hubLike) *SystemManager {
|
||||||
return &SystemManager{
|
sm := &SystemManager{
|
||||||
systems: store.New(map[string]*System{}),
|
systems: store.New(map[string]*System{}),
|
||||||
hub: hub,
|
hub: hub,
|
||||||
smartFetchMap: expirymap.New[smartFetchState](time.Hour),
|
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
|
// GetSystem returns a system by ID from the store
|
||||||
@@ -103,7 +108,9 @@ func (sm *SystemManager) Initialize() error {
|
|||||||
sleepTime := time.Duration(delta) * time.Millisecond
|
sleepTime := time.Duration(delta) * time.Millisecond
|
||||||
|
|
||||||
for _, system := range systems {
|
for _, system := range systems {
|
||||||
time.Sleep(sleepTime)
|
if !waitForContext(sm.ctx, sleepTime) {
|
||||||
|
return
|
||||||
|
}
|
||||||
_ = sm.AddSystem(system)
|
_ = sm.AddSystem(system)
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
@@ -121,6 +128,13 @@ func (sm *SystemManager) bindEventHooks() {
|
|||||||
sm.hub.OnRecordAfterUpdateSuccess("fingerprints").BindFunc(sm.onTokenRotated)
|
sm.hub.OnRecordAfterUpdateSuccess("fingerprints").BindFunc(sm.onTokenRotated)
|
||||||
sm.hub.OnRealtimeSubscribeRequest().BindFunc(sm.onRealtimeSubscribeRequest)
|
sm.hub.OnRealtimeSubscribeRequest().BindFunc(sm.onRealtimeSubscribeRequest)
|
||||||
sm.hub.OnRealtimeConnectRequest().BindFunc(sm.onRealtimeConnectRequest)
|
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.
|
// onTokenRotated handles fingerprint token rotation events.
|
||||||
@@ -247,7 +261,7 @@ func (sm *SystemManager) AddSystem(sys *System) error {
|
|||||||
|
|
||||||
// Initialize system for monitoring
|
// Initialize system for monitoring
|
||||||
sys.manager = sm
|
sys.manager = sm
|
||||||
sys.ctx, sys.cancel = sys.getContext()
|
sys.ctx, sys.cancel = sys.getContext(sm.ctx)
|
||||||
sys.data = &system.CombinedData{}
|
sys.data = &system.CombinedData{}
|
||||||
sm.systems.Set(sys.Id, sys)
|
sm.systems.Set(sys.Id, sys)
|
||||||
|
|
||||||
@@ -372,3 +386,15 @@ func deactivateAlerts(app core.App, systemID string) error {
|
|||||||
}
|
}
|
||||||
return nil
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ func TestSystemManagerNew(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
synctest.Test(t, func(t *testing.T) {
|
synctest.Test(t, func(t *testing.T) {
|
||||||
|
sm.ResetContextForTesting()
|
||||||
sm.Initialize()
|
sm.Initialize()
|
||||||
|
|
||||||
record, err := tests.CreateRecord(hub, "systems", map[string]any{
|
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")
|
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)
|
testOld(t, hub)
|
||||||
|
|
||||||
synctest.Test(t, func(t *testing.T) {
|
synctest.Test(t, func(t *testing.T) {
|
||||||
|
|||||||
@@ -117,6 +117,11 @@ func (sm *SystemManager) RemoveAllSystems() {
|
|||||||
sm.smartFetchMap.StopCleaner()
|
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() {
|
func (s *System) StopUpdater() {
|
||||||
s.cancel()
|
s.cancel()
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user