fix race in parallel tests stretchr/testify#1165

SaveTheRbtz · Apr 11, 2022 · 1b20c8a · 1b20c8a
1 parent 7158c70
commit 1b20c8a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/module/lifecycle/lifecycle.go b/module/lifecycle/lifecycle.go
@@ -49,17 +49,22 @@ func (lm *LifecycleManager) OnStart(startupFns ...func()) {
 // starting up, we will wait for startup to complete before shutting down. After the first
 // call, subsequent calls to OnStop do nothing.
 func (lm *LifecycleManager) OnStop(shutdownFns ...func()) {
+	var waitForStartup bool
+
 	lm.stateTransition.Lock()
 	if lm.shutdownCommenced {
 		lm.stateTransition.Unlock()
 		return
 	}
 	lm.shutdownCommenced = true
+	if lm.startupCommenced {
+		waitForStartup = true
+	}
 	lm.stateTransition.Unlock()
 
 	close(lm.shutdownSignal)
 	go func() {
-		if lm.startupCommenced {
+		if waitForStartup {
 			<-lm.started
 			for _, fn := range shutdownFns {
 				fn()

diff --git a/module/lifecycle/lifecycle_test.go b/module/lifecycle/lifecycle_test.go
@@ -25,9 +25,10 @@ func (suite *LifecycleManagerSuite) SetupTest() {
 func (suite *LifecycleManagerSuite) TestConcurrentStart() {
 	var numStarts uint32
 
+	lm := suite.lm
 	for i := 0; i < 10; i++ {
 		go func() {
-			suite.lm.OnStart(func() {
+			lm.OnStart(func() {
 				atomic.AddUint32(&numStarts, 1)
 			})
 		}()
@@ -46,9 +47,10 @@ func (suite *LifecycleManagerSuite) TestConcurrentStop() {
 
 	var numStops uint32
 
+	lm := suite.lm
 	for i := 0; i < 10; i++ {
 		go func() {
-			suite.lm.OnStop(func() {
+			lm.OnStop(func() {
 				atomic.AddUint32(&numStops, 1)
 			})
 		}()