dapr_client/saga/orchestrator.go

package saga

import (
	"fmt"
	"time"
)

type StepFunc func() error
type StepRollbackFunc func() error
type OutputFunc func() (interface{}, error)

type Step struct {
	StepFunc            StepFunc
	StepRollbackFunc    StepRollbackFunc
	RollbackContextData string
	rollbackDone        chan interface{}
}

type Orchestrator struct {
	stateStoreName         string
	sagaName               string
	name                   string
	rollbackRetryPeriodSec time.Duration
	steps                  []*Step
	OutputFunc             OutputFunc
}

func (orchestrator *Orchestrator) AddStep(step *Step) *Orchestrator {
	orchestrator.steps = append(orchestrator.steps, step)

	return orchestrator
}

func (orchestrator *Orchestrator) Output(outputFunc OutputFunc) *Orchestrator {
	orchestrator.OutputFunc = outputFunc

	return orchestrator
}

func (orchestrator *Orchestrator) Run() (interface{}, error) {
	for index, step := range orchestrator.steps {
		if step.StepFunc == nil {
			continue
		}

		err := step.StepFunc()
		if err == nil {
			continue
		}

		if err != nil {
			err = saveOrchestratorState(orchestrator.stateStoreName, orchestrator.sagaName,
				orchestrator.name, step.RollbackContextData, index)
			if err != nil {
				return nil, err
			}

			go orchestrator.Rollback(index)
			return nil, err
		}
	}

	if orchestrator.OutputFunc != nil {
		return orchestrator.OutputFunc()
	}

	return nil, nil
}

func (orchestrator *Orchestrator) Rollback(startIndex int) {
	for i := startIndex; i >= 0; i-- {
		rollbackStep := orchestrator.steps[i]
		err := orchestrator.rollbackStep(rollbackStep)
		if err != nil {
			fmt.Println("Rollback", "orchestrator.rollbackStep", err)
		}

		err = saveOrchestratorState(orchestrator.stateStoreName, orchestrator.sagaName, orchestrator.name,
			rollbackStep.RollbackContextData, i-1)
		if err != nil {
			fmt.Println("Rollback", "saveOrchestratorState", err)
			return
		}
	}

	err := deleteOrchestratorState(orchestrator.stateStoreName, orchestrator.sagaName, orchestrator.name)
	if err != nil {
		fmt.Println("Rollback", "deleteOrchestratorState", err)
		return
	}
}

func (orchestrator *Orchestrator) rollbackStep(rollbackStep *Step) error {
	if rollbackStep.StepRollbackFunc == nil {
		return nil
	}

	err := rollbackStep.StepRollbackFunc()
	if err == nil {
		return nil
	}

	orchestrator.rollbackRetry(rollbackStep)

	return nil
}

func (orchestrator *Orchestrator) rollbackRetry(step *Step) {
	ticker := time.NewTicker(orchestrator.rollbackRetryPeriodSec)

	for {
		select {
		case <-step.rollbackDone:
			ticker.Stop()
			ticker = nil
		case <-ticker.C:
			err := step.StepRollbackFunc()
			if err == nil {
				return
			}
		}
	}
}

func (orchestrator *Orchestrator) stop() {
	for _, step := range orchestrator.steps {
		if step.rollbackDone != nil {
			step.rollbackDone <- true
			close(step.rollbackDone)
			step.rollbackDone = nil
		}
	}
}