try to intercept messages

Author: HDT3213

README.md

@@ -53,14 +53,14 @@ func main() {
 	}).WithConcurrent(4) // set the number of concurrent consumers 
 	// send delay message
 	for i := 0; i < 10; i++ {
-		err := queue.SendDelayMsg(strconv.Itoa(i), time.Hour, delayqueue.WithRetryCount(3))
+		_, err := queue.SendDelayMsgV2(strconv.Itoa(i), time.Hour, delayqueue.WithRetryCount(3))
 		if err != nil {
 			panic(err)
 		}
 	}
 	// send schedule message
 	for i := 0; i < 10; i++ {
-		err := queue.SendScheduleMsg(strconv.Itoa(i), time.Now().Add(time.Hour))
+		_, err := queue.SendScheduleMsgV2(strconv.Itoa(i), time.Now().Add(time.Hour))
 		if err != nil {
 			panic(err)
 		}
@@ -71,6 +71,8 @@ func main() {
 }
 ```
 
+> `SendScheduleMsgV2` (`SendDelayMsgV2`) is fully compatible with `SendScheduleMsg` (`SendDelayMsg`)
+
 > Please note that redis/v8 is not compatible with redis cluster 7.x. [detail](https://github.com/redis/go-redis/issues/2085)
 
 > If you are using redis client other than go-redis, you could wrap your redis client into [RedisCli](https://pkg.go.dev/github.com/hdt3213/delayqueue#RedisCli) interface
@@ -95,6 +97,30 @@ func producer() {
 }
 ```
 
+## Intercept/Delete Messages
+
+```go
+msg, err := queue.SendScheduleMsgV2(strconv.Itoa(i), time.Now().Add(time.Second))
+if err != nil {
+	panic(err)
+}
+result, err := queue.TryIntercept(msg)
+if err != nil {
+	panic(err)
+}
+if result.Intercepted {
+	println("interception success!")
+} else {
+	println("interception failed, message has been consumed!")
+}
+```
+
+`SendScheduleMsgV2` and `SendDelayMsgV2` return a structure which contains message tracking information.Then passing it to `TryIntercept` to try to intercept the consumption of the message.
+
+If the message is pending or waiting to consume the interception will succeed.If the message has been consumed or is awaiting retry, the interception will fail, but TryIntercept will prevent subsequent retries.
+
+TryIntercept returns a InterceptResult, which Intercepted field indicates whether it is successful.
+
 ## Options
 
 ### Consume Function

README_CN.md

@@ -48,14 +48,14 @@ func main() {
 	}).WithConcurrent(4) // 设置消费者并发数
 	// 发送延时投递消息
 	for i := 0; i < 10; i++ {
-		err := queue.SendDelayMsg(strconv.Itoa(i), time.Hour, delayqueue.WithRetryCount(3))
+		_, err := queue.SendDelayMsgV2(strconv.Itoa(i), time.Hour, delayqueue.WithRetryCount(3))
 		if err != nil {
 			panic(err)
 		}
 	}
 	// 发送定时投递消息
 	for i := 0; i < 10; i++ {
-		err := queue.SendScheduleMsg(strconv.Itoa(i), time.Now().Add(time.Hour))
+		_, err := queue.SendScheduleMsg(strconv.Itoa(i), time.Now().Add(time.Hour))
 		if err != nil {
 			panic(err)
 		}
@@ -90,6 +90,30 @@ func producer() {
 }
 ```
 
+## 拦截消息/删除消息
+
+```go
+msg, err := queue.SendScheduleMsgV2(strconv.Itoa(i), time.Now().Add(time.Second))
+if err != nil {
+	panic(err)
+}
+result, err := queue.TryIntercept(msg)
+if err != nil {
+	panic(err)
+}
+if result.Intercepted {
+	println("拦截成功!")
+} else {
+	println("拦截失败，消息已被消费!")
+}
+```
+
+`SendScheduleMsgV2` 和 `SendDelayMsgV2` 返回一个可以跟踪消息的结构体。然后将其传递给 `TryIntercept` 就可以尝试拦截消息的消费。
+
+如果消息处于待处理状态(pending)或等待消费(ready)，则可以成功拦截。如果消息已被消费或正在等待重试，则无法拦截，但 TryIntercept 将阻止后续重试。
+
+TryIntercept 返回一个 InterceptResult，其中的 Intercepted 字段会表示拦截是否成功。
+
 ## 选项
 
 ### 回调函数

delayqueue.go

@@ -66,9 +66,11 @@ type RedisCli interface {
 	SMembers(key string) ([]string, error)
 	SRem(key string, members []string) error
 	ZAdd(key string, values map[string]float64) error
-	ZRem(key string, fields []string) error
+	ZRem(key string, fields []string) (int64, error)
 	ZCard(key string) (int64, error)
+	ZScore(key string, member string) (float64, error)
 	LLen(key string) (int64, error)
+	LRem(key string, count int64, value string) (int64, error)
 
 	// Publish used for monitor only
 	Publish(channel string, payload string) error
@@ -206,7 +208,7 @@ func (q *DelayQueue) WithDefaultRetryCount(count uint) *DelayQueue {
 	return q
 }
 
-// WithNackRedeliveryDelay customizes the interval between redelivery and nack (callback returns false) 
+// WithNackRedeliveryDelay customizes the interval between redelivery and nack (callback returns false)
 // If consumption exceeded deadline, the message will be redelivered immediately
 func (q *DelayQueue) WithNackRedeliveryDelay(d time.Duration) *DelayQueue {
 	q.nackRedeliveryDelay = d
@@ -236,8 +238,25 @@ func WithMsgTTL(d time.Duration) interface{} {
 	return msgTTLOpt(d)
 }
 
-// SendScheduleMsg submits a message delivered at given time
-func (q *DelayQueue) SendScheduleMsg(payload string, t time.Time, opts ...interface{}) error {
+// MessageInfo stores information to trace a message
+type MessageInfo struct {
+	id string
+}
+
+func (msg *MessageInfo) ID() string {
+	return msg.id
+}
+
+const (
+	StatePending    = "pending"
+	StateReady      = "ready"
+	StateReadyRetry = "ready_to_retry"
+	StateConsuming  = "consuming"
+	StateUnknown    = "unknown"
+)
+
+// SendScheduleMsgV2 submits a message delivered at given time
+func (q *DelayQueue) SendScheduleMsgV2(payload string, t time.Time, opts ...interface{}) (*MessageInfo, error) {
 	// parse options
 	retryCount := q.defaultRetryCount
 	for _, opt := range opts {
@@ -255,28 +274,90 @@ func (q *DelayQueue) SendScheduleMsg(payload string, t time.Time, opts ...interf
 	msgTTL := t.Sub(now) + q.msgTTL // delivery + q.msgTTL
 	err := q.redisCli.Set(q.genMsgKey(idStr), payload, msgTTL)
 	if err != nil {
-		return fmt.Errorf("store msg failed: %v", err)
+		return nil, fmt.Errorf("store msg failed: %v", err)
 	}
 	// store retry count
 	err = q.redisCli.HSet(q.retryCountKey, idStr, strconv.Itoa(int(retryCount)))
 	if err != nil {
-		return fmt.Errorf("store retry count failed: %v", err)
+		return nil, fmt.Errorf("store retry count failed: %v", err)
 	}
 	// put to pending
 	err = q.redisCli.ZAdd(q.pendingKey, map[string]float64{idStr: float64(t.Unix())})
 	if err != nil {
-		return fmt.Errorf("push to pending failed: %v", err)
+		return nil, fmt.Errorf("push to pending failed: %v", err)
 	}
 	q.reportEvent(NewMessageEvent, 1)
-	return nil
+	return &MessageInfo{
+		id: idStr,
+	}, nil
 }
 
 // SendDelayMsg submits a message delivered after given duration
+func (q *DelayQueue) SendDelayMsgV2(payload string, duration time.Duration, opts ...interface{}) (*MessageInfo, error) {
+	t := time.Now().Add(duration)
+	return q.SendScheduleMsgV2(payload, t, opts...)
+}
+
+// SendScheduleMsg submits a message delivered at given time
+// It is compatible with SendScheduleMsgV2, but does not return MessageInfo
+func (q *DelayQueue) SendScheduleMsg(payload string, t time.Time, opts ...interface{}) error {
+	_, err := q.SendScheduleMsgV2(payload, t, opts...)
+	return err
+}
+
+// SendDelayMsg submits a message delivered after given duration
+// It is compatible with SendDelayMsgV2, but does not return MessageInfo
 func (q *DelayQueue) SendDelayMsg(payload string, duration time.Duration, opts ...interface{}) error {
 	t := time.Now().Add(duration)
 	return q.SendScheduleMsg(payload, t, opts...)
 }
 
+type InterceptResult struct {
+	Intercepted bool
+	State       string
+}
+
+// TryIntercept trys to intercept a message
+func (q *DelayQueue) TryIntercept(msg *MessageInfo) (*InterceptResult, error) {
+	id := msg.ID()
+	// try to intercept at ready
+	removed, err := q.redisCli.LRem(q.readyKey, 0, id)
+	if err != nil {
+		q.logger.Printf("intercept %s from ready failed: %v", id, err)
+	}
+	if removed > 0 {
+		_ = q.redisCli.Del([]string{q.genMsgKey(id)})
+		_ = q.redisCli.HDel(q.retryCountKey, []string{id})
+		return &InterceptResult{
+			Intercepted: true,
+			State:       StateReady,
+		}, nil
+	}	
+	// try to intercept at pending
+	removed, err = q.redisCli.ZRem(q.pendingKey, []string{id})
+	if err != nil {
+		q.logger.Printf("intercept %s from pending failed: %v", id, err)
+	}
+	if removed > 0 {
+		_ = q.redisCli.Del([]string{q.genMsgKey(id)})
+		_ = q.redisCli.HDel(q.retryCountKey, []string{id})
+		return &InterceptResult{
+			Intercepted: true,
+			State:       StatePending,
+		}, nil
+	}
+	// message may be being consumed or has been successfully consumed
+	// if the message has been successfully consumed, the following action will cause nothing
+	// if the message is being consumed，the following action will prevent it from being retried
+	q.redisCli.HDel(q.retryCountKey, []string{id}) 
+	q.redisCli.LRem(q.retryKey, 0, id)
+
+	return &InterceptResult{
+		Intercepted: false,
+		State:       StateUnknown,
+	}, nil
+}
+
 func (q *DelayQueue) loadScript(script string) (string, error) {
 	sha1, err := q.redisCli.ScriptLoad(script)
 	if err != nil {
@@ -420,7 +501,7 @@ func (q *DelayQueue) callback(idStr string) error {
 
 func (q *DelayQueue) ack(idStr string) error {
 	atomic.AddInt32(&q.fetchCount, -1)
-	err := q.redisCli.ZRem(q.unAckKey, []string{idStr})
+	_, err := q.redisCli.ZRem(q.unAckKey, []string{idStr})
 	if err != nil {
 		return fmt.Errorf("remove from unack failed: %v", err)
 	}

delayqueue_test.go

@@ -445,3 +445,84 @@ func TestDelayQueue_NackRedeliveryDelay(t *testing.T) {
 		return
 	}
 }
+
+func TestDelayQueue_TryIntercept(t *testing.T) {
+	redisCli := redis.NewClient(&redis.Options{
+		Addr: "127.0.0.1:6379",
+	})
+	redisCli.FlushDB(context.Background())
+	cb := func(s string) bool {
+		return false
+	}
+	queue := NewQueue("test", redisCli, cb).
+		WithDefaultRetryCount(3).
+		WithNackRedeliveryDelay(time.Minute)
+
+	// intercept pending message
+	msg, err := queue.SendDelayMsgV2("foobar", time.Minute)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	result, err := queue.TryIntercept(msg)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if !result.Intercepted {
+		t.Error("expect intercepted")
+	}
+
+	// intercept ready message
+	msg, err = queue.SendScheduleMsgV2("foobar2", time.Now().Add(-time.Minute))
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	err = queue.pending2Ready()
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	result, err = queue.TryIntercept(msg)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if !result.Intercepted {
+		t.Error("expect intercepted")
+	}
+
+	// prevent from retry
+	msg, err = queue.SendScheduleMsgV2("foobar3", time.Now().Add(-time.Minute))
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	ids, err := queue.beforeConsume()
+	if err != nil {
+		t.Errorf("consume error: %v", err)
+		return
+	}
+	for _, id := range ids {
+		queue.nack(id)
+	}
+	queue.afterConsume()
+	result, err = queue.TryIntercept(msg)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if result.Intercepted {
+		t.Error("expect not intercepted")
+		return
+	}
+	ids, err = queue.beforeConsume()
+	if err != nil {
+		t.Errorf("consume error: %v", err)
+		return
+	}
+	if len(ids) > 0 {
+		t.Error("expect empty messages")
+	}
+}

example/getstarted/main.go

@@ -19,14 +19,14 @@ func main() {
 	}).WithConcurrent(4)
 	// send delay message
 	for i := 0; i < 10; i++ {
-		err := queue.SendDelayMsg(strconv.Itoa(i), time.Second, delayqueue.WithRetryCount(3))
+		_, err := queue.SendDelayMsgV2(strconv.Itoa(i), time.Second, delayqueue.WithRetryCount(3))
 		if err != nil {
 			panic(err)
 		}
 	}
 	// send schedule message
 	for i := 0; i < 10; i++ {
-		err := queue.SendScheduleMsg(strconv.Itoa(i), time.Now().Add(time.Second))
+		_, err := queue.SendScheduleMsgV2(strconv.Itoa(i), time.Now().Add(time.Second))
 		if err != nil {
 			panic(err)
 		}