From e668a855b71a8311a3f1862326bcbbf0223dd35b Mon Sep 17 00:00:00 2001
From: Emin <9283660+eminden@users.noreply.github.com>
Date: Thu, 2 Feb 2023 05:59:09 +0000
Subject: [PATCH] race fixes (#94)

* race fixes

* race flag

* ci race flag

* update transition within state mutex, this prevents races at Can calls

---------

Co-authored-by: Emin Tasgetiren <emin.tasgetiren@abex.capital>
---
 .github/workflows/main.yml |  2 +-
 Makefile                   |  2 +-
 fsm.go                     |  3 +--
 fsm_test.go                | 19 ++++++-------------
 4 files changed, 9 insertions(+), 17 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 9967359..afb9ff4 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -18,7 +18,7 @@ jobs:
           go-version: 1.16
 
       - name: Test
-        run: go test -coverprofile=coverage.out ./...
+        run: go test -race -coverprofile=coverage.out ./...
 
       - name: Convert coverage
         uses: jandelgado/gcov2lcov-action@v1.0.5
diff --git a/Makefile b/Makefile
index 70e816f..ffd4820 100644
--- a/Makefile
+++ b/Makefile
@@ -2,7 +2,7 @@ default: services test
 
 .PHONY: test
 test:
-	go test ./...
+	go test -race ./...
 
 .PHONY: lint
 lint:
diff --git a/fsm.go b/fsm.go
index e8662f6..bd269d7 100644
--- a/fsm.go
+++ b/fsm.go
@@ -349,6 +349,7 @@ func (f *FSM) Event(ctx context.Context, event string, args ...interface{}) erro
 
 			f.stateMu.Lock()
 			f.current = dst
+			f.transition = nil // treat the state transition as done
 			f.stateMu.Unlock()
 
 			// at this point, we unlock the event mutex in order to allow
@@ -359,7 +360,6 @@ func (f *FSM) Event(ctx context.Context, event string, args ...interface{}) erro
 				f.eventMu.Unlock()
 				unlocked = true
 			}
-			f.transition = nil // treat the state transition as done
 			f.enterStateCallbacks(ctx, e)
 			f.afterEventCallbacks(ctx, e)
 		}
@@ -420,7 +420,6 @@ func (t transitionerStruct) transition(f *FSM) error {
 		return NotInTransitionError{}
 	}
 	f.transition()
-	f.transition = nil
 	return nil
 }
 
diff --git a/fsm_test.go b/fsm_test.go
index 985b0ea..8e0c4ed 100644
--- a/fsm_test.go
+++ b/fsm_test.go
@@ -544,20 +544,17 @@ func TestCancelAsyncTransition(t *testing.T) {
 	if !ok {
 		t.Errorf("expected error to be 'AsyncError', got %v", err)
 	}
-	var asyncStateTransitionWasCanceled bool
+	var asyncStateTransitionWasCanceled = make(chan struct{})
 	go func() {
 		<-asyncError.Ctx.Done()
-		asyncStateTransitionWasCanceled = true
+		close(asyncStateTransitionWasCanceled)
 	}()
 	asyncError.CancelTransition()
-	time.Sleep(20 * time.Millisecond)
+	<-asyncStateTransitionWasCanceled
 
 	if err = fsm.Transition(); err != nil {
 		t.Errorf("expected no error, got %v", err)
 	}
-	if !asyncStateTransitionWasCanceled {
-		t.Error("expected async state transition cancelation to have propagated")
-	}
 	if fsm.Current() != "start" {
 		t.Error("expected state to be 'start'")
 	}
@@ -775,7 +772,7 @@ func TestTransitionInCallbacks(t *testing.T) {
 
 func TestContextInCallbacks(t *testing.T) {
 	var fsm *FSM
-	var enterEndAsyncWorkDone bool
+	var enterEndAsyncWorkDone = make(chan struct{})
 	fsm = NewFSM(
 		"start",
 		Events{
@@ -787,7 +784,7 @@ func TestContextInCallbacks(t *testing.T) {
 			"enter_end": func(ctx context.Context, e *Event) {
 				go func() {
 					<-ctx.Done()
-					enterEndAsyncWorkDone = true
+					close(enterEndAsyncWorkDone)
 				}()
 
 				<-ctx.Done()
@@ -806,11 +803,7 @@ func TestContextInCallbacks(t *testing.T) {
 	if !errors.Is(err, context.Canceled) {
 		t.Errorf("expected 'context canceled' error, got %v", err)
 	}
-	time.Sleep(20 * time.Millisecond)
-
-	if !enterEndAsyncWorkDone {
-		t.Error("expected asynchronous work in callback to be done but it wasn't")
-	}
+	<-enterEndAsyncWorkDone
 
 	currentState := fsm.Current()
 	if currentState != "end" {