diff --git a/cmd/vGPUmonitor/feedback.go b/cmd/vGPUmonitor/feedback.go index 3d119e6c4c..cb86a6960b 100644 --- a/cmd/vGPUmonitor/feedback.go +++ b/cmd/vGPUmonitor/feedback.go @@ -71,7 +71,15 @@ func CheckPriority(utSwitchOn map[string]UtilizationPerDevice, p int, c *nvidia. return false } +var observeTestHook func() + func Observe(lister *nvidia.ContainerLister) { + if observeTestHook != nil { + observeTestHook() + } + lister.Lock() + defer lister.UnLock() + utSwitchOn := map[string]UtilizationPerDevice{} containers := lister.ListContainers() diff --git a/cmd/vGPUmonitor/feedback_test.go b/cmd/vGPUmonitor/feedback_test.go index 10c2177790..62074c8cbd 100644 --- a/cmd/vGPUmonitor/feedback_test.go +++ b/cmd/vGPUmonitor/feedback_test.go @@ -18,6 +18,7 @@ package main import ( "testing" + "time" "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia" ) @@ -160,3 +161,46 @@ func TestCheckBlocking_MultiDevice(t *testing.T) { }) } } + +func TestObserve(t *testing.T) { + // Call Observe with an empty lister to cover the missing lines for codecov + // and to ensure no panics occur with locking/unlocking. + lister := &nvidia.ContainerLister{} + + // Test that Observe actually acquires the lock deterministically. + lister.Lock() + + reachedLock := make(chan struct{}) + observeTestHook = func() { + close(reachedLock) + } + defer func() { observeTestHook = nil }() + + done := make(chan struct{}) + go func() { + Observe(lister) + close(done) + }() + + // Wait until Observe reaches the lock boundary + <-reachedLock + + // Ensure Observe is now blocked on the lock + select { + case <-done: + t.Fatal("Observe completed while lock was held by another goroutine, indicating it did not acquire the lock!") + default: + // Expected: Observe is blocked + } + + // Release the lock + lister.UnLock() + + // Now Observe should complete + select { + case <-done: + // Success + case <-time.After(1 * time.Second): + t.Fatal("Observe did not complete after lock was released") + } +}