fantasticbin/go-study
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fantasticbin 2024-12-10 12:51:17 +08:00
parent c9632e7ab1
commit 1647f91ba4
3 changed files with 174 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
lock_free/option.go Normal file
View File

@ -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
}
}

112
lock_free/queues.go
View File

@ -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] {
return &Queues[TKey, TValue, TRoute]{}
func NewQueues[TKey comparable, TValue, TRoute any](opt ...Option) *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 {
queue.(*DelayLkQueue[TKey, TValue]).ContinuousDequeueExecute(fn)
for value := range q.ContinuousDequeue(route) {
fn(value)
}
}
// ContinuousDequeueNotify 持续监听出队通知
func (q *Queues[TKey, TValue, TRoute]) ContinuousDequeueNotify(route TRoute, chs ...chan TValue) {
if queue, ok := q.queues.Load(route); ok {
queue.(*DelayLkQueue[TKey, TValue]).ContinuousDequeueNotify(chs...)
}
q.ContinuousDequeueExecute(route, func(value TValue) {
for _, ch := range chs {
ch <- value
}
})
}

38
lock_free/queues_ack_test.go Normal file
View File

@ -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))
}
}