Move redis task completion to script

Author: cschleiden

backend/redis/delete.go

@@ -14,9 +14,12 @@ import (
 // KEYS[4] - payload key
 // KEYS[5] - active-instance-execution key
 // KEYS[6] - instances-by-creation key
+// KEYS[7] - instances
 // ARGV[1] - instance segment
+// ARGV[2] - instance id
 var deleteCmd = redis.NewScript(
 	`redis.call("DEL", KEYS[1], KEYS[2], KEYS[3], KEYS[4], KEYS[5])
+	redis.call("HDEL", KEYS[7], ARGV[1])
 	return redis.call("ZREM", KEYS[6], ARGV[1])`)
 
 // deleteInstance deletes an instance from Redis. It does not attempt to remove any future events or pending
@@ -31,7 +34,8 @@ func deleteInstance(ctx context.Context, rdb redis.UniversalClient, instance *co
 		payloadKey(instance),
 		activeInstanceExecutionKey(instance.InstanceID),
 		instancesByCreation(),
-	}, instanceSegment(instance)).Err(); err != nil {
+		instanceIDs(),
+	}, instanceSegment(instance), instance.InstanceID).Err(); err != nil {
 		return fmt.Errorf("failed to delete instance: %w", err)
 	}
 

backend/redis/events.go

@@ -72,34 +72,3 @@ func addEventToStreamP(ctx context.Context, p redis.Pipeliner, streamKey string,
 		},
 	}).Err()
 }
-
-// addEventsToStream adds the given events to the given event stream. If successful, the message id of the last event added
-// is returned
-// KEYS[1] - stream key
-// ARGV[1] - event data as serialized strings
-var addEventsToStreamCmd = redis.NewScript(`
-	local msgID = ""
-	for i = 1, #ARGV, 2 do
-		msgID = redis.call("XADD", KEYS[1], ARGV[i], "event", ARGV[i + 1])
-	end
-	return msgID
-`)
-
-func addEventsToStreamP(ctx context.Context, p redis.Pipeliner, streamKey string, events []*history.Event) error {
-	eventsData := make([]string, 0)
-	for _, event := range events {
-		eventData, err := marshalEventWithoutAttributes(event)
-		if err != nil {
-			return err
-		}
-
-		// log.Println("addEventsToHistoryStreamP:", event.SequenceID, string(eventData))
-
-		eventsData = append(eventsData, historyID(event.SequenceID))
-		eventsData = append(eventsData, string(eventData))
-	}
-
-	addEventsToStreamCmd.Run(ctx, p, []string{streamKey}, eventsData)
-
-	return nil
-}

backend/redis/expire.go

@@ -18,17 +18,20 @@ import (
 // KEYS[4] - pending events key
 // KEYS[5] - history key
 // KEYS[6] - payload key
+// KEYS[7] - instances key
 // ARGV[1] - current timestamp
 // ARGV[2] - expiration time in seconds
 // ARGV[3] - expiration timestamp in unix milliseconds
 // ARGV[4] - instance segment
+// ARGV[5] - instance id
 var expireCmd = redis.NewScript(
 	`-- Find instances which have already expired and remove from the index set
 	local expiredInstances = redis.call("ZRANGE", KEYS[2], "-inf", ARGV[1], "BYSCORE")
 	for i = 1, #expiredInstances do
 		local instanceSegment = expiredInstances[i]
 		redis.call("ZREM", KEYS[1], instanceSegment) -- index set
 		redis.call("ZREM", KEYS[2], instanceSegment) -- expiration set
+		redis.call("HDEL", KEYS[7], ARGV[5])
 	end
 
 	-- Add expiration time for future cleanup
@@ -57,10 +60,12 @@ func setWorkflowInstanceExpiration(ctx context.Context, rdb redis.UniversalClien
 		pendingEventsKey(instance),
 		historyKey(instance),
 		payloadKey(instance),
+		instanceIDs(),
 	},
 		nowStr,
 		expiration.Seconds(),
 		expStr,
 		instanceSegment(instance),
+		instance.InstanceID,
 	).Err()
 }

backend/redis/instance.go

@@ -26,7 +26,7 @@ func (rb *redisBackend) CreateWorkflowInstance(ctx context.Context, instance *wo
 
 	p := rb.rdb.TxPipeline()
 
-	if err := createInstanceP(ctx, p, instance, event.Attributes.(*history.ExecutionStartedAttributes).Metadata, false); err != nil {
+	if err := createInstanceP(ctx, p, instance, event.Attributes.(*history.ExecutionStartedAttributes).Metadata); err != nil {
 		return err
 	}
 
@@ -145,7 +145,7 @@ type instanceState struct {
 	LastSequenceID int64 `json:"last_sequence_id,omitempty"`
 }
 
-func createInstanceP(ctx context.Context, p redis.Pipeliner, instance *core.WorkflowInstance, metadata *metadata.WorkflowMetadata, ignoreDuplicate bool) error {
+func createInstanceP(ctx context.Context, p redis.Pipeliner, instance *core.WorkflowInstance, metadata *metadata.WorkflowMetadata) error {
 	key := instanceKey(instance)
 
 	createdAt := time.Now()
@@ -165,6 +165,9 @@ func createInstanceP(ctx context.Context, p redis.Pipeliner, instance *core.Work
 	// The newly created instance is going to be the active execution
 	setActiveInstanceExecutionP(ctx, p, instance)
 
+	// Record instance id
+	p.HSet(ctx, instanceIDs(), instance.InstanceID, 1)
+
 	p.ZAdd(ctx, instancesByCreation(), redis.Z{
 		Member: instanceSegment(instance),
 		Score:  float64(createdAt.UnixMilli()),

backend/redis/keys.go

@@ -37,6 +37,10 @@ func instancesExpiring() string {
 	return "instances-expiring"
 }
 
+func instanceIDs() string {
+	return "instances"
+}
+
 func pendingEventsKey(instance *core.WorkflowInstance) string {
 	return fmt.Sprintf("pending-events:%v", instanceSegment(instance))
 }

backend/redis/queue.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log"
 	"time"
 
 	"github.com/google/uuid"
@@ -119,7 +120,6 @@ var createGroupCmd = redis.NewScript(`
     return true
 `)
 
-
 func (q *taskQueue[T]) Enqueue(ctx context.Context, p redis.Pipeliner, id string, data *T) error {
 	ds, err := json.Marshal(data)
 	if err != nil {
@@ -139,6 +139,7 @@ func (q *taskQueue[T]) Dequeue(ctx context.Context, rdb redis.UniversalClient, l
 	}
 
 	if task != nil {
+		log.Println("Recovered task", task.ID)
 		return task, nil
 	}
 

backend/redis/redis.go

@@ -2,7 +2,9 @@ package redis
 
 import (
 	"context"
+	"embed"
 	"fmt"
+	"io/fs"
 	"log/slog"
 	"time"
 
@@ -19,6 +21,11 @@ import (
 
 var _ backend.Backend = (*redisBackend)(nil)
 
+//go:embed scripts/*.lua
+var luaScripts embed.FS
+
+var completeWorkflowTaskCmd *redis.Script
+
 func NewRedisBackend(client redis.UniversalClient, opts ...RedisBackendOption) (*redisBackend, error) {
 	workflowQueue, err := newTaskQueue[any](client, "workflows")
 	if err != nil {
@@ -52,15 +59,13 @@ func NewRedisBackend(client redis.UniversalClient, opts ...RedisBackendOption) (
 	// them, loads them. This doesn't work when using (transactional) pipelines, so eagerly load them on startup.
 	ctx := context.Background()
 	cmds := map[string]*redis.StringCmd{
-		"addEventsToStreamCmd":   addEventsToStreamCmd.Load(ctx, rb.rdb),
-		"addFutureEventCmd":      addFutureEventCmd.Load(ctx, rb.rdb),
-		"futureEventsCmd":        futureEventsCmd.Load(ctx, rb.rdb),
-		"removeFutureEventCmd":   removeFutureEventCmd.Load(ctx, rb.rdb),
-		"removePendingEventsCmd": removePendingEventsCmd.Load(ctx, rb.rdb),
-		"requeueInstanceCmd":     requeueInstanceCmd.Load(ctx, rb.rdb),
-		"deleteInstanceCmd":      deleteCmd.Load(ctx, rb.rdb),
-		"expireInstanceCmd":      expireCmd.Load(ctx, rb.rdb),
-		"addPayloadsCmd":         addPayloadsCmd.Load(ctx, rb.rdb),
+		"addEventsToStreamCmd": addEventsToStreamCmd.Load(ctx, rb.rdb),
+		"addFutureEventCmd":    addFutureEventCmd.Load(ctx, rb.rdb),
+		"futureEventsCmd":      futureEventsCmd.Load(ctx, rb.rdb),
+		"removeFutureEventCmd": removeFutureEventCmd.Load(ctx, rb.rdb),
+		"deleteInstanceCmd":    deleteCmd.Load(ctx, rb.rdb),
+		"expireInstanceCmd":    expireCmd.Load(ctx, rb.rdb),
+		"addPayloadsCmd":       addPayloadsCmd.Load(ctx, rb.rdb),
 	}
 	for name, cmd := range cmds {
 		// fmt.Println(name, cmd.Val())
@@ -70,6 +75,25 @@ func NewRedisBackend(client redis.UniversalClient, opts ...RedisBackendOption) (
 		}
 	}
 
+	// Load all Lua scripts
+
+	cmdMapping := map[string]**redis.Script{
+		"complete_workflow_task.lua": &completeWorkflowTaskCmd,
+	}
+
+	for scriptFile, cmd := range cmdMapping {
+		scriptContent, err := fs.ReadFile(luaScripts, "scripts/"+scriptFile)
+		if err != nil {
+			return nil, fmt.Errorf("reading Lua script %s: %w", scriptFile, err)
+		}
+
+		*cmd = redis.NewScript(string(scriptContent))
+
+		if c := (*cmd).Load(ctx, rb.rdb); c.Err() != nil {
+			return nil, fmt.Errorf("loading Lua script %s: %w", scriptFile, c.Err())
+		}
+	}
+
 	return rb, nil
 }