pkg/utils: deprecate StartStopOnce in favor of StateMachine (#203)

Author: jmank88

pkg/loop/internal/plugin_service.go

@@ -28,7 +28,7 @@ type grpcPlugin interface {
 // pluginService is a [types.Service] wrapper that maintains an internal [types.Service] created from a [grpcPlugin]
 // client instance by launching and re-launching as necessary.
 type PluginService[P grpcPlugin, S services.Service] struct {
-	utils.StartStopOnce
+	services.StateMachine
 
 	pluginName string
 

pkg/services/state.go

@@ -13,7 +13,7 @@ import (
 const defaultErrorBufferCap = 50
 
 type errNotStarted struct {
-	state State
+	state state
 }
 
 func (e *errNotStarted) Error() string {
@@ -25,7 +25,7 @@ var (
 	ErrCannotStopUnstarted = errors.New("cannot stop unstarted service")
 )
 
-// StateMachine contains a State integer
+// StateMachine contains a state integer
 type StateMachine struct {
 	state        atomic.Int32
 	sync.RWMutex // lock is held during startup/shutdown, RLock is held while executing functions dependent on a particular state
@@ -38,12 +38,12 @@ type StateMachine struct {
 	SvcErrBuffer ErrorBuffer
 }
 
-// State holds the state for StateMachine
-type State int32
+// state holds the state for StateMachine
+type state int32
 
 // nolint
 const (
-	stateUnstarted State = iota
+	stateUnstarted state = iota
 	stateStarted
 	stateStarting
 	stateStartFailed
@@ -52,7 +52,7 @@ const (
 	stateStopFailed
 )
 
-func (s State) String() string {
+func (s state) String() string {
 	switch s {
 	case stateUnstarted:
 		return "Unstarted"
@@ -80,7 +80,7 @@ func (once *StateMachine) StartOnce(name string, fn func() error) error {
 	success := once.state.CompareAndSwap(int32(stateUnstarted), int32(stateStarting))
 
 	if !success {
-		return pkgerrors.Errorf("%v has already been started once; state=%v", name, State(once.state.Load()))
+		return pkgerrors.Errorf("%v has already been started once; state=%v", name, state(once.state.Load()))
 	}
 
 	once.Lock()
@@ -116,14 +116,14 @@ func (once *StateMachine) StopOnce(name string, fn func() error) error {
 	success := once.state.CompareAndSwap(int32(stateStarted), int32(stateStopping))
 
 	if !success {
-		state := once.state.Load()
-		switch state {
-		case int32(stateStopped):
+		s := once.loadState()
+		switch s {
+		case stateStopped:
 			return pkgerrors.Wrapf(ErrAlreadyStopped, "%s has already been stopped", name)
-		case int32(stateUnstarted):
+		case stateUnstarted:
 			return pkgerrors.Wrapf(ErrCannotStopUnstarted, "%s has not been started", name)
 		default:
-			return pkgerrors.Errorf("%v cannot be stopped from this state; state=%v", name, State(state))
+			return pkgerrors.Errorf("%v cannot be stopped from this state; state=%v", name, s)
 		}
 	}
 
@@ -145,19 +145,20 @@ func (once *StateMachine) StopOnce(name string, fn func() error) error {
 }
 
 // State retrieves the current state
-func (once *StateMachine) State() State {
-	state := once.state.Load()
-	return State(state)
+func (once *StateMachine) State() string {
+	return once.loadState().String()
+}
+
+func (once *StateMachine) loadState() state {
+	return state(once.state.Load())
 }
 
 // IfStarted runs the func and returns true only if started, otherwise returns false
 func (once *StateMachine) IfStarted(f func()) (ok bool) {
 	once.RLock()
 	defer once.RUnlock()
 
-	state := once.state.Load()
-
-	if State(state) == stateStarted {
+	if once.loadState() == stateStarted {
 		f()
 		return true
 	}
@@ -169,9 +170,7 @@ func (once *StateMachine) IfNotStopped(f func()) (ok bool) {
 	once.RLock()
 	defer once.RUnlock()
 
-	state := once.state.Load()
-
-	if State(state) == stateStopped {
+	if once.loadState() == stateStopped {
 		return false
 	}
 	f()
@@ -180,7 +179,7 @@ func (once *StateMachine) IfNotStopped(f func()) (ok bool) {
 
 // Ready returns ErrNotStarted if the state is not started.
 func (once *StateMachine) Ready() error {
-	state := once.State()
+	state := once.loadState()
 	if state == stateStarted {
 		return nil
 	}
@@ -190,7 +189,7 @@ func (once *StateMachine) Ready() error {
 // Healthy returns ErrNotStarted if the state is not started.
 // Override this per-service with more specific implementations.
 func (once *StateMachine) Healthy() error {
-	state := once.State()
+	state := once.loadState()
 	if state == stateStarted {
 		return once.SvcErrBuffer.Flush()
 	}

pkg/services/state_test.go

@@ -0,0 +1,178 @@
+package services
+
+import (
+	"errors"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestStateMachine_StartOnce_StopOnce(t *testing.T) {
+	t.Parallel()
+
+	var sm StateMachine
+
+	ch := make(chan int, 3)
+
+	ready := make(chan bool)
+
+	go func() {
+		assert.NoError(t, sm.StartOnce("slow service", func() (err error) {
+			ch <- 1
+			ready <- true
+			<-time.After(time.Millisecond * 500) // wait for StopOnce to happen
+			ch <- 2
+
+			return nil
+		}))
+	}()
+
+	go func() {
+		<-ready // try stopping halfway through startup
+		assert.NoError(t, sm.StopOnce("slow service", func() (err error) {
+			ch <- 3
+
+			return nil
+		}))
+	}()
+
+	require.Equal(t, 1, <-ch)
+	require.Equal(t, 2, <-ch)
+	require.Equal(t, 3, <-ch)
+}
+
+func TestStateMachine_MultipleStartNoBlock(t *testing.T) {
+	t.Parallel()
+
+	var sm StateMachine
+
+	ch := make(chan int, 3)
+
+	ready := make(chan bool)
+	next := make(chan bool)
+
+	go func() {
+		ch <- 1
+		assert.NoError(t, sm.StartOnce("slow service", func() (err error) {
+			ready <- true
+			<-next // continue after the other StartOnce call fails
+
+			return nil
+		}))
+		<-next
+		ch <- 2
+
+	}()
+
+	go func() {
+		<-ready // try starting halfway through startup
+		assert.Error(t, sm.StartOnce("slow service", func() (err error) {
+			return nil
+		}))
+		next <- true
+		ch <- 3
+		next <- true
+
+	}()
+
+	require.Equal(t, 1, <-ch)
+	require.Equal(t, 3, <-ch) // 3 arrives before 2 because it returns immediately
+	require.Equal(t, 2, <-ch)
+}
+
+func TestStateMachine(t *testing.T) {
+	t.Parallel()
+
+	var callsCount atomic.Int32
+	incCount := func() {
+		callsCount.Add(1)
+	}
+
+	var s StateMachine
+	ok := s.IfStarted(incCount)
+	assert.False(t, ok)
+	ok = s.IfNotStopped(incCount)
+	assert.True(t, ok)
+	assert.Equal(t, int32(1), callsCount.Load())
+
+	err := s.StartOnce("foo", func() error { return nil })
+	assert.NoError(t, err)
+
+	assert.True(t, s.IfStarted(incCount))
+	assert.Equal(t, int32(2), callsCount.Load())
+
+	err = s.StopOnce("foo", func() error { return nil })
+	assert.NoError(t, err)
+	ok = s.IfNotStopped(incCount)
+	assert.False(t, ok)
+	assert.Equal(t, int32(2), callsCount.Load())
+}
+
+func TestStateMachine_StartErrors(t *testing.T) {
+	var s StateMachine
+
+	err := s.StartOnce("foo", func() error { return errors.New("foo") })
+	assert.Error(t, err)
+
+	var callsCount atomic.Int32
+	incCount := func() {
+		callsCount.Add(1)
+	}
+
+	assert.False(t, s.IfStarted(incCount))
+	assert.Equal(t, int32(0), callsCount.Load())
+
+	err = s.StartOnce("foo", func() error { return nil })
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "foo has already been started once")
+	err = s.StopOnce("foo", func() error { return nil })
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "foo cannot be stopped from this state; state=StartFailed")
+
+	assert.Equal(t, stateStartFailed, s.loadState())
+}
+
+func TestStateMachine_StopErrors(t *testing.T) {
+	var s StateMachine
+
+	err := s.StartOnce("foo", func() error { return nil })
+	require.NoError(t, err)
+
+	var callsCount atomic.Int32
+	incCount := func() {
+		callsCount.Add(1)
+	}
+
+	err = s.StopOnce("foo", func() error { return errors.New("explodey mcsplode") })
+	assert.Error(t, err)
+
+	assert.False(t, s.IfStarted(incCount))
+	assert.Equal(t, int32(0), callsCount.Load())
+	assert.True(t, s.IfNotStopped(incCount))
+	assert.Equal(t, int32(1), callsCount.Load())
+
+	err = s.StartOnce("foo", func() error { return nil })
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "foo has already been started once")
+	err = s.StopOnce("foo", func() error { return nil })
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "foo cannot be stopped from this state; state=StopFailed")
+
+	assert.Equal(t, stateStopFailed, s.loadState())
+}
+
+func TestStateMachine_Ready_Healthy(t *testing.T) {
+	t.Parallel()
+
+	var s StateMachine
+	assert.Error(t, s.Ready())
+	assert.Error(t, s.Healthy())
+
+	err := s.StartOnce("foo", func() error { return nil })
+	assert.NoError(t, err)
+	assert.NoError(t, s.Ready())
+	assert.NoError(t, s.Healthy())
+}