diff --git a/core/application/watchdog.go b/core/application/watchdog.go index c71871d87..330c95353 100644 --- a/core/application/watchdog.go +++ b/core/application/watchdog.go @@ -61,7 +61,15 @@ func extractModelGroupsFromConfigs(configs []config.ModelConfig) map[string][]st return out } +// StopWatchdog signals the watchdog to stop. It takes the watchdogMutex because +// startWatchdog/RestartWatchdog reassign and close a.watchdogStop under that same +// lock, and POST /api/settings can reach either this or RestartWatchdog depending +// on WatchdogShouldRun(). Without the lock, two callers can both observe a non-nil +// channel and close it twice, which panics. func (a *Application) StopWatchdog() error { + a.watchdogMutex.Lock() + defer a.watchdogMutex.Unlock() + if a.watchdogStop != nil { close(a.watchdogStop) a.watchdogStop = nil diff --git a/core/application/watchdog_test.go b/core/application/watchdog_test.go index ddb1a6849..09e856e10 100644 --- a/core/application/watchdog_test.go +++ b/core/application/watchdog_test.go @@ -1,6 +1,8 @@ package application import ( + "sync" + "github.com/mudler/LocalAI/core/config" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -45,3 +47,32 @@ var _ = Describe("extractModelGroupsFromConfigs", func() { Expect(out).ToNot(HaveKey("a")) }) }) + +var _ = Describe("StopWatchdog", func() { + It("closes the stop channel exactly once under concurrent callers", func() { + // POST /api/settings routes to StopWatchdog whenever the new settings make + // WatchdogShouldRun() false, so concurrent requests land here in parallel. + // Before the mutex was taken, both could observe a non-nil watchdogStop and + // close it twice, panicking with "close of closed channel". + app := &Application{watchdogStop: make(chan bool, 1)} + + var wg sync.WaitGroup + for i := 0; i < 64; i++ { + wg.Add(1) + go func() { + defer GinkgoRecover() + defer wg.Done() + Expect(app.StopWatchdog()).To(Succeed()) + }() + } + wg.Wait() + + Expect(app.watchdogStop).To(BeNil()) + }) + + It("is a no-op when the watchdog was never started", func() { + app := &Application{} + Expect(app.StopWatchdog()).To(Succeed()) + Expect(app.watchdogStop).To(BeNil()) + }) +})