Skip to content

Commit 4baa054

Browse files
committed
Rollback in progress.
1 parent 8290106 commit 4baa054

8 files changed

Lines changed: 530 additions & 5 deletions

File tree

builder.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ func (builder *Builder) Step(name, handler string, opts ...StepOption) *Builder
5555
Handler: handler,
5656
MaxRetries: builder.defaultMaxRetries,
5757
Next: []string{},
58+
Prev: builder.currentStep,
5859
Metadata: make(map[string]string),
5960
}
6061

@@ -97,6 +98,7 @@ func (builder *Builder) OnFailure(name, handler string, opts ...StepOption) *Bui
9798
Type: StepTypeTask,
9899
Handler: handler,
99100
MaxRetries: builder.defaultMaxRetries,
101+
Prev: "", // Compensation steps don't have prev in the main flow
100102
}
101103
for _, opt := range opts {
102104
opt(compensation)
@@ -163,6 +165,7 @@ func (builder *Builder) Parallel(name string, tasks ...*StepDefinition) *Builder
163165
parallelStep := &StepDefinition{
164166
Name: name,
165167
Type: StepTypeParallel,
168+
Prev: builder.currentStep,
166169
Metadata: make(map[string]string),
167170
Parallel: []string{},
168171
}
@@ -195,6 +198,8 @@ func (builder *Builder) Parallel(name string, tasks ...*StepDefinition) *Builder
195198
return builder
196199
}
197200

201+
// Set the parallel step as prev for each task
202+
task.Prev = name
198203
builder.steps[task.Name] = task
199204

200205
parallelStep.Parallel = append(parallelStep.Parallel, task.Name)
@@ -220,6 +225,7 @@ func (builder *Builder) Fork(name string, branches ...func(branch *Builder)) *Bu
220225
forkStep := &StepDefinition{
221226
Name: name,
222227
Type: StepTypeFork,
228+
Prev: builder.currentStep,
223229
Metadata: make(map[string]string),
224230
}
225231

@@ -262,6 +268,11 @@ func (builder *Builder) Fork(name string, branches ...func(branch *Builder)) *Bu
262268
return builder
263269
}
264270

271+
// Set the fork step as prev for the first step of each branch
272+
if stepName == sub.startStep {
273+
stepDef.Prev = name
274+
}
275+
265276
builder.steps[stepName] = stepDef
266277
}
267278

@@ -299,6 +310,7 @@ func (builder *Builder) JoinStep(name string, waitFor []string, strategy JoinStr
299310
Type: StepTypeJoin,
300311
WaitFor: waitFor,
301312
JoinStrategy: strategy,
313+
Prev: builder.currentStep,
302314
Metadata: make(map[string]string),
303315
}
304316

@@ -344,6 +356,7 @@ func (builder *Builder) SavePoint(name string) *Builder {
344356
step := &StepDefinition{
345357
Name: name,
346358
Type: StepTypeSavePoint,
359+
Prev: builder.currentStep,
347360
Metadata: make(map[string]string),
348361
}
349362

@@ -484,6 +497,7 @@ func NewTask(name, handler string, opts ...StepOption) *StepDefinition {
484497
Name: name,
485498
Handler: handler,
486499
Type: StepTypeTask,
500+
Prev: "", // NewTask doesn't have currentStep context
487501
Metadata: make(map[string]string),
488502
MaxRetries: defaultMaxRetries,
489503
}

engine.go

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,8 @@ func (engine *Engine) executeStep(ctx context.Context, instance *WorkflowInstanc
159159
output, stepErr = engine.executeJoin(ctx, instance, step, stepDef)
160160
case StepTypeParallel:
161161
output, stepErr = engine.executeFork(ctx, instance, step, stepDef)
162+
case StepTypeSavePoint:
163+
output = step.Input
162164
default:
163165
stepErr = fmt.Errorf("unsupported step type: %s", stepDef.Type)
164166
}
@@ -410,6 +412,18 @@ func (engine *Engine) handleStepFailure(
410412
return fmt.Errorf("notify join steps: %w", err)
411413
}
412414

415+
// Try to rollback to save point before handling failure
416+
def, err := engine.store.GetWorkflowDefinition(ctx, instance.WorkflowID)
417+
if err == nil {
418+
if rollbackErr := engine.rollbackToSavePoint(ctx, instance.ID, step.StepName, def); rollbackErr != nil {
419+
// Log rollback error but continue with failure handling
420+
_ = engine.store.LogEvent(ctx, instance.ID, &step.ID, EventStepFailed, map[string]any{
421+
KeyStepName: step.StepName,
422+
KeyError: fmt.Sprintf("rollback failed: %v", rollbackErr),
423+
})
424+
}
425+
}
426+
413427
if stepDef.OnFailure != "" {
414428
return engine.enqueueNextSteps(ctx, instance.ID, []string{stepDef.OnFailure}, step.Input)
415429
}
@@ -664,3 +678,158 @@ func (engine *Engine) validateDefinition(def *WorkflowDefinition) error {
664678

665679
return nil
666680
}
681+
682+
func (engine *Engine) rollbackToSavePoint(
683+
ctx context.Context,
684+
instanceID int64,
685+
failedStepName string,
686+
def *WorkflowDefinition,
687+
) error {
688+
savePointName := engine.findNearestSavePoint(failedStepName, def)
689+
if savePointName == "" {
690+
return engine.rollbackAllSteps(ctx, instanceID, failedStepName, def)
691+
}
692+
693+
return engine.rollbackStepsToSavePoint(ctx, instanceID, failedStepName, savePointName, def)
694+
}
695+
696+
func (engine *Engine) findNearestSavePoint(stepName string, def *WorkflowDefinition) string {
697+
visited := make(map[string]bool)
698+
699+
for stepName != "" {
700+
if visited[stepName] {
701+
break // Prevent infinite loops
702+
}
703+
visited[stepName] = true
704+
705+
stepDef, ok := def.Definition.Steps[stepName]
706+
if !ok {
707+
break
708+
}
709+
710+
if stepDef.Type == StepTypeSavePoint {
711+
return stepName
712+
}
713+
714+
stepName = stepDef.Prev
715+
}
716+
717+
return ""
718+
}
719+
720+
func (engine *Engine) rollbackAllSteps(
721+
ctx context.Context,
722+
instanceID int64,
723+
failedStepName string,
724+
def *WorkflowDefinition,
725+
) error {
726+
steps, err := engine.store.GetStepsByInstance(ctx, instanceID)
727+
if err != nil {
728+
return fmt.Errorf("get steps by instance: %w", err)
729+
}
730+
731+
for _, step := range steps {
732+
if step.Status == StepStatusCompleted {
733+
if err := engine.rollbackStep(ctx, step, def); err != nil {
734+
return fmt.Errorf("rollback step %s: %w", step.StepName, err)
735+
}
736+
}
737+
}
738+
739+
return nil
740+
}
741+
742+
func (engine *Engine) rollbackStepsToSavePoint(
743+
ctx context.Context,
744+
instanceID int64,
745+
failedStepName, savePointName string,
746+
def *WorkflowDefinition,
747+
) error {
748+
steps, err := engine.store.GetStepsByInstance(ctx, instanceID)
749+
if err != nil {
750+
return fmt.Errorf("get steps by instance: %w", err)
751+
}
752+
753+
stepMap := make(map[string]*WorkflowStep)
754+
for _, step := range steps {
755+
stepMap[step.StepName] = step
756+
}
757+
758+
return engine.rollbackStepChain(ctx, failedStepName, savePointName, def, stepMap)
759+
}
760+
761+
func (engine *Engine) rollbackStepChain(
762+
ctx context.Context,
763+
currentStep, savePointName string,
764+
def *WorkflowDefinition,
765+
stepMap map[string]*WorkflowStep,
766+
) error {
767+
if currentStep == savePointName {
768+
return nil // Reached save point
769+
}
770+
771+
stepDef, ok := def.Definition.Steps[currentStep]
772+
if !ok {
773+
return fmt.Errorf("step definition not found: %s", currentStep)
774+
}
775+
776+
if step, exists := stepMap[currentStep]; exists && step.Status == StepStatusCompleted {
777+
if err := engine.rollbackStep(ctx, step, def); err != nil {
778+
return fmt.Errorf("rollback step %s: %w", currentStep, err)
779+
}
780+
}
781+
782+
// Handle parallel steps (fork branches)
783+
if stepDef.Type == StepTypeFork || stepDef.Type == StepTypeParallel {
784+
for _, parallelStepName := range stepDef.Parallel {
785+
if err := engine.rollbackStepChain(ctx, parallelStepName, savePointName, def, stepMap); err != nil {
786+
return err
787+
}
788+
}
789+
}
790+
791+
// Continue with a previous step
792+
if stepDef.Prev != "" {
793+
return engine.rollbackStepChain(ctx, stepDef.Prev, savePointName, def, stepMap)
794+
}
795+
796+
return nil
797+
}
798+
799+
func (engine *Engine) rollbackStep(ctx context.Context, step *WorkflowStep, def *WorkflowDefinition) error {
800+
stepDef, ok := def.Definition.Steps[step.StepName]
801+
if !ok {
802+
return fmt.Errorf("step definition not found: %s", step.StepName)
803+
}
804+
805+
handler, exists := engine.handlers[stepDef.OnFailure]
806+
if !exists {
807+
return fmt.Errorf("handler not found: %s", stepDef.Handler)
808+
}
809+
810+
stepCtx := &executionContext{
811+
instanceID: step.InstanceID,
812+
stepName: step.StepName,
813+
retryCount: step.RetryCount,
814+
variables: stepDef.Metadata,
815+
}
816+
817+
// Execute the handler in compensation mode
818+
_, err := handler.Execute(ctx, stepCtx, step.Input)
819+
if err != nil {
820+
return fmt.Errorf("execute compensation for step %q: %w", step.StepName, err)
821+
}
822+
823+
// Update step status to rolled back
824+
if err := engine.store.UpdateStep(ctx, step.ID, StepStatusRolledBack, step.Input, nil); err != nil {
825+
return fmt.Errorf("update step status: %w", err)
826+
}
827+
828+
_ = engine.store.LogEvent(ctx, step.InstanceID, &step.ID, EventStepFailed, map[string]any{
829+
KeyStepName: step.StepName,
830+
KeyStepType: step.StepType,
831+
KeyError: "step rolled back",
832+
})
833+
834+
return nil
835+
}

examples/Makefile

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,6 @@ data-pipeline: ## Run data pipeline example
1515

1616
microservices: ## Run microservices example
1717
cd microservices && go mod tidy && go run main.go
18+
19+
savepoint: ## Run rollback to savepoint example
20+
cd savepoint_demo && go mod tidy && go run main.go

examples/savepoint_demo/go.mod

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
module savepoint-demo
2+
3+
go 1.24
4+
5+
require (
6+
github.com/jackc/pgx/v5 v5.7.6
7+
github.com/rom8726/floxy v0.0.0
8+
)
9+
10+
require (
11+
github.com/google/uuid v1.6.0 // indirect
12+
github.com/jackc/pgpassfile v1.0.0 // indirect
13+
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
14+
github.com/jackc/puddle/v2 v2.2.2 // indirect
15+
github.com/lib/pq v1.10.9 // indirect
16+
golang.org/x/crypto v0.37.0 // indirect
17+
golang.org/x/sync v0.13.0 // indirect
18+
golang.org/x/text v0.24.0 // indirect
19+
)
20+
21+
replace github.com/rom8726/floxy => ../../

examples/savepoint_demo/go.sum

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
2+
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
3+
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
4+
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
5+
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
6+
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
7+
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
8+
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
9+
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
10+
github.com/jackc/pgx/v5 v5.7.6 h1:rWQc5FwZSPX58r1OQmkuaNicxdmExaEz5A2DO2hUuTk=
11+
github.com/jackc/pgx/v5 v5.7.6/go.mod h1:aruU7o91Tc2q2cFp5h4uP3f6ztExVpyVv88Xl/8Vl8M=
12+
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
13+
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
14+
github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
15+
github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
16+
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
17+
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
18+
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
19+
github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
20+
github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
21+
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
22+
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
23+
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
24+
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
25+
golang.org/x/crypto v0.37.0 h1:kJNSjF/Xp7kU0iB2Z+9viTPMW4EqqsrywMXLJOOsXSE=
26+
golang.org/x/crypto v0.37.0/go.mod h1:vg+k43peMZ0pUMhYmVAWysMK35e6ioLh3wB8ZCAfbVc=
27+
golang.org/x/sync v0.13.0 h1:AauUjRAJ9OSnvULf/ARrrVywoJDy0YS2AwQ98I37610=
28+
golang.org/x/sync v0.13.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
29+
golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0=
30+
golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU=
31+
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
32+
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
33+
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
34+
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

0 commit comments

Comments
 (0)