Support cancel / delete of in-flight stream messages (#216)

* Add cancelation metadata to Redis

Add some cancelation metadata to redis that allows API to introspect the
stream and message ID from the prediction ID. This allows api to perform
an `XDEL` on the redis message when a cancelation is received. The only
time an `XDEL` should happen is if the prediction has not started.

API will be responsible to check `XPENDING` and/or the prediciton cache.
Additionally API should `DEL` the meta key when it performs a
cancelation or when the prediction completes.

key format: `meta:cancelation:<prediction_id>` and it contains json
in the form of: `{"stream_id": "<stream key>", "msg_id": "<redis message
id>"}`

* Move message tracking to separate write script

and add client `Del` command for the other side.

* Shore up tests for queue client Del

and switch meta cancelation key prefix to start with `_` instead of `:`
given the latter tends to be expected as a separator, imho.

* Renaming to "track field" and using the sha1 of the field value

for better alignment with current names and a smidge of safety.

* Combine two string ops for efficiency

Co-authored-by: Mark Phelps <[email protected]>

* A few more refinements based on PR feedback

---------

Co-authored-by: Morgan Fainberg <[email protected]>
Co-authored-by: Mark Phelps <[email protected]>

Author: 3 people

queue/client.go

@@ -2,6 +2,9 @@ package queue
 
 import (
 	"context"
+	"crypto/sha1"
+	"encoding/json"
+	"errors"
 	"fmt"
 	"regexp"
 	"strings"
@@ -13,15 +16,18 @@ import (
 )
 
 var (
-	ErrInvalidReadArgs  = fmt.Errorf("queue: invalid read arguments")
-	ErrInvalidWriteArgs = fmt.Errorf("queue: invalid write arguments")
+	ErrInvalidReadArgs           = errors.New("queue: invalid read arguments")
+	ErrInvalidWriteArgs          = errors.New("queue: invalid write arguments")
+	ErrNoMatchingMessageInStream = errors.New("queue: no matching message in stream")
 
 	streamSuffixPattern = regexp.MustCompile(`\A:s(\d+)\z`)
 )
 
 type Client struct {
 	rdb redis.Cmdable
 	ttl time.Duration // ttl for all keys in queue
+
+	trackField string
 }
 
 type Stats struct {
@@ -32,10 +38,11 @@ type Stats struct {
 }
 
 func NewClient(rdb redis.Cmdable, ttl time.Duration) *Client {
-	return &Client{
-		rdb: rdb,
-		ttl: ttl,
-	}
+	return &Client{rdb: rdb, ttl: ttl}
+}
+
+func NewTrackingClient(rdb redis.Cmdable, ttl time.Duration, field string) *Client {
+	return &Client{rdb: rdb, ttl: ttl, trackField: field}
 }
 
 // Prepare stores the write and read scripts in the Redis script cache so that
@@ -252,16 +259,64 @@ func (c *Client) write(ctx context.Context, args *WriteArgs) (string, error) {
 	cmdArgs = append(cmdArgs, int(c.ttl.Seconds()))
 	cmdArgs = append(cmdArgs, args.Streams)
 	cmdArgs = append(cmdArgs, len(shard))
+
+	if c.trackField != "" {
+		cmdArgs = append(cmdArgs, c.trackField)
+	}
+
 	for _, s := range shard {
 		cmdArgs = append(cmdArgs, s)
 	}
 	for k, v := range args.Values {
 		cmdArgs = append(cmdArgs, k, v)
 	}
 
+	if c.trackField != "" {
+		return writeTrackingScript.Run(ctx, c.rdb, cmdKeys, cmdArgs...).Text()
+	}
+
 	return writeScript.Run(ctx, c.rdb, cmdKeys, cmdArgs...).Text()
 }
 
+type metaCancelation struct {
+	StreamID string `json:"stream_id"`
+	MsgID    string `json:"msg_id"`
+}
+
+// Del supports removal of a message when the given `fieldValue` matches a "meta
+// cancelation" key as written when using a client with tracking support.
+func (c *Client) Del(ctx context.Context, fieldValue string) error {
+	metaCancelationKey := fmt.Sprintf("_meta:cancelation:%x", sha1.Sum([]byte(fieldValue)))
+
+	msgBytes, err := c.rdb.Get(ctx, metaCancelationKey).Bytes()
+	if err != nil {
+		return err
+	}
+
+	msg := &metaCancelation{}
+	if err := json.Unmarshal(msgBytes, msg); err != nil {
+		return err
+	}
+
+	n, err := c.rdb.XDel(ctx, msg.StreamID, msg.MsgID).Result()
+	if err != nil {
+		return err
+	}
+
+	if n == 0 {
+		return fmt.Errorf(
+			"key=%q field-value=%q stream=%q message-id=%q: %w",
+			metaCancelationKey,
+			fieldValue,
+			msg.StreamID,
+			msg.MsgID,
+			ErrNoMatchingMessageInStream,
+		)
+	}
+
+	return nil
+}
+
 func parse(v any) (*Message, error) {
 	result, err := parseSliceWithLength(v, 1)
 	if err != nil {

queue/client_test.go

@@ -1,7 +1,9 @@
 package queue_test
 
 import (
+	"context"
 	crand "crypto/rand"
+	"crypto/sha1"
 	"errors"
 	"fmt"
 	"math/rand"
@@ -19,16 +21,34 @@ import (
 
 	"github.com/replicate/go/queue"
 	"github.com/replicate/go/test"
+	"github.com/replicate/go/uuid"
 )
 
 func TestClientIntegration(t *testing.T) {
-	ctx := test.Context(t)
-	rdb := test.Redis(ctx, t)
+	t.Run("vanilla", func(t *testing.T) {
+		ctx := test.Context(t)
+		rdb := test.Redis(ctx, t)
 
-	ttl := 24 * time.Hour
-	client := queue.NewClient(rdb, ttl)
-	require.NoError(t, client.Prepare(ctx))
+		ttl := 24 * time.Hour
+		client := queue.NewClient(rdb, ttl)
+		require.NoError(t, client.Prepare(ctx))
+
+		runClientIntegrationTest(ctx, t, client)
+	})
+
+	t.Run("with-tracking", func(t *testing.T) {
+		ctx := test.Context(t)
+		rdb := test.Redis(ctx, t)
 
+		ttl := 24 * time.Hour
+		client := queue.NewTrackingClient(rdb, ttl, "id")
+		require.NoError(t, client.Prepare(ctx))
+
+		runClientIntegrationTest(ctx, t, client)
+	})
+}
+
+func runClientIntegrationTest(ctx context.Context, t *testing.T, client *queue.Client) {
 	id := 0
 
 	for range 10 {
@@ -250,36 +270,67 @@ func TestClientReadIntegration(t *testing.T) {
 }
 
 func TestClientWriteIntegration(t *testing.T) {
-	ctx := test.Context(t)
-	rdb := test.Redis(ctx, t)
+	t.Run("vanilla", func(t *testing.T) {
+		ctx := test.Context(t)
+		rdb := test.Redis(ctx, t)
 
-	ttl := 24 * time.Hour
-	client := queue.NewClient(rdb, ttl)
-	require.NoError(t, client.Prepare(ctx))
+		ttl := 24 * time.Hour
+		client := queue.NewClient(rdb, ttl)
+		require.NoError(t, client.Prepare(ctx))
+
+		runClientWriteIntegrationTest(ctx, t, rdb, client, false)
+	})
+
+	t.Run("with-tracking", func(t *testing.T) {
+		ctx := test.Context(t)
+		rdb := test.Redis(ctx, t)
+
+		ttl := 24 * time.Hour
+		client := queue.NewTrackingClient(rdb, ttl, "tracketytrack")
+		require.NoError(t, client.Prepare(ctx))
+
+		runClientWriteIntegrationTest(ctx, t, rdb, client, true)
+	})
+}
+
+func runClientWriteIntegrationTest(ctx context.Context, t *testing.T, rdb *redis.Client, client *queue.Client, withTracking bool) {
+	trackIDs := []string{}
 
 	for i := range 10 {
-		_, err := client.Write(ctx, &queue.WriteArgs{
+		trackID, err := uuid.NewV7()
+		require.NoError(t, err)
+
+		trackIDs = append(trackIDs, trackID.String())
+
+		_, err = client.Write(ctx, &queue.WriteArgs{
 			Name:            "myqueue",
 			Streams:         2,
 			StreamsPerShard: 1,
 			ShardKey:        []byte("panda"),
 			Values: map[string]any{
-				"name": "panda",
-				"idx":  i,
+				"idx":           i,
+				"name":          "panda",
+				"tracketytrack": trackID.String(),
 			},
 		})
 		require.NoError(t, err)
 	}
 
 	for i := range 5 {
-		_, err := client.Write(ctx, &queue.WriteArgs{
+		trackID, err := uuid.NewV7()
+		require.NoError(t, err)
+
+		trackIDs = append(trackIDs, trackID.String())
+
+		_, err = client.Write(ctx, &queue.WriteArgs{
 			Name:            "myqueue",
 			Streams:         2,
 			StreamsPerShard: 1,
 			ShardKey:        []byte("giraffe"),
 			Values: map[string]any{
-				"name": "giraffe",
-				"idx":  i,
+				"idx":           i,
+				"name":          "giraffe",
+				"tracketytrack": trackID.String(),
 			},
 		})
 		require.NoError(t, err)
@@ -306,10 +357,10 @@ func TestClientWriteIntegration(t *testing.T) {
 		require.NoError(t, err)
 		assert.Len(t, values, 5)
 		for i, v := range values {
-			assert.Equal(t, map[string]interface{}{
-				"name": "giraffe",
-				"idx":  strconv.Itoa(i),
-			}, v.Values)
+			assert.Contains(t, v.Values, "name")
+			assert.Contains(t, v.Values, "idx")
+			assert.Equal(t, v.Values["name"], "giraffe")
+			assert.Equal(t, v.Values["idx"], strconv.Itoa(i))
 		}
 	}
 
@@ -323,10 +374,10 @@ func TestClientWriteIntegration(t *testing.T) {
 		require.NoError(t, err)
 		assert.Len(t, values, 10)
 		for i, v := range values {
-			assert.Equal(t, map[string]interface{}{
-				"name": "panda",
-				"idx":  strconv.Itoa(i),
-			}, v.Values)
+			assert.Contains(t, v.Values, "name")
+			assert.Contains(t, v.Values, "idx")
+			assert.Equal(t, v.Values["name"], "panda")
+			assert.Equal(t, v.Values["idx"], strconv.Itoa(i))
 		}
 	}
 
@@ -343,6 +394,62 @@ func TestClientWriteIntegration(t *testing.T) {
 		require.NoError(t, err)
 		assert.Greater(t, ttl, 23*time.Hour)
 	}
+
+	if !withTracking {
+		return
+	}
+
+	for _, trackID := range trackIDs {
+		require.NoError(t, client.Del(ctx, trackID))
+	}
+}
+
+func TestClientDelIntegration(t *testing.T) {
+	ctx := test.Context(t)
+	rdb := test.Redis(ctx, t)
+
+	ttl := 24 * time.Hour
+	client := queue.NewTrackingClient(rdb, ttl, "tracketytrack")
+	require.NoError(t, client.Prepare(ctx))
+
+	trackIDs := []string{}
+
+	for i := range 3 {
+		trackID, err := uuid.NewV7()
+		require.NoError(t, err)
+
+		trackIDs = append(trackIDs, trackID.String())
+
+		_, err = client.Write(ctx, &queue.WriteArgs{
+			Name:            "myqueue",
+			Streams:         2,
+			StreamsPerShard: 1,
+			ShardKey:        []byte("capybara"),
+			Values: map[string]any{
+				"idx":           i,
+				"name":          "capybara",
+				"tracketytrack": trackID.String(),
+			},
+		})
+		require.NoError(t, err)
+	}
+
+	require.NoError(t, client.Del(ctx, trackIDs[0]))
+	require.Error(t, client.Del(ctx, trackIDs[0]))
+	require.Error(t, client.Del(ctx, trackIDs[0]+"oops"))
+	require.Error(t, client.Del(ctx, "bogustown"))
+
+	metaCancelationKey := "_meta:cancelation:" + fmt.Sprintf("%x", sha1.Sum([]byte(trackIDs[1])))
+
+	metaCancel, err := rdb.Get(ctx, metaCancelationKey).Result()
+	require.NoError(t, err)
+
+	rdb.SetEx(ctx, metaCancelationKey, "{{[,"+metaCancel, 5*time.Second)
+
+	require.Error(t, client.Del(ctx, trackIDs[1]))
+
+	require.NoError(t, client.Del(ctx, trackIDs[2]))
+	require.ErrorIs(t, client.Del(ctx, trackIDs[2]), queue.ErrNoMatchingMessageInStream)
 }
 
 // TestPickupLatencyIntegration runs a test with a mostly-empty queue -- by

queue/queue.go

@@ -59,6 +59,10 @@ var (
 	//go:embed write.lua
 	writeCmd    string
 	writeScript = redis.NewScript(writeCmd)
+
+	//go:embed writetracking.lua
+	writeTrackingCmd    string
+	writeTrackingScript = redis.NewScript(writeTrackingCmd)
 )
 
 func prepare(ctx context.Context, rdb redis.Cmdable) error {
@@ -77,5 +81,8 @@ func prepare(ctx context.Context, rdb redis.Cmdable) error {
 	if err := writeScript.Load(ctx, rdb).Err(); err != nil {
 		return err
 	}
+	if err := writeTrackingScript.Load(ctx, rdb).Err(); err != nil {
+		return err
+	}
 	return nil
 }