Merge pull request #361 from cschleiden/remove-instances2

Support auto-expiration for all backends

Author: cschleiden

backend/backend.go

@@ -96,4 +96,7 @@ type Backend interface {
 
 	// Close closes any underlying resources
 	Close() error
+
+	// FeatureSupported returns true if the given feature is supported by the backend
+	FeatureSupported(feature Feature) bool
 }

backend/features.go

@@ -0,0 +1,7 @@
+package backend
+
+type Feature int
+
+const (
+	Feature_Expiration Feature = iota
+)

backend/mock_Backend.go

@@ -75,6 +75,10 @@ type mysqlBackend struct {
 	options    *options
 }
 
+func (mb *mysqlBackend) FeatureSupported(feature backend.Feature) bool {
+	return true
+}
+
 func (mb *mysqlBackend) Close() error {
 	return mb.db.Close()
 }
@@ -191,9 +195,9 @@ func (b *mysqlBackend) RemoveWorkflowInstance(ctx context.Context, instance *cor
 }
 
 func (b *mysqlBackend) RemoveWorkflowInstances(ctx context.Context, options ...backend.RemovalOption) error {
-	ro := &backend.RemovalOptions{}
+	ro := backend.DefaultRemovalOptions
 	for _, opt := range options {
-		opt(ro)
+		opt(&ro)
 	}
 
 	rows, err := b.db.QueryContext(ctx, `SELECT instance_id, execution_id FROM instances WHERE completed_at < ?`, ro.FinishedBefore)
@@ -213,7 +217,7 @@ func (b *mysqlBackend) RemoveWorkflowInstances(ctx context.Context, options ...b
 		executionIDs = append(executionIDs, executionID)
 	}
 
-	batchSize := 100
+	batchSize := ro.BatchSize
 	for i := 0; i < len(instanceIDs); i += batchSize {
 		instanceIDs := instanceIDs[i:min(i+batchSize, len(instanceIDs))]
 		executionIDs := executionIDs[i:min(i+batchSize, len(executionIDs))]
@@ -365,7 +369,11 @@ func (b *mysqlBackend) GetWorkflowInstanceState(ctx context.Context, instance *w
 
 func createInstance(ctx context.Context, tx *sql.Tx, queue workflow.Queue, wfi *workflow.Instance, metadata *workflow.Metadata) error {
 	// Check for existing instance
-	if err := tx.QueryRowContext(ctx, "SELECT 1 FROM `instances` WHERE instance_id = ? AND state = ? LIMIT 1", wfi.InstanceID, core.WorkflowInstanceStateActive).
+	if err := tx.QueryRowContext(
+		ctx,
+		"SELECT 1 FROM `instances` WHERE instance_id = ? AND state = ? LIMIT 1",
+		wfi.InstanceID,
+		core.WorkflowInstanceStateActive).
 		Scan(new(int)); err != sql.ErrNoRows {
 		return backend.ErrInstanceAlreadyExists
 	}

backend/mysql/mysql.go

@@ -143,3 +143,12 @@ func (b *redisBackend) Options() *backend.Options {
 func (rb *redisBackend) Close() error {
 	return rb.rdb.Close()
 }
+
+func (rb *redisBackend) FeatureSupported(feature backend.Feature) bool {
+	switch feature {
+	case backend.Feature_Expiration:
+		return false
+	}
+
+	return true
+}

backend/redis/redis.go

@@ -6,6 +6,11 @@ import (
 
 type RemovalOptions struct {
 	FinishedBefore time.Time
+	BatchSize      int
+}
+
+var DefaultRemovalOptions = RemovalOptions{
+	BatchSize: 100,
 }
 
 type RemovalOption func(o *RemovalOptions)
@@ -15,3 +20,9 @@ func RemoveFinishedBefore(t time.Time) RemovalOption {
 		o.FinishedBefore = t
 	}
 }
+
+func RemoveFinishedBatchSize(size int) RemovalOption {
+	return func(o *RemovalOptions) {
+		o.BatchSize = size
+	}
+}

backend/removal.go

@@ -100,6 +100,10 @@ type sqliteBackend struct {
 
 var _ backend.Backend = (*sqliteBackend)(nil)
 
+func (sb *sqliteBackend) FeatureSupported(feature backend.Feature) bool {
+	return true
+}
+
 func (sb *sqliteBackend) Close() error {
 	if sb.memConn != nil {
 		if err := sb.memConn.Close(); err != nil {
@@ -256,9 +260,9 @@ func (sb *sqliteBackend) RemoveWorkflowInstance(ctx context.Context, instance *c
 }
 
 func (sb *sqliteBackend) RemoveWorkflowInstances(ctx context.Context, options ...backend.RemovalOption) error {
-	ro := &backend.RemovalOptions{}
+	ro := backend.DefaultRemovalOptions
 	for _, opt := range options {
-		opt(ro)
+		opt(&ro)
 	}
 
 	rows, err := sb.db.QueryContext(ctx, `SELECT id, execution_id FROM instances WHERE completed_at < ?`, ro.FinishedBefore)
@@ -278,7 +282,7 @@ func (sb *sqliteBackend) RemoveWorkflowInstances(ctx context.Context, options ..
 		executionIDs = append(executionIDs, executionID)
 	}
 
-	batchSize := 100
+	batchSize := ro.BatchSize
 	for i := 0; i < len(instanceIDs); i += batchSize {
 		instanceIDs := instanceIDs[i:min(i+batchSize, len(instanceIDs))]
 		executionIDs := executionIDs[i:min(i+batchSize, len(executionIDs))]

backend/sqlite/sqlite.go

@@ -15,7 +15,7 @@ import (
 
 var e2eRemovalTests = []backendTest{
 	{
-		name: "RemoveWorkflowInstances_Removes",
+		name: "RemoveWorkflowInstances/Removes",
 		f: func(t *testing.T, ctx context.Context, c *client.Client, w *worker.Worker, b TestBackend) {
 			wf := func(ctx workflow.Context) (bool, error) {
 				return true, nil
@@ -28,20 +28,100 @@ var e2eRemovalTests = []backendTest{
 			require.NoError(t, err)
 
 			now := time.Now()
-			time.Sleep(300 * time.Millisecond)
 
-			_, err = runWorkflowWithResult[bool](t, ctx, c, wf)
+			for i := 0; i < 10; i++ {
+				time.Sleep(300 * time.Millisecond)
+
+				err = b.RemoveWorkflowInstances(ctx, backend.RemoveFinishedBefore(now))
+				if errors.As(err, &backend.ErrNotSupported{}) {
+					t.Skip()
+					return
+				}
+
+				require.NoError(t, err)
+
+				_, err = c.GetWorkflowInstanceState(ctx, workflowA)
+				if errors.Is(err, backend.ErrInstanceNotFound) {
+					break
+				}
+			}
+
+			require.ErrorIs(t, err, backend.ErrInstanceNotFound)
+		},
+	},
+	{
+		name: "AutoExpiration/StartsWorkflowAndRemoves",
+		f: func(t *testing.T, ctx context.Context, c *client.Client, w *worker.Worker, b TestBackend) {
+			wf := func(ctx workflow.Context) (bool, error) {
+				return true, nil
+			}
+
+			register(t, ctx, w, []interface{}{wf}, nil)
+
+			err := c.StartAutoExpiration(ctx, time.Millisecond)
+			if errors.As(err, &backend.ErrNotSupported{}) {
+				t.Skip()
+				return
+			}
+
 			require.NoError(t, err)
 
-			err = b.RemoveWorkflowInstances(ctx, backend.RemoveFinishedBefore(now))
+			workflowA := runWorkflow(t, ctx, c, wf)
+			_, err = client.GetWorkflowResult[bool](ctx, c, workflowA, time.Second*10)
+			require.NoError(t, err)
+
+			for i := 0; i < 10; i++ {
+				time.Sleep(100 * time.Millisecond)
+				_, err = c.GetWorkflowInstanceState(ctx, workflowA)
+				if err != backend.ErrInstanceNotFound {
+					continue
+				} else {
+					break
+				}
+			}
+
+			require.ErrorIs(t, err, backend.ErrInstanceNotFound)
+		},
+	},
+	{
+		name: "AutoExpiration/UpdateDelay",
+		f: func(t *testing.T, ctx context.Context, c *client.Client, w *worker.Worker, b TestBackend) {
+			wf := func(ctx workflow.Context) (bool, error) {
+				return true, nil
+			}
+
+			register(t, ctx, w, []interface{}{wf}, nil)
+
+			err := c.StartAutoExpiration(ctx, time.Hour)
 			if errors.As(err, &backend.ErrNotSupported{}) {
 				t.Skip()
 				return
 			}
 
 			require.NoError(t, err)
 
+			workflowA := runWorkflow(t, ctx, c, wf)
+			_, err = client.GetWorkflowResult[bool](ctx, c, workflowA, time.Second*10)
+			require.NoError(t, err)
+
+			time.Sleep(1000 * time.Millisecond)
 			_, err = c.GetWorkflowInstanceState(ctx, workflowA)
+			require.NotErrorIs(t, err, backend.ErrInstanceNotFound)
+
+			// update delay
+			err = c.StartAutoExpiration(ctx, time.Millisecond)
+			require.NoError(t, err)
+
+			for i := 0; i < 10; i++ {
+				time.Sleep(100 * time.Millisecond)
+				_, err = c.GetWorkflowInstanceState(ctx, workflowA)
+				if err != backend.ErrInstanceNotFound {
+					continue
+				} else {
+					break
+				}
+			}
+
 			require.ErrorIs(t, err, backend.ErrInstanceNotFound)
 		},
 	},

backend/test/e2e_removal.go

@@ -0,0 +1,40 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/cschleiden/go-workflows/backend"
+	"github.com/cschleiden/go-workflows/internal/workflows"
+)
+
+const expirationWorkflowInstanceID = "expiration"
+
+// StartAutoExpiration starts a system workflow that will automatically expire workflow instances.
+//
+// The workflow will run every `delay` and remove all workflow instances finished before Now() - `delay`.
+func (c *Client) StartAutoExpiration(ctx context.Context, delay time.Duration) error {
+	if !c.backend.FeatureSupported(backend.Feature_Expiration) {
+		return backend.ErrNotSupported{}
+	}
+
+	_, err := c.CreateWorkflowInstance(ctx, WorkflowInstanceOptions{
+		InstanceID: expirationWorkflowInstanceID,
+	}, workflows.ExpireWorkflowInstances, delay)
+	if err != nil {
+		if errors.Is(err, backend.ErrInstanceAlreadyExists) {
+			err = c.SignalWorkflow(ctx, expirationWorkflowInstanceID, workflows.UpdateExpirationSignal, delay)
+			if err != nil {
+				return fmt.Errorf("updating expiration workflow: %w", err)
+			}
+
+			return nil
+		}
+
+		return fmt.Errorf("starting expiration workflow: %w", err)
+	}
+
+	return nil
+}

client/expiration.go

@@ -0,0 +1,79 @@
+package workflows
+
+import (
+	"context"
+	"errors"
+	"log/slog"
+	"time"
+
+	"github.com/cschleiden/go-workflows/backend"
+	"github.com/cschleiden/go-workflows/core"
+	"github.com/cschleiden/go-workflows/internal/sync"
+	"github.com/cschleiden/go-workflows/workflow"
+)
+
+const (
+	maxIterations = 10
+
+	UpdateExpirationSignal = "update-expiration"
+)
+
+func ExpireWorkflowInstances(ctx workflow.Context, delay time.Duration) error {
+	logger := workflow.Logger(ctx)
+
+	updates := workflow.NewSignalChannel[time.Duration](ctx, UpdateExpirationSignal)
+
+	for i := 0; i < maxIterations; i++ {
+		tctx, cancelTimer := workflow.WithCancel(ctx)
+		t := workflow.ScheduleTimer(tctx, delay)
+
+		timerFired := false
+		for !timerFired {
+			workflow.Select(ctx,
+				workflow.Receive(updates, func(ctx workflow.Context, s time.Duration, _ bool) {
+					delay = s
+
+					cancelTimer()
+					tctx, cancelTimer = workflow.WithCancel(ctx)
+					t = workflow.ScheduleTimer(tctx, delay)
+				}),
+				workflow.Await(t, func(ctx sync.Context, _ workflow.Future[any]) {
+					timerFired = true
+				}),
+			)
+		}
+
+		before := workflow.Now(ctx).Add(-delay)
+
+		logger.Info("removing workflow instances", slog.Time("before", before))
+
+		var a *Activities
+		_, err := workflow.ExecuteActivity[any](
+			ctx, workflow.ActivityOptions{
+				Queue: core.QueueSystem,
+				RetryOptions: workflow.RetryOptions{
+					MaxAttempts: 2,
+				},
+			}, a.RemoveWorkflowInstances, before).Get(ctx)
+		if err != nil {
+			if errors.As(err, &backend.ErrNotSupported{}) {
+				logger.Warn("removing workflow instances not supported")
+
+				// Stop execution
+				return nil
+			}
+
+			logger.Error("removing workflow instances", slog.Any("error", err))
+		}
+	}
+
+	return workflow.ContinueAsNew(ctx, delay)
+}
+
+type Activities struct {
+	Backend backend.Backend
+}
+
+func (a *Activities) RemoveWorkflowInstances(ctx context.Context, before time.Time) error {
+	return a.Backend.RemoveWorkflowInstances(ctx, backend.RemoveFinishedBefore(before))
+}