Merge pull request #364 from cschleiden/cleanup-continue-as-new

Support removing ContinuedAsNew instances immediately

Author: cschleiden

backend/mysql/mysql.go

@@ -166,6 +166,14 @@ func (b *mysqlBackend) RemoveWorkflowInstance(ctx context.Context, instance *cor
 	}
 	defer tx.Rollback()
 
+	if err := b.removeWorkflowInstance(ctx, instance, tx); err != nil {
+		return err
+	}
+
+	return tx.Commit()
+}
+
+func (b *mysqlBackend) removeWorkflowInstance(ctx context.Context, instance *core.WorkflowInstance, tx *sql.Tx) error {
 	row := tx.QueryRowContext(ctx, "SELECT state FROM `instances` WHERE instance_id = ? AND execution_id = ? LIMIT 1", instance.InstanceID, instance.ExecutionID)
 	var state core.WorkflowInstanceState
 	if err := row.Scan(&state); err != nil {
@@ -191,7 +199,7 @@ func (b *mysqlBackend) RemoveWorkflowInstance(ctx context.Context, instance *cor
 		return err
 	}
 
-	return tx.Commit()
+	return nil
 }
 
 func (b *mysqlBackend) RemoveWorkflowInstances(ctx context.Context, options ...backend.RemovalOption) error {
@@ -752,6 +760,12 @@ func (b *mysqlBackend) CompleteWorkflowTask(
 		}
 	}
 
+	if b.options.RemoveContinuedAsNewInstances && state == core.WorkflowInstanceStateContinuedAsNew {
+		if err := b.removeWorkflowInstance(ctx, instance, tx); err != nil {
+			return fmt.Errorf("removing old instance: %w", err)
+		}
+	}
+
 	if err := tx.Commit(); err != nil {
 		return fmt.Errorf("committing complete workflow transaction: %w", err)
 	}

backend/options.go

@@ -37,6 +37,11 @@ type Options struct {
 	// ActivityLockTimeout determines how long an activity task can be locked for. If the activity task is not completed
 	// by that timeframe, it's considered abandoned and another worker might pick it up
 	ActivityLockTimeout time.Duration
+
+	// RemoveContinuedAsNewInstances determines whether instances that were completed using ContinueAsNew should be
+	// removed immediately, including their history. If set to false, the instance will be removed after the configured
+	// retention period or never.
+	RemoveContinuedAsNewInstances bool
 }
 
 var DefaultOptions Options = Options{
@@ -50,6 +55,8 @@ var DefaultOptions Options = Options{
 	Converter:      converter.DefaultConverter,
 
 	ContextPropagators: []workflow.ContextPropagator{&tracing.TracingContextPropagator{}},
+
+	RemoveContinuedAsNewInstances: false,
 }
 
 type BackendOption func(*Options)
@@ -90,6 +97,12 @@ func WithContextPropagator(prop workflow.ContextPropagator) BackendOption {
 	}
 }
 
+func WithRemoveContinuedAsNewInstances() BackendOption {
+	return func(o *Options) {
+		o.RemoveContinuedAsNewInstances = true
+	}
+}
+
 func ApplyOptions(opts ...BackendOption) *Options {
 	options := DefaultOptions
 

backend/redis/instance.go

@@ -149,8 +149,7 @@ func (rb *redisBackend) RemoveWorkflowInstance(ctx context.Context, instance *co
 		return err
 	}
 
-	// Check state
-	if i.State != core.WorkflowInstanceStateFinished {
+	if i.State != core.WorkflowInstanceStateFinished && i.State != core.WorkflowInstanceStateContinuedAsNew {
 		return backend.ErrInstanceNotFinished
 	}
 

backend/redis/workflow.go

@@ -323,6 +323,12 @@ func (rb *redisBackend) CompleteWorkflowTask(
 				return fmt.Errorf("setting workflow instance expiration: %w", err)
 			}
 		}
+
+		if rb.options.RemoveContinuedAsNewInstances && state == core.WorkflowInstanceStateContinuedAsNew {
+			if err := rb.RemoveWorkflowInstance(ctx, instance); err != nil {
+				return fmt.Errorf("removing workflow instance: %w", err)
+			}
+		}
 	}
 
 	return nil

backend/sqlite/sqlite.go

@@ -227,10 +227,17 @@ func (sb *sqliteBackend) RemoveWorkflowInstance(ctx context.Context, instance *c
 	}
 	defer tx.Rollback()
 
+	if err := sb.removeWorkflowInstance(ctx, instance, tx); err != nil {
+		return err
+	}
+
+	return tx.Commit()
+}
+
+func (sb *sqliteBackend) removeWorkflowInstance(ctx context.Context, instance *core.WorkflowInstance, tx *sql.Tx) error {
 	instanceID := instance.InstanceID
 	executionID := instance.ExecutionID
 
-	// Check status of the instance
 	row := tx.QueryRowContext(ctx, "SELECT state FROM `instances` WHERE id = ? AND execution_id = ? LIMIT 1", instanceID, executionID)
 	var state core.WorkflowInstanceState
 	if err := row.Scan(&state); err != nil {
@@ -243,7 +250,6 @@ func (sb *sqliteBackend) RemoveWorkflowInstance(ctx context.Context, instance *c
 		return backend.ErrInstanceNotFinished
 	}
 
-	// Delete from instances and history tables
 	if _, err := tx.ExecContext(ctx, "DELETE FROM `instances` WHERE id = ? AND execution_id = ?", instanceID, executionID); err != nil {
 		return err
 	}
@@ -256,7 +262,7 @@ func (sb *sqliteBackend) RemoveWorkflowInstance(ctx context.Context, instance *c
 		return err
 	}
 
-	return tx.Commit()
+	return nil
 }
 
 func (sb *sqliteBackend) RemoveWorkflowInstances(ctx context.Context, options ...backend.RemovalOption) error {
@@ -539,6 +545,7 @@ func (sb *sqliteBackend) GetWorkflowTask(ctx context.Context, queues []workflow.
 	return t, nil
 }
 
+// MARK: CompleteWorkflowTask
 func (sb *sqliteBackend) CompleteWorkflowTask(
 	ctx context.Context,
 	task *backend.WorkflowTask,
@@ -670,6 +677,12 @@ func (sb *sqliteBackend) CompleteWorkflowTask(
 		}
 	}
 
+	if sb.options.RemoveContinuedAsNewInstances && state == core.WorkflowInstanceStateContinuedAsNew {
+		if err := sb.removeWorkflowInstance(ctx, instance, tx); err != nil {
+			return fmt.Errorf("removing old instance: %w", err)
+		}
+	}
+
 	return tx.Commit()
 }