micro/fsm/default.go

package fsm

import (
	"context"
	"fmt"
	"sync"
)

type state struct {
	body interface{}
	name string
}

var _ State = &state{}

func (s *state) Name() string {
	return s.name
}

func (s *state) Body() interface{} {
	return s.body
}

// fsm is a finite state machine
type fsm struct {
	statesMap   map[string]StateFunc
	current     string
	statesOrder []string
	opts        Options
	mu          sync.Mutex
}

// NewFSM creates a new finite state machine having the specified initial state
// with specified options
func NewFSM(opts ...Option) FSM {
	return &fsm{
		statesMap: map[string]StateFunc{},
		opts:      NewOptions(opts...),
	}
}

// Current returns the current state
func (f *fsm) Current() string {
	f.mu.Lock()
	s := f.current
	f.mu.Unlock()
	return s
}

// Current returns the current state
func (f *fsm) Reset() {
	f.mu.Lock()
	f.current = f.opts.Initial
	f.mu.Unlock()
}

// State adds state to fsm
func (f *fsm) State(state string, fn StateFunc) {
	f.mu.Lock()
	f.statesMap[state] = fn
	f.statesOrder = append(f.statesOrder, state)
	f.mu.Unlock()
}

// Start runs state machine with provided data
func (f *fsm) Start(ctx context.Context, args interface{}, opts ...Option) (interface{}, error) {
	var err error

	f.mu.Lock()
	options := f.opts

	for _, opt := range opts {
		opt(&options)
	}

	sopts := []StateOption{StateDryRun(options.DryRun)}

	cstate := options.Initial
	states := make(map[string]StateFunc, len(f.statesMap))
	for k, v := range f.statesMap {
		states[k] = v
	}
	f.current = cstate
	f.mu.Unlock()

	var s State
	s = &state{name: cstate, body: args}
	nstate := s.Name()

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			fn, ok := states[nstate]
			if !ok {
				return nil, fmt.Errorf(`state "%s" %w`, nstate, ErrInvalidState)
			}
			f.mu.Lock()
			f.current = nstate
			f.mu.Unlock()

			// wrap the handler func
			for i := len(options.Wrappers); i > 0; i-- {
				fn = options.Wrappers[i-1](fn)
			}

			s, err = fn(ctx, s, sopts...)

			switch {
			case err != nil:
				return s.Body(), err
			case s.Name() == StateEnd:
				return s.Body(), nil
			case s.Name() == "":
				for idx := range f.statesOrder {
					if f.statesOrder[idx] == nstate && len(f.statesOrder) > idx+1 {
						nstate = f.statesOrder[idx+1]
					}
				}
			default:
				nstate = s.Name()
			}
		}
	}
}