fsm: add WaitForStateAsync to the cached observer

By adding WaitForStateAsync to the observer we can always observe state
changes in an atomic way without relying on the observer's internal
cache.

Author: bhandras

fsm/example_fsm_test.go

@@ -140,71 +140,52 @@ func getTestContext() (*ExampleFSM, *CachedObserver) {
 	return exampleContext, cachedObserver
 }
 
-// TestExampleFSMFlow tests different flows that the example FSM can go through.
-func TestExampleFSMFlow(t *testing.T) {
+// TestObserverAsyncWait tests the observer's WaitForStateAsync function.
+func TestObserverAsyncWait(t *testing.T) {
 	testCases := []struct {
-		name              string
-		expectedStateFlow []StateType
-		expectedEventFlow []EventType
-		storeError        error
-		serviceError      error
+		name          string
+		waitTime      time.Duration
+		blockTime     time.Duration
+		expectTimeout bool
 	}{
 		{
-			name: "success",
-			expectedStateFlow: []StateType{
-				InitFSM,
-				StuffSentOut,
-				StuffSuccess,
-			},
-			expectedEventFlow: []EventType{
-				OnRequestStuff,
-				OnStuffSentOut,
-				OnStuffSuccess,
-			},
-		},
-		{
-			name: "failure on store",
-			expectedStateFlow: []StateType{
-				InitFSM,
-				StuffFailed,
-			},
-			expectedEventFlow: []EventType{
-				OnRequestStuff,
-				OnError,
-			},
-			storeError: errStore,
+			name:          "success",
+			waitTime:      time.Second,
+			blockTime:     time.Millisecond,
+			expectTimeout: false,
 		},
 		{
-			name: "failure on service",
-			expectedStateFlow: []StateType{
-				InitFSM,
-				StuffSentOut,
-				StuffFailed,
-			},
-			expectedEventFlow: []EventType{
-				OnRequestStuff,
-				OnStuffSentOut,
-				OnError,
-			},
-			serviceError: errService,
+			name:          "timeout",
+			waitTime:      time.Millisecond,
+			blockTime:     time.Second,
+			expectTimeout: true,
 		},
 	}
 
 	for _, tc := range testCases {
 		tc := tc
 
 		t.Run(tc.name, func(t *testing.T) {
-			exampleContext, cachedObserver := getTestContext()
-
-			if tc.storeError != nil {
-				exampleContext.store.(*mockStore).
-					storeErr = tc.storeError
+			service := &mockService{
+				respondChan: make(chan bool),
 			}
 
-			if tc.serviceError != nil {
-				exampleContext.service.(*mockService).
-					respondErr = tc.serviceError
-			}
+			store := &mockStore{}
+
+			exampleContext := NewExampleFSMContext(service, store)
+			cachedObserver := NewCachedObserver(100)
+			exampleContext.RegisterObserver(cachedObserver)
+
+			t0 := time.Now()
+			timeoutCtx, cancel := context.WithTimeout(
+				context.Background(), tc.waitTime,
+			)
+			defer cancel()
+
+			// Wait for the final state.
+			errChan := cachedObserver.WaitForStateAsync(
+				timeoutCtx, StuffSuccess, true,
+			)
 
 			go func() {
 				err := exampleContext.SendEvent(
@@ -213,32 +194,26 @@ func TestExampleFSMFlow(t *testing.T) {
 				)
 
 				require.NoError(t, err)
-			}()
 
-			// Wait for the final state.
-			err := cachedObserver.WaitForState(
-				context.Background(),
-				time.Second,
-				tc.expectedStateFlow[len(
-					tc.expectedStateFlow,
-				)-1],
-			)
-			require.NoError(t, err)
+				time.Sleep(tc.blockTime)
+				service.respondChan <- true
+			}()
 
-			allNotifications := cachedObserver.
-				GetCachedNotifications()
+			timeout := false
+			select {
+			case <-timeoutCtx.Done():
+				timeout = true
 
-			for index, notification := range allNotifications {
-				require.Equal(
-					t,
-					tc.expectedStateFlow[index],
-					notification.NextState,
-				)
-				require.Equal(
-					t,
-					tc.expectedEventFlow[index],
-					notification.Event,
-				)
+			case <-errChan:
+			}
+			require.Equal(t, tc.expectTimeout, timeout)
+
+			t1 := time.Now()
+			diff := t1.Sub(t0)
+			if tc.expectTimeout {
+				require.Less(t, diff, tc.blockTime)
+			} else {
+				require.Less(t, diff, tc.waitTime)
 			}
 		})
 	}

fsm/fsm.go

@@ -306,23 +306,36 @@ func NoOpAction(_ EventContext) EventType {
 }
 
 // ErrConfigError is an error returned when the state machine is misconfigured.
-type ErrConfigError error
+type ErrConfigError struct {
+	msg string
+}
+
+// Error returns the error message.
+func (e ErrConfigError) Error() string {
+	return fmt.Sprintf("config error: %s", e.msg)
+}
 
 // NewErrConfigError creates a new ErrConfigError.
 func NewErrConfigError(msg string) ErrConfigError {
-	return (ErrConfigError)(fmt.Errorf("config error: %s", msg))
+	return ErrConfigError{
+		msg: msg,
+	}
 }
 
 // ErrWaitingForStateTimeout is an error returned when the state machine times
 // out while waiting for a state.
-type ErrWaitingForStateTimeout error
+type ErrWaitingForStateTimeout struct {
+	expected StateType
+}
 
-// NewErrWaitingForStateTimeout creates a new ErrWaitingForStateTimeout.
-func NewErrWaitingForStateTimeout(expected,
-	actual StateType) ErrWaitingForStateTimeout {
+// Error returns the error message.
+func (e ErrWaitingForStateTimeout) Error() string {
+	return fmt.Sprintf("waiting for state timed out: %s", e.expected)
+}
 
-	return (ErrWaitingForStateTimeout)(fmt.Errorf(
-		"waiting for state timeout: expected %s, actual: %s",
-		expected, actual,
-	))
+// NewErrWaitingForStateTimeout creates a new ErrWaitingForStateTimeout.
+func NewErrWaitingForStateTimeout(expected StateType) ErrWaitingForStateTimeout {
+	return ErrWaitingForStateTimeout{
+		expected: expected,
+	}
 }

fsm/observer.go

@@ -105,7 +105,6 @@ func (s *CachedObserver) WaitForState(ctx context.Context,
 	opts ...WaitForStateOption) error {
 
 	var options fsmOptions
-
 	for _, opt := range opts {
 		opt.apply(&options)
 	}
@@ -120,61 +119,80 @@ func (s *CachedObserver) WaitForState(ctx context.Context,
 		}
 	}
 
+	// Create a new context with a timeout.
 	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
-	// Channel to notify when the desired state is reached
-	// or an error occurred.
+	ch := s.WaitForStateAsync(timeoutCtx, state, options.abortEarlyOnError)
+
+	// Wait for either the condition to be met or for a timeout.
+	select {
+	case <-timeoutCtx.Done():
+		return NewErrWaitingForStateTimeout(state)
+
+	case err := <-ch:
+		return err
+	}
+}
+
+// WaitForStateAsync waits asynchronously until the passed context is canceled
+// or the expected state is reached. The function returns a channel that will
+// receive an error if the expected state is reached or an error occurred. If
+// the context is canceled before the expected state is reached, the channel
+// will receive an ErrWaitingForStateTimeout error.
+func (s *CachedObserver) WaitForStateAsync(ctx context.Context, state StateType,
+	abortOnEarlyError bool) chan error {
+
+	// Channel to notify when the desired state is reached or an error
+	// occurred.
 	ch := make(chan error)
 
-	// Goroutine to wait on condition variable
+	// Wait on the notification condition variable asynchronously to avoid
+	// blocking the caller.
 	go func() {
 		s.notificationMx.Lock()
 		defer s.notificationMx.Unlock()
 
-		for {
-			// Check if the last state is the desired state
-			if s.lastNotification.NextState == state {
+		notifyAsync := func(err error) {
+			// To avoid potential deadlock when sending to the
+			// channel, while still holding the notificationMx
+			// lock, we use a goroutine to send the error to the
+			// channel.
+			go func() {
 				select {
-				case <-timeoutCtx.Done():
-					return
+				case <-ctx.Done():
+					ch <- NewErrWaitingForStateTimeout(
+						state,
+					)
 
-				case ch <- nil:
-					return
+				case ch <- err:
 				}
+			}()
+		}
+
+		for {
+			// Check if the last state is the desired state.
+			if s.lastNotification.NextState == state {
+				notifyAsync(nil)
+				return
 			}
 
-			// Check if an error occurred
+			// Check if an error has occurred.
 			if s.lastNotification.Event == OnError {
-				if options.abortEarlyOnError {
-					select {
-					case <-timeoutCtx.Done():
-						return
-
-					case ch <- s.lastNotification.LastActionError:
-						return
-					}
+				lastErr := s.lastNotification.LastActionError
+				if abortOnEarlyError {
+					notifyAsync(lastErr)
+					return
 				}
 			}
 
-			// Otherwise, wait for the next notification
+			// Otherwise use the conditonal variable to wait for
+			// the next notification.
 			s.notificationCond.Wait()
 		}
 	}()
 
-	// Wait for either the condition to be met or for a timeout
-	select {
-	case <-timeoutCtx.Done():
-		return NewErrWaitingForStateTimeout(
-			state, s.lastNotification.NextState,
-		)
-
-	case lastActionErr := <-ch:
-		if lastActionErr != nil {
-			return lastActionErr
-		}
-		return nil
-	}
+	return ch
 }
 
 // FixedSizeSlice is a slice with a fixed size.