Merge pull request #316 from cschleiden/continue-as-new-expiration

Support custom expiration for continue-as-new workflows

Author: cschleiden

backend/redis/expire_test.go

@@ -121,3 +121,77 @@ func Test_AutoExpiration_SubWorkflow(t *testing.T) {
 	_, err = b.GetWorkflowInstanceState(ctx, wfi)
 	require.ErrorIs(t, err, backend.ErrInstanceNotFound)
 }
+
+func Test_AutoExpiration_ContinueAsNew_SubWorkflow(t *testing.T) {
+	if testing.Short() {
+		t.Skip()
+	}
+
+	autoExpirationTime := time.Second * 2
+
+	redisClient := getClient()
+	setup := getCreateBackend(redisClient, WithAutoExpiration(0), WithAutoExpirationContinueAsNew(autoExpirationTime))
+	b := setup()
+
+	c := client.New(b)
+	w := worker.New(b, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	require.NoError(t, w.Start(ctx))
+	defer func() {
+		cancel()
+
+		require.NoError(t, w.WaitForCompletion())
+	}()
+
+	var swfInstances []*workflow.Instance
+
+	swf := func(ctx workflow.Context, iteration int) (int, error) {
+		if iteration > 3 {
+			return 42, nil
+		}
+
+		// Keep track of continuedasnew instances
+		swfInstances = append(swfInstances, workflow.WorkflowInstance(ctx))
+
+		return 0, workflow.ContinueAsNew(ctx, iteration+1)
+	}
+
+	swfInstanceID := uuid.NewString()
+
+	wf := func(ctx workflow.Context) (int, error) {
+		l := workflow.Logger(ctx)
+		l.Debug("Starting sub workflow", "instanceID", swfInstanceID)
+
+		r, err := workflow.CreateSubWorkflowInstance[int](ctx, workflow.SubWorkflowOptions{
+			InstanceID: swfInstanceID,
+		}, swf, 0).Get(ctx)
+
+		workflow.ScheduleTimer(ctx, time.Second*2).Get(ctx)
+
+		return r, err
+	}
+
+	w.RegisterWorkflow(wf)
+	w.RegisterWorkflow(swf)
+
+	wfi, err := c.CreateWorkflowInstance(ctx, client.WorkflowInstanceOptions{
+		InstanceID: uuid.NewString(),
+	}, wf)
+	require.NoError(t, err)
+
+	// Wait for redis to expire the keys
+	time.Sleep(autoExpirationTime * 2)
+
+	// Main workflow should still be there
+	r, err := client.GetWorkflowResult[int](ctx, c, wfi, time.Second*10)
+	require.NoError(t, err)
+	require.Equal(t, 42, r)
+
+	// All continued-as-new sub-workflow instances should be expired
+	for _, swfInstance := range swfInstances {
+		_, err = b.GetWorkflowInstanceState(ctx, swfInstance)
+		require.ErrorIs(t, err, backend.ErrInstanceNotFound)
+	}
+}

backend/redis/options.go

@@ -11,7 +11,8 @@ type RedisOptions struct {
 
 	BlockTimeout time.Duration
 
-	AutoExpiration time.Duration
+	AutoExpiration              time.Duration
+	AutoExpirationContinueAsNew time.Duration
 }
 
 type RedisBackendOption func(*RedisOptions)
@@ -37,3 +38,12 @@ func WithAutoExpiration(expireFinishedRunsAfter time.Duration) RedisBackendOptio
 		o.AutoExpiration = expireFinishedRunsAfter
 	}
 }
+
+// WithAutoExpirationContinueAsNew sets the duration after which runs that were completed with `ContinueAsNew`
+// automatically expire.
+// If set to 0 (default), the overall expiration setting set with `WithAutoExpiration` will be used.
+func WithAutoExpirationContinueAsNew(expireContinuedAsNewRunsAfter time.Duration) RedisBackendOption {
+	return func(o *RedisOptions) {
+		o.AutoExpirationContinueAsNew = expireContinuedAsNewRunsAfter
+	}
+}

backend/redis/workflow.go

@@ -279,8 +279,13 @@ func (rb *redisBackend) CompleteWorkflowTask(
 		span.End()
 
 		// Auto expiration
-		if rb.options.AutoExpiration > 0 {
-			if err := setWorkflowInstanceExpiration(ctx, rb.rdb, instance, rb.options.AutoExpiration); err != nil {
+		expiration := rb.options.AutoExpiration
+		if state == core.WorkflowInstanceStateContinuedAsNew && rb.options.AutoExpirationContinueAsNew > 0 {
+			expiration = rb.options.AutoExpirationContinueAsNew
+		}
+
+		if expiration > 0 {
+			if err := setWorkflowInstanceExpiration(ctx, rb.rdb, instance, expiration); err != nil {
 				return fmt.Errorf("setting workflow instance expiration: %w", err)
 			}
 		}