Merge pull request #2815 from OffchainLabs/fix-stop-only

Fix StopOnly followed by StopAndWait

util/stopwaiter/stopwaiter.go

@@ -96,20 +96,12 @@ func (s *StopWaiterSafe) Start(ctx context.Context, parent any) error {
 }
 
 func (s *StopWaiterSafe) StopOnly() {
-	_ = s.stopOnly()
-}
-
-// returns true if stop function was called
-func (s *StopWaiterSafe) stopOnly() bool {
-	stopWasCalled := false
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 	if s.started && !s.stopped {
 		s.stopFunc()
-		stopWasCalled = true
 	}
 	s.stopped = true
-	return stopWasCalled
 }
 
 // StopAndWait may be called multiple times, even before start.
@@ -126,9 +118,15 @@ func getAllStackTraces() string {
 }
 
 func (s *StopWaiterSafe) stopAndWaitImpl(warningTimeout time.Duration) error {
-	if !s.stopOnly() {
+	s.StopOnly()
+	if !s.Started() {
+		// No need to wait, because nothing can be started if it's already stopped.
 		return nil
 	}
+	// Even if StopOnly has been previously called, make sure we wait for everything to shut down.
+	// Otherwise, a StopOnly call followed by StopAndWait might return early without waiting.
+	// At this point started must be true (because it was true above and cannot go back to false),
+	// so GetWaitChannel won't return an error.
 	waitChan, err := s.GetWaitChannel()
 	if err != nil {
 		return err

util/stopwaiter/stopwaiter_test.go

@@ -5,6 +5,7 @@ package stopwaiter
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -73,3 +74,19 @@ func TestStopWaiterStopAndWaitMultipleTimes(t *testing.T) {
 	sw.StopAndWait()
 	sw.StopAndWait()
 }
+
+func TestStopWaiterStopOnlyThenStopAndWait(t *testing.T) {
+	t.Parallel()
+	sw := StopWaiter{}
+	sw.Start(context.Background(), &TestStruct{})
+	var threadStopping atomic.Bool
+	sw.LaunchThread(func(context.Context) {
+		time.Sleep(time.Second)
+		threadStopping.Store(true)
+	})
+	sw.StopOnly()
+	sw.StopAndWait()
+	if !threadStopping.Load() {
+		t.Error("StopAndWait returned before background thread stopped")
+	}
+}