完善队列集合的ACK确认机制

lock_free/option.go

@@ -0,0 +1,35 @@
+package lock_free
+
+import "time"
+
+type Option func(*options)
+
+type options struct {
+	autoAck        bool
+	expireAutoFail time.Duration
+}
+
+func loadOptions(opt ...Option) options {
+	opts := options{
+		autoAck:        true,
+		expireAutoFail: 10 * time.Second,
+	}
+	for _, o := range opt {
+		o(&opts)
+	}
+	return opts
+}
+
+// WithAutoAck 设置自动确认
+func WithAutoAck(autoAck bool) Option {
+	return func(o *options) {
+		o.autoAck = autoAck
+	}
+}
+
+// WithExpireAutoFail 设置手动确认截止时间
+func WithExpireAutoFail(expireAutoFail time.Duration) Option {
+	return func(o *options) {
+		o.expireAutoFail = expireAutoFail
+	}
+}

lock_free/queues.go

@@ -1,6 +1,7 @@
 package lock_free
 
 import (
+	"fmt"
 	"iter"
 	"sync"
 	"time"
@@ -8,16 +9,46 @@ import (
 
 // Queues 队列集合
 type Queues[TKey comparable, TValue, TRoute any] struct {
-	queues sync.Map // 读多写少的场景，适合用 sync.Map 路由队列
+	queues         sync.Map      // 读多写少的场景，适合用 sync.Map 路由队列
+	ack            sync.Map      // 存储队列确认channel
+	msgs           sync.Map      // 存储消费消息，用于ack失败时重新入队
+	autoAck        bool          // 是否自动确认
+	expireAutoFail time.Duration // 手动确认截止时间
 }
 
 // NewQueues 创建队列集合
-func NewQueues[TKey comparable, TValue, TRoute any]() *Queues[TKey, TValue, TRoute] {
+func NewQueues[TKey comparable, TValue, TRoute any](opt ...Option) *Queues[TKey, TValue, TRoute] {
-	return &Queues[TKey, TValue, TRoute]{}
+	opts := loadOptions(opt...)
+	return &Queues[TKey, TValue, TRoute]{autoAck: opts.autoAck, expireAutoFail: opts.expireAutoFail}
+}
+
+// SetAutoAck 设置自动确认
+func (q *Queues[TKey, TValue, TRoute]) SetAutoAck(autoAck bool) {
+	q.autoAck = autoAck
+}
+
+// SetExpireAutoFail 设置手动确认截止时间
+func (q *Queues[TKey, TValue, TRoute]) SetExpireAutoFail(expireAutoFail time.Duration) {
+	q.expireAutoFail = expireAutoFail
+}
+
+// checkAckStatus 检查确认状态
+func (q *Queues[TKey, TValue, TRoute]) checkAckStatus(route TRoute) {
+	if ack, ok := q.ack.Load(route); ok && !q.autoAck {
+		if status, exist := <-ack.(chan bool); !status && exist {
+			if msg, ok := q.msgs.Load(route); ok {
+				// 重新入队
+				q.Enqueue(route, msg.(TValue))
+			}
+		}
+	}
 }
 
 // Enqueue 入队
 func (q *Queues[TKey, TValue, TRoute]) Enqueue(route TRoute, value TValue) {
+	// 入队时检查上一次出队是否确认成功
+	q.checkAckStatus(route)
+
 	if queue, ok := q.queues.Load(route); ok {
 		queue.(*DelayLkQueue[TKey, TValue]).Enqueue(value)
 	} else {
@@ -27,14 +58,57 @@ func (q *Queues[TKey, TValue, TRoute]) Enqueue(route TRoute, value TValue) {
 	}
 }
 
+// makeAckData 生成确认数据，并持续监听
+func (q *Queues[TKey, TValue, TRoute]) makeAckData(route TRoute, value TValue, valid bool) {
+	if !q.autoAck && valid {
+		var ack chan bool
+		// ack数据只需生成一次
+		if a, ok := q.ack.Load(route); !ok {
+			ack = make(chan bool, 1)
+			q.ack.Store(route, ack)
+		} else {
+			ack = a.(chan bool)
+		}
+
+		q.msgs.Store(route, value)
+
+		// todo 可考虑使用协程池
+		go func() {
+			if queue, ok := q.queues.Load(route); ok {
+				select {
+				case status := <-ack:
+					if !status {
+						// 重新入队
+						queue.(*DelayLkQueue[TKey, TValue]).Enqueue(value)
+					}
+
+				case <-time.After(q.expireAutoFail):
+					// 重新入队
+					fmt.Println("ack timeout")
+					queue.(*DelayLkQueue[TKey, TValue]).Enqueue(value)
+				}
+			}
+		}()
+	}
+}
+
 // Dequeue 出队
-func (q *Queues[TKey, TValue, TRoute]) Dequeue(route TRoute) (value TValue, ok bool) {
+func (q *Queues[TKey, TValue, TRoute]) Dequeue(route TRoute) (value TValue, valid bool) {
 	if queue, ok := q.queues.Load(route); ok {
-		return queue.(*DelayLkQueue[TKey, TValue]).Dequeue()
+		value, valid = queue.(*DelayLkQueue[TKey, TValue]).Dequeue()
+		q.makeAckData(route, value, valid)
+		return value, valid
 	}
 	return value, false
 }
 
+// Ack 确认消息
+func (q *Queues[TKey, TValue, TRoute]) Ack(route TRoute, status bool) {
+	if ack, ok := q.ack.Load(route); ok && !q.autoAck {
+		ack.(chan bool) <- status
+	}
+}
+
 // Len 队列长度
 func (q *Queues[TKey, TValue, TRoute]) Len(route TRoute) int {
 	if queue, ok := q.queues.Load(route); ok {
@@ -45,6 +119,9 @@ func (q *Queues[TKey, TValue, TRoute]) Len(route TRoute) int {
 
 // DelayEnqueue 延迟入队
 func (q *Queues[TKey, TValue, TRoute]) DelayEnqueue(route TRoute, value TValue, duration time.Duration) {
+	// 入队时检查上一次出队是否确认成功
+	q.checkAckStatus(route)
+
 	if queue, ok := q.queues.Load(route); ok {
 		queue.(*DelayLkQueue[TKey, TValue]).DelayEnqueue(value, duration)
 	} else {
@@ -56,6 +133,9 @@ func (q *Queues[TKey, TValue, TRoute]) DelayEnqueue(route TRoute, value TValue,
 
 // CancellableDelayEnqueue 可取消的延迟入队
 func (q *Queues[TKey, TValue, TRoute]) CancellableDelayEnqueue(route TRoute, key TKey, value TValue, duration time.Duration) {
+	// 入队时检查上一次出队是否确认成功
+	q.checkAckStatus(route)
+
 	if queue, ok := q.queues.Load(route); ok {
 		queue.(*DelayLkQueue[TKey, TValue]).CancellableDelayEnqueue(key, value, duration)
 	} else {
@@ -75,21 +155,31 @@ func (q *Queues[TKey, TValue, TRoute]) CancelDelayEnqueue(route TRoute, key TKey
 // ContinuousDequeue 持续监听出队
 func (q *Queues[TKey, TValue, TRoute]) ContinuousDequeue(route TRoute) iter.Seq[TValue] {
 	if queue, ok := q.queues.Load(route); ok {
-		return queue.(*DelayLkQueue[TKey, TValue]).ContinuousDequeue()
+		// 因需增加确认状态，故重新封装
+		return func(yield func(TValue) bool) {
+			for value := range queue.(*DelayLkQueue[TKey, TValue]).ContinuousDequeue() {
+				q.makeAckData(route, value, true)
+				if !yield(value) {
+					return
+				}
+			}
+		}
 	}
 	return nil
 }
 
 // ContinuousDequeueExecute 持续监听出队执行函数
 func (q *Queues[TKey, TValue, TRoute]) ContinuousDequeueExecute(route TRoute, fn func(TValue)) {
-	if queue, ok := q.queues.Load(route); ok {
+	for value := range q.ContinuousDequeue(route) {
-		queue.(*DelayLkQueue[TKey, TValue]).ContinuousDequeueExecute(fn)
+		fn(value)
 	}
 }
 
 // ContinuousDequeueNotify 持续监听出队通知
 func (q *Queues[TKey, TValue, TRoute]) ContinuousDequeueNotify(route TRoute, chs ...chan TValue) {
-	if queue, ok := q.queues.Load(route); ok {
+	q.ContinuousDequeueExecute(route, func(value TValue) {
-		queue.(*DelayLkQueue[TKey, TValue]).ContinuousDequeueNotify(chs...)
+		for _, ch := range chs {
-	}
+			ch <- value
+		}
+	})
 }

lock_free/queues_ack_test.go

@@ -0,0 +1,38 @@
+package lock_free
+
+import (
+	"testing"
+	"time"
+)
+
+func TestQueuesAck(t *testing.T) {
+	cases := []int{1, 2, 3}
+	route := "test_ack"
+	q := NewQueues[struct{}, int, string](
+		WithAutoAck(false),
+		WithExpireAutoFail(time.Second),
+	)
+
+	for _, c := range cases {
+		q.Enqueue(route, c)
+	}
+
+	for range cases {
+		q.Dequeue(route)
+	}
+
+	time.Sleep(2 * time.Second)
+
+	if q.Len(route) != len(cases) {
+		t.Errorf("queue length error, want %d, got %d", len(cases), q.Len(route))
+	}
+
+	for range cases {
+		q.Dequeue(route)
+		q.Ack(route, true)
+	}
+
+	if q.Len(route) != 0 {
+		t.Errorf("queue length error, want %d, got %d", 0, q.Len(route))
+	}
+}