From 1403c8874a47c3b534fd24def6f935eafc848d31 Mon Sep 17 00:00:00 2001
From: Lee Bousfield <ljbousfield@gmail.com>
Date: Tue, 3 Dec 2024 10:28:04 -0700
Subject: [PATCH 1/2] Fix StopOnly followed by StopAndWait

---
 util/stopwaiter/stopwaiter.go | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/util/stopwaiter/stopwaiter.go b/util/stopwaiter/stopwaiter.go
index 993768dd85..974178e1cf 100644
--- a/util/stopwaiter/stopwaiter.go
+++ b/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,9 @@ func getAllStackTraces() string {
 }
 
 func (s *StopWaiterSafe) stopAndWaitImpl(warningTimeout time.Duration) error {
-	if !s.stopOnly() {
-		return nil
-	}
+	s.StopOnly()
+	// 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.
 	waitChan, err := s.GetWaitChannel()
 	if err != nil {
 		return err

From 99830437ba13b1bab4d0a83fb700f1d9bd0e8a84 Mon Sep 17 00:00:00 2001
From: Lee Bousfield <ljbousfield@gmail.com>
Date: Tue, 3 Dec 2024 12:14:34 -0700
Subject: [PATCH 2/2] Fix stop before start and add test

---
 util/stopwaiter/stopwaiter.go      |  6 ++++++
 util/stopwaiter/stopwaiter_test.go | 17 +++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/util/stopwaiter/stopwaiter.go b/util/stopwaiter/stopwaiter.go
index 974178e1cf..c242ac26ab 100644
--- a/util/stopwaiter/stopwaiter.go
+++ b/util/stopwaiter/stopwaiter.go
@@ -119,8 +119,14 @@ func getAllStackTraces() string {
 
 func (s *StopWaiterSafe) stopAndWaitImpl(warningTimeout time.Duration) error {
 	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
diff --git a/util/stopwaiter/stopwaiter_test.go b/util/stopwaiter/stopwaiter_test.go
index c561e1f43b..68e49ac2be 100644
--- a/util/stopwaiter/stopwaiter_test.go
+++ b/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")
+	}
+}