From 1b20c8a8720ecf84637fa881274abae771d196ca Mon Sep 17 00:00:00 2001
From: Alexey Ivanov <rbtz@ph34r.me>
Date: Mon, 11 Apr 2022 03:03:16 -0700
Subject: [PATCH] fix race in parallel tests
 https://github.com/stretchr/testify/pull/1165

---
 module/lifecycle/lifecycle.go      | 7 ++++++-
 module/lifecycle/lifecycle_test.go | 6 ++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/module/lifecycle/lifecycle.go b/module/lifecycle/lifecycle.go
index be10638b2b4..16d6f7593d4 100644
--- 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
index c88207b31c0..ed3707960a2 100644
--- 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)
 			})
 		}()