Merge pull request #324 from cschleiden/prefix-keys

Support key prefix for redis backend

Author: cschleiden

backend/redis/delete.go

@@ -23,14 +23,14 @@ var deleteCmd = redis.NewScript(
 // workflow tasks. It's assumed that the instance is in the finished state.
 //
 // Note: might want to revisit this in the future if we want to support removing hung instances.
-func deleteInstance(ctx context.Context, rdb redis.UniversalClient, instance *core.WorkflowInstance) error {
-	if err := deleteCmd.Run(ctx, rdb, []string{
-		instanceKey(instance),
-		pendingEventsKey(instance),
-		historyKey(instance),
-		payloadKey(instance),
-		activeInstanceExecutionKey(instance.InstanceID),
-		instancesByCreation(),
+func (rb *redisBackend) deleteInstance(ctx context.Context, instance *core.WorkflowInstance) error {
+	if err := deleteCmd.Run(ctx, rb.rdb, []string{
+		rb.keys.instanceKey(instance),
+		rb.keys.pendingEventsKey(instance),
+		rb.keys.historyKey(instance),
+		rb.keys.payloadKey(instance),
+		rb.keys.activeInstanceExecutionKey(instance.InstanceID),
+		rb.keys.instancesByCreation(),
 	}, instanceSegment(instance)).Err(); err != nil {
 		return fmt.Errorf("failed to delete instance: %w", err)
 	}

backend/redis/diagnostics.go

@@ -18,7 +18,7 @@ func (rb *redisBackend) GetWorkflowInstances(ctx context.Context, afterInstanceI
 
 	if afterInstanceID != "" {
 		afterID := instanceSegment(core.NewWorkflowInstance(afterInstanceID, afterExecutionID))
-		scores, err := rb.rdb.ZMScore(ctx, instancesByCreation(), afterID).Result()
+		scores, err := rb.rdb.ZMScore(ctx, rb.keys.instancesByCreation(), afterID).Result()
 		if err != nil {
 			return nil, fmt.Errorf("getting instance score for %v: %w", afterID, err)
 		}
@@ -35,7 +35,7 @@ func (rb *redisBackend) GetWorkflowInstances(ctx context.Context, afterInstanceI
 	}
 
 	result, err := rb.rdb.ZRangeArgs(ctx, redis.ZRangeArgs{
-		Key:     instancesByCreation(),
+		Key:     rb.keys.instancesByCreation(),
 		Stop:    max,
 		Start:   "-inf",
 		ByScore: true,
@@ -49,7 +49,7 @@ func (rb *redisBackend) GetWorkflowInstances(ctx context.Context, afterInstanceI
 	instanceKeys := make([]string, 0)
 	for _, r := range result {
 		instanceSegment := r
-		instanceKeys = append(instanceKeys, instanceKeyFromSegment(instanceSegment))
+		instanceKeys = append(instanceKeys, rb.keys.instanceKeyFromSegment(instanceSegment))
 	}
 
 	if len(instanceKeys) == 0 {
@@ -85,7 +85,7 @@ func (rb *redisBackend) GetWorkflowInstances(ctx context.Context, afterInstanceI
 }
 
 func (rb *redisBackend) GetWorkflowInstance(ctx context.Context, instance *core.WorkflowInstance) (*diag.WorkflowInstanceRef, error) {
-	instanceState, err := readInstance(ctx, rb.rdb, instanceKey(instance))
+	instanceState, err := readInstance(ctx, rb.rdb, rb.keys.instanceKey(instance))
 	if err != nil {
 		return nil, err
 	}

backend/redis/events.go

@@ -43,7 +43,7 @@ var addPayloadsCmd = redis.NewScript(`
 	return 0
 `)
 
-func addEventPayloadsP(ctx context.Context, p redis.Pipeliner, instance *core.WorkflowInstance, events []*history.Event) error {
+func (rb *redisBackend) addEventPayloadsP(ctx context.Context, p redis.Pipeliner, instance *core.WorkflowInstance, events []*history.Event) error {
 	args := make([]interface{}, 0)
 
 	for _, event := range events {
@@ -55,7 +55,7 @@ func addEventPayloadsP(ctx context.Context, p redis.Pipeliner, instance *core.Wo
 		args = append(args, event.ID, string(payload))
 	}
 
-	return addPayloadsCmd.Run(ctx, p, []string{payloadKey(instance)}, args...).Err()
+	return addPayloadsCmd.Run(ctx, p, []string{rb.keys.payloadKey(instance)}, args...).Err()
 }
 
 func addEventToStreamP(ctx context.Context, p redis.Pipeliner, streamKey string, event *history.Event) error {

backend/redis/events_future.go

@@ -54,7 +54,7 @@ func scheduleFutureEvents(ctx context.Context, rb *redisBackend) error {
 	queueKeys := rb.workflowQueue.Keys()
 
 	if _, err := futureEventsCmd.Run(ctx, rb.rdb, []string{
-		futureEventsKey(),
+		rb.keys.futureEventsKey(),
 		queueKeys.StreamKey,
 		queueKeys.SetKey,
 	}, nowStr).Result(); err != nil && err != redis.Nil {

backend/redis/expire.go

@@ -43,20 +43,20 @@ var expireCmd = redis.NewScript(
 	`,
 )
 
-func setWorkflowInstanceExpiration(ctx context.Context, rdb redis.UniversalClient, instance *core.WorkflowInstance, expiration time.Duration) error {
+func (rb *redisBackend) setWorkflowInstanceExpiration(ctx context.Context, instance *core.WorkflowInstance, expiration time.Duration) error {
 	now := time.Now().UnixMilli()
 	nowStr := strconv.FormatInt(now, 10)
 
 	exp := time.Now().Add(expiration).UnixMilli()
 	expStr := strconv.FormatInt(exp, 10)
 
-	return expireCmd.Run(ctx, rdb, []string{
-		instancesByCreation(),
-		instancesExpiring(),
-		instanceKey(instance),
-		pendingEventsKey(instance),
-		historyKey(instance),
-		payloadKey(instance),
+	return expireCmd.Run(ctx, rb.rdb, []string{
+		rb.keys.instancesByCreation(),
+		rb.keys.instancesExpiring(),
+		rb.keys.instanceKey(instance),
+		rb.keys.pendingEventsKey(instance),
+		rb.keys.historyKey(instance),
+		rb.keys.payloadKey(instance),
 	},
 		nowStr,
 		expiration.Seconds(),

backend/redis/instance.go

@@ -43,12 +43,12 @@ func (rb *redisBackend) CreateWorkflowInstance(ctx context.Context, instance *wo
 	}
 
 	_, err = createWorkflowInstanceCmd.Run(ctx, rb.rdb, []string{
-		instanceKey(instance),
-		activeInstanceExecutionKey(instance.InstanceID),
-		pendingEventsKey(instance),
-		payloadKey(instance),
-		instancesActive(),
-		instancesByCreation(),
+		rb.keys.instanceKey(instance),
+		rb.keys.activeInstanceExecutionKey(instance.InstanceID),
+		rb.keys.pendingEventsKey(instance),
+		rb.keys.payloadKey(instance),
+		rb.keys.instancesActive(),
+		rb.keys.instancesByCreation(),
 		keyInfo.SetKey,
 		keyInfo.StreamKey,
 	},
@@ -81,7 +81,7 @@ func (rb *redisBackend) GetWorkflowInstanceHistory(ctx context.Context, instance
 		start = "(" + historyID(*lastSequenceID)
 	}
 
-	msgs, err := rb.rdb.XRange(ctx, historyKey(instance), start, "+").Result()
+	msgs, err := rb.rdb.XRange(ctx, rb.keys.historyKey(instance), start, "+").Result()
 	if err != nil {
 		return nil, err
 	}
@@ -98,7 +98,7 @@ func (rb *redisBackend) GetWorkflowInstanceHistory(ctx context.Context, instance
 		events = append(events, event)
 	}
 
-	res, err := rb.rdb.HMGet(ctx, payloadKey(instance), payloadKeys...).Result()
+	res, err := rb.rdb.HMGet(ctx, rb.keys.payloadKey(instance), payloadKeys...).Result()
 	if err != nil {
 		return nil, fmt.Errorf("reading payloads: %w", err)
 	}
@@ -114,7 +114,7 @@ func (rb *redisBackend) GetWorkflowInstanceHistory(ctx context.Context, instance
 }
 
 func (rb *redisBackend) GetWorkflowInstanceState(ctx context.Context, instance *core.WorkflowInstance) (core.WorkflowInstanceState, error) {
-	instanceState, err := readInstance(ctx, rb.rdb, instanceKey(instance))
+	instanceState, err := readInstance(ctx, rb.rdb, rb.keys.instanceKey(instance))
 	if err != nil {
 		return core.WorkflowInstanceStateActive, err
 	}
@@ -124,7 +124,7 @@ func (rb *redisBackend) GetWorkflowInstanceState(ctx context.Context, instance *
 
 func (rb *redisBackend) CancelWorkflowInstance(ctx context.Context, instance *core.WorkflowInstance, event *history.Event) error {
 	// Read the instance to check if it exists
-	_, err := readInstance(ctx, rb.rdb, instanceKey(instance))
+	_, err := readInstance(ctx, rb.rdb, rb.keys.instanceKey(instance))
 	if err != nil {
 		return err
 	}
@@ -141,7 +141,7 @@ func (rb *redisBackend) CancelWorkflowInstance(ctx context.Context, instance *co
 }
 
 func (rb *redisBackend) RemoveWorkflowInstance(ctx context.Context, instance *core.WorkflowInstance) error {
-	i, err := readInstance(ctx, rb.rdb, instanceKey(instance))
+	i, err := readInstance(ctx, rb.rdb, rb.keys.instanceKey(instance))
 	if err != nil {
 		return err
 	}
@@ -151,7 +151,7 @@ func (rb *redisBackend) RemoveWorkflowInstance(ctx context.Context, instance *co
 		return backend.ErrInstanceNotFinished
 	}
 
-	return deleteInstance(ctx, rb.rdb, instance)
+	return rb.deleteInstance(ctx, instance)
 }
 
 type instanceState struct {
@@ -199,8 +199,8 @@ func readInstancePipelineCmd(cmd *redis.StringCmd) (*instanceState, error) {
 	return &state, nil
 }
 
-func readActiveInstanceExecution(ctx context.Context, rdb redis.UniversalClient, instanceID string) (*core.WorkflowInstance, error) {
-	val, err := rdb.Get(ctx, activeInstanceExecutionKey(instanceID)).Result()
+func (rb *redisBackend) readActiveInstanceExecution(ctx context.Context, instanceID string) (*core.WorkflowInstance, error) {
+	val, err := rb.rdb.Get(ctx, rb.keys.activeInstanceExecutionKey(instanceID)).Result()
 	if err != nil {
 		if err == redis.Nil {
 			return nil, nil

backend/redis/keys.go

@@ -6,59 +6,72 @@ import (
 	"github.com/cschleiden/go-workflows/core"
 )
 
+type keys struct {
+	// Ensure prefix ends with `:`
+	prefix string
+}
+
+func newKeys(prefix string) *keys {
+	if prefix != "" && prefix[len(prefix)-1] != ':' {
+		prefix += ":"
+	}
+
+	return &keys{prefix: prefix}
+}
+
 // activeInstanceExecutionKey returns the key for the latest execution of the given instance
-func activeInstanceExecutionKey(instanceID string) string {
-	return fmt.Sprintf("active-instance-execution:%v", instanceID)
+func (k *keys) activeInstanceExecutionKey(instanceID string) string {
+	return fmt.Sprintf("%sactive-instance-execution:%v", k.prefix, instanceID)
 }
 
 func instanceSegment(instance *core.WorkflowInstance) string {
 	return fmt.Sprintf("%v:%v", instance.InstanceID, instance.ExecutionID)
 }
 
-func instanceKey(instance *core.WorkflowInstance) string {
-	return instanceKeyFromSegment(instanceSegment(instance))
+func (k *keys) instanceKey(instance *core.WorkflowInstance) string {
+	return k.instanceKeyFromSegment(instanceSegment(instance))
 }
 
-func instanceKeyFromSegment(segment string) string {
-	return fmt.Sprintf("instance:%v", segment)
+func (k *keys) instanceKeyFromSegment(segment string) string {
+	return fmt.Sprintf("%sinstance:%v", k.prefix, segment)
 }
 
 // instancesByCreation returns the key for the ZSET that contains all instances sorted by creation date. The score is the
 // creation time as a unix timestamp. Used for listing all workflow instances in the diagnostics UI.
-func instancesByCreation() string {
-	return "instances-by-creation"
+func (k *keys) instancesByCreation() string {
+	return fmt.Sprintf("%sinstances-by-creation", k.prefix)
 }
 
 // instancesActive returns the key for the SET that contains all active instances. Used for reporting active workflow
 // instances in stats.
-func instancesActive() string {
-	return "instances-active"
+func (k *keys) instancesActive() string {
+	return fmt.Sprintf("%sinstances-active", k.prefix)
 }
 
-func instancesExpiring() string {
-	return "instances-expiring"
+func (k *keys) instancesExpiring() string {
+	return fmt.Sprintf("%sinstances-expiring", k.prefix)
 }
 
-func pendingEventsKey(instance *core.WorkflowInstance) string {
-	return fmt.Sprintf("pending-events:%v", instanceSegment(instance))
+func (k *keys) pendingEventsKey(instance *core.WorkflowInstance) string {
+	return fmt.Sprintf("%spending-events:%v", k.prefix, instanceSegment(instance))
 }
 
-func historyKey(instance *core.WorkflowInstance) string {
-	return fmt.Sprintf("history:%v", instanceSegment(instance))
+func (k *keys) historyKey(instance *core.WorkflowInstance) string {
+	return fmt.Sprintf("%shistory:%v", k.prefix, instanceSegment(instance))
 }
 
 func historyID(sequenceID int64) string {
 	return fmt.Sprintf("%v-0", sequenceID)
 }
 
-func futureEventsKey() string {
-	return "future-events"
+func (k *keys) futureEventsKey() string {
+	return fmt.Sprintf("%sfuture-events", k.prefix)
 }
 
-func futureEventKey(instance *core.WorkflowInstance, scheduleEventID int64) string {
-	return fmt.Sprintf("future-event:%v:%v", instanceSegment(instance), scheduleEventID)
+func (k *keys) futureEventKey(instance *core.WorkflowInstance, scheduleEventID int64) string {
+	return fmt.Sprintf("%sfuture-event:%v:%v", k.prefix, instanceSegment(instance), scheduleEventID)
 }
 
-func payloadKey(instance *core.WorkflowInstance) string {
-	return fmt.Sprintf("payload:%v", instanceSegment(instance))
+func (k *keys) payloadKey(instance *core.WorkflowInstance) string {
+	return fmt.Sprintf("%spayload:%v", k.prefix, instanceSegment(instance))
 }

backend/redis/keys_test.go

@@ -0,0 +1,24 @@
+package redis
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_newKeys(t *testing.T) {
+	t.Run("WithEmptyPrefix", func(t *testing.T) {
+		k := newKeys("")
+		require.Equal(t, "", k.prefix)
+	})
+
+	t.Run("WithNonEmptyPrefixWithoutColon", func(t *testing.T) {
+		k := newKeys("prefix")
+		require.Equal(t, "prefix:", k.prefix)
+	})
+
+	t.Run("WithNonEmptyPrefixWithColon", func(t *testing.T) {
+		k := newKeys("prefix:")
+		require.Equal(t, "prefix:", k.prefix)
+	})
+}

backend/redis/options.go

@@ -13,21 +13,23 @@ type RedisOptions struct {
 
 	AutoExpiration              time.Duration
 	AutoExpirationContinueAsNew time.Duration
+
+	KeyPrefix string
 }
 
 type RedisBackendOption func(*RedisOptions)
 
-func WithBlockTimeout(timeout time.Duration) RedisBackendOption {
+// WithKeyPrefix sets the prefix for all keys used in the Redis backend.
+func WithKeyPrefix(prefix string) RedisBackendOption {
 	return func(o *RedisOptions) {
-		o.BlockTimeout = timeout
+		o.KeyPrefix = prefix
 	}
 }
 
-func WithBackendOptions(opts ...backend.BackendOption) RedisBackendOption {
+// WithBlockTimeout sets the timeout for blocking operations like dequeuing a workflow or activity task
+func WithBlockTimeout(timeout time.Duration) RedisBackendOption {
 	return func(o *RedisOptions) {
-		for _, opt := range opts {
-			opt(&o.Options)
-		}
+		o.BlockTimeout = timeout
 	}
 }
 
@@ -47,3 +49,11 @@ func WithAutoExpirationContinueAsNew(expireContinuedAsNewRunsAfter time.Duration
 		o.AutoExpirationContinueAsNew = expireContinuedAsNewRunsAfter
 	}
 }
+
+func WithBackendOptions(opts ...backend.BackendOption) RedisBackendOption {
+	return func(o *RedisOptions) {
+		for _, opt := range opts {
+			opt(&o.Options)
+		}
+	}
+}

backend/redis/queue.go

@@ -34,16 +34,21 @@ type KeyInfo struct {
 	SetKey    string
 }
 
-func newTaskQueue[T any](rdb redis.UniversalClient, tasktype string) (*taskQueue[T], error) {
+func newTaskQueue[T any](rdb redis.UniversalClient, keyPrefix string, tasktype string) (*taskQueue[T], error) {
+	// Ensure the key prefix ends with a colon
+	if keyPrefix != "" && keyPrefix[len(keyPrefix)-1] != ':' {
+		keyPrefix += ":"
+	}
+
 	tq := &taskQueue[T]{
 		tasktype:   tasktype,
-		setKey:     "task-set:" + tasktype,
-		streamKey:  "task-stream:" + tasktype,
+		setKey:     keyPrefix + "task-set:" + tasktype,
+		streamKey:  keyPrefix + "task-stream:" + tasktype,
 		groupName:  "task-workers",
 		workerName: uuid.NewString(),
 	}
 
-	// Pre-load script
+	// Pre-load scripts
 	cmds := map[string]*redis.StringCmd{
 		"createGroupCmd": createGroupCmd.Load(context.Background(), rdb),
 		"enqueueCmd":     enqueueCmd.Load(context.Background(), rdb),