diff --git a/modules/queue/backoff.go b/modules/queue/backoff.go index 02fc88574a6..c44c06ac18b 100644 --- a/modules/queue/backoff.go +++ b/modules/queue/backoff.go @@ -13,10 +13,18 @@ var ( backoffUpper = 2 * time.Second ) -type ( - backoffFuncRetErr[T any] func() (retry bool, ret T, err error) - backoffFuncErr func() (retry bool, err error) -) +type backoffFunc[T any] func() (retry bool, ret T, err error) + +type backoffOptions struct { + begin, upper time.Duration + + notify <-chan struct{} + end <-chan time.Time +} + +func backoffOptionsDefault(notify <-chan struct{}, end <-chan time.Time) backoffOptions { + return backoffOptions{begin: backoffBegin, upper: backoffUpper, notify: notify, end: end} +} func mockBackoffDuration(d time.Duration) func() { oldBegin, oldUpper := backoffBegin, backoffUpper @@ -26,18 +34,9 @@ func mockBackoffDuration(d time.Duration) func() { } } -func backoffRetErr[T any](ctx context.Context, begin, upper time.Duration, end <-chan time.Time, fn backoffFuncRetErr[T]) (ret T, err error) { - d := begin +func backoffCall[T any](ctx context.Context, opts backoffOptions, fn backoffFunc[T]) (ret T, err error) { + d := opts.begin for { - // check whether the context has been cancelled or has reached the deadline, return early - select { - case <-ctx.Done(): - return ret, ctx.Err() - case <-end: - return ret, context.DeadlineExceeded - default: - } - // call the target function retry, ret, err := fn() if err != nil { @@ -47,25 +46,19 @@ func backoffRetErr[T any](ctx context.Context, begin, upper time.Duration, end < return ret, nil } - // wait for a while before retrying, and also respect the context & deadline + // wait for a while before retrying, and also respect the context & deadline & notify select { case <-ctx.Done(): return ret, ctx.Err() + case <-opts.end: + return ret, context.DeadlineExceeded + case <-opts.notify: + continue case <-time.After(d): d *= 2 - if d > upper { - d = upper + if d > opts.upper { + d = opts.upper } - case <-end: - return ret, context.DeadlineExceeded } } } - -func backoffErr(ctx context.Context, begin, upper time.Duration, end <-chan time.Time, fn backoffFuncErr) error { - _, err := backoffRetErr(ctx, begin, upper, end, func() (retry bool, ret any, err error) { - retry, err = fn() - return retry, nil, err - }) - return err -} diff --git a/modules/queue/base.go b/modules/queue/base.go index 102e79e5416..46033ca2735 100644 --- a/modules/queue/base.go +++ b/modules/queue/base.go @@ -40,3 +40,28 @@ func popItemByChan(ctx context.Context, popItemFn func(ctx context.Context) ([]b }() return chanItem, chanErr } + +type baseQueueNotifiableInterface interface { + getNotifySignalChan() chan struct{} +} + +type baseQueueNotifiable struct { + notifySignal chan struct{} +} + +var _ baseQueueNotifiableInterface = (*baseQueueNotifiable)(nil) + +func (n *baseQueueNotifiable) notifyPushItem() { + select { + case n.notifySignal <- struct{}{}: + default: + } +} + +func (n *baseQueueNotifiable) getNotifySignalChan() chan struct{} { + return n.notifySignal +} + +func newBaseQueueNotifiable() *baseQueueNotifiable { + return &baseQueueNotifiable{notifySignal: make(chan struct{}, 1)} +} diff --git a/modules/queue/base_channel_test.go b/modules/queue/base_channel_test.go index 5d0a2ed0a7c..cee864657f7 100644 --- a/modules/queue/base_channel_test.go +++ b/modules/queue/base_channel_test.go @@ -6,6 +6,6 @@ package queue import "testing" func TestBaseChannel(t *testing.T) { - testQueueBasic(t, newBaseChannelSimple, &BaseConfig{ManagedName: "baseChannel", Length: 10}, false) - testQueueBasic(t, newBaseChannelUnique, &BaseConfig{ManagedName: "baseChannel", Length: 10}, true) + testQueueBasic(t, newBaseChannelSimple, &BaseConfig{ManagedName: "baseChannel", Length: 10}, testQueueBasicOptions{}) + testQueueBasic(t, newBaseChannelUnique, &BaseConfig{ManagedName: "baseChannel", Length: 10}, testQueueBasicOptions{UniqueQueue: true}) } diff --git a/modules/queue/base_levelqueue.go b/modules/queue/base_levelqueue.go index 72295ce3a0f..24f382913a0 100644 --- a/modules/queue/base_levelqueue.go +++ b/modules/queue/base_levelqueue.go @@ -15,6 +15,7 @@ import ( ) type baseLevelQueue struct { + *baseLevelQueueCommonImpl internal atomic.Pointer[levelqueue.Queue] conn string @@ -22,7 +23,10 @@ type baseLevelQueue struct { db *leveldb.DB } -var _ baseQueue = (*baseLevelQueue)(nil) +var ( + _ baseQueue = (*baseLevelQueue)(nil) + _ baseQueueNotifiableInterface = (*baseLevelQueue)(nil) +) func newBaseLevelQueueGeneric(cfg *BaseConfig, unique bool) (baseQueue, error) { if unique { @@ -42,19 +46,10 @@ func newBaseLevelQueueSimple(cfg *BaseConfig) (baseQueue, error) { return nil, err } q.internal.Store(lq) + q.baseLevelQueueCommonImpl = baseLevelQueueCommon(q.cfg, nil, func() baseLevelQueuePushPoper { return q.internal.Load() }) return q, nil } -func (q *baseLevelQueue) PushItem(ctx context.Context, data []byte) error { - c := baseLevelQueueCommon(q.cfg, nil, func() baseLevelQueuePushPoper { return q.internal.Load() }) - return c.PushItem(ctx, data) -} - -func (q *baseLevelQueue) PopItem(ctx context.Context) ([]byte, error) { - c := baseLevelQueueCommon(q.cfg, nil, func() baseLevelQueuePushPoper { return q.internal.Load() }) - return c.PopItem(ctx) -} - func (q *baseLevelQueue) HasItem(ctx context.Context, data []byte) (bool, error) { return false, nil } diff --git a/modules/queue/base_levelqueue_common.go b/modules/queue/base_levelqueue_common.go index 5ac899cc189..9d5e578d8c3 100644 --- a/modules/queue/base_levelqueue_common.go +++ b/modules/queue/base_levelqueue_common.go @@ -25,35 +25,40 @@ type baseLevelQueuePushPoper interface { } type baseLevelQueueCommonImpl struct { + *baseQueueNotifiable length int internalFunc func() baseLevelQueuePushPoper - mu *sync.Mutex + muCommon *sync.Mutex } func (q *baseLevelQueueCommonImpl) PushItem(ctx context.Context, data []byte) error { - return backoffErr(ctx, backoffBegin, backoffUpper, time.After(pushBlockTime), func() (retry bool, err error) { - if q.mu != nil { - q.mu.Lock() - defer q.mu.Unlock() + _, err := backoffCall(ctx, backoffOptionsDefault(noNotifyChan, time.After(pushBlockTime)), func() (retry bool, ret any, err error) { + if q.muCommon != nil { + q.muCommon.Lock() + defer q.muCommon.Unlock() } cnt := int(q.internalFunc().Len()) if cnt >= q.length { - return true, nil + return true, nil, nil } retry, err = false, q.internalFunc().RPush(data) if err == levelqueue.ErrAlreadyInQueue { err = ErrAlreadyInQueue } - return retry, err + if err == nil { + q.notifyPushItem() + } + return retry, nil, err }) + return err } func (q *baseLevelQueueCommonImpl) PopItem(ctx context.Context) ([]byte, error) { - return backoffRetErr(ctx, backoffBegin, backoffUpper, infiniteTimerC, func() (retry bool, data []byte, err error) { - if q.mu != nil { - q.mu.Lock() - defer q.mu.Unlock() + return backoffCall(ctx, backoffOptionsDefault(q.notifySignal, infiniteTimerC), func() (retry bool, data []byte, err error) { + if q.muCommon != nil { + q.muCommon.Lock() + defer q.muCommon.Unlock() } data, err = q.internalFunc().LPop() @@ -68,7 +73,7 @@ func (q *baseLevelQueueCommonImpl) PopItem(ctx context.Context) ([]byte, error) } func baseLevelQueueCommon(cfg *BaseConfig, mu *sync.Mutex, internalFunc func() baseLevelQueuePushPoper) *baseLevelQueueCommonImpl { - return &baseLevelQueueCommonImpl{length: cfg.Length, mu: mu, internalFunc: internalFunc} + return &baseLevelQueueCommonImpl{length: cfg.Length, muCommon: mu, internalFunc: internalFunc, baseQueueNotifiable: newBaseQueueNotifiable()} } func prepareLevelDB(cfg *BaseConfig) (conn string, db *leveldb.DB, err error) { diff --git a/modules/queue/base_levelqueue_test.go b/modules/queue/base_levelqueue_test.go index afcb5f67902..7c4166bc5e3 100644 --- a/modules/queue/base_levelqueue_test.go +++ b/modules/queue/base_levelqueue_test.go @@ -22,8 +22,10 @@ func TestBaseLevelDB(t *testing.T) { _, err = newBaseLevelQueueGeneric(&BaseConfig{DataFullDir: "relative"}, false) assert.ErrorContains(t, err, "invalid leveldb data dir") - testQueueBasic(t, newBaseLevelQueueSimple, toBaseConfig("baseLevelQueue", setting.QueueSettings{Datadir: t.TempDir() + "/queue-test", Length: 10}), false) - testQueueBasic(t, newBaseLevelQueueUnique, toBaseConfig("baseLevelQueueUnique", setting.QueueSettings{ConnStr: "leveldb://" + t.TempDir() + "/queue-test", Length: 10}), true) + optsSimple := testQueueBasicOptions{NotifiableQueue: true} + optsUnique := testQueueBasicOptions{UniqueQueue: true, NotifiableQueue: true} + testQueueBasic(t, newBaseLevelQueueSimple, toBaseConfig("baseLevelQueue", setting.QueueSettings{Datadir: t.TempDir() + "/queue-test", Length: 10}), optsSimple) + testQueueBasic(t, newBaseLevelQueueUnique, toBaseConfig("baseLevelQueueUnique", setting.QueueSettings{ConnStr: "leveldb://" + t.TempDir() + "/queue-test", Length: 10}), optsUnique) } func TestCorruptedLevelQueue(t *testing.T) { diff --git a/modules/queue/base_levelqueue_unique.go b/modules/queue/base_levelqueue_unique.go index 23c76bc5aed..60de3785a15 100644 --- a/modules/queue/base_levelqueue_unique.go +++ b/modules/queue/base_levelqueue_unique.go @@ -16,16 +16,20 @@ import ( ) type baseLevelQueueUnique struct { + *baseLevelQueueCommonImpl internal atomic.Pointer[levelqueue.UniqueQueue] conn string cfg *BaseConfig db *leveldb.DB - mu sync.Mutex // the levelqueue.UniqueQueue is not thread-safe, there is no mutex protecting the underlying queue&set together + muBase sync.Mutex // the levelqueue.UniqueQueue is not thread-safe, there is no mutex protecting the underlying queue&set together } -var _ baseQueue = (*baseLevelQueueUnique)(nil) +var ( + _ baseQueue = (*baseLevelQueueUnique)(nil) + _ baseQueueNotifiableInterface = (*baseLevelQueueUnique)(nil) +) func newBaseLevelQueueUnique(cfg *BaseConfig) (baseQueue, error) { conn, db, err := prepareLevelDB(cfg) @@ -38,34 +42,25 @@ func newBaseLevelQueueUnique(cfg *BaseConfig) (baseQueue, error) { return nil, err } q.internal.Store(lq) + q.baseLevelQueueCommonImpl = baseLevelQueueCommon(q.cfg, &q.muBase, func() baseLevelQueuePushPoper { return q.internal.Load() }) return q, nil } -func (q *baseLevelQueueUnique) PushItem(ctx context.Context, data []byte) error { - c := baseLevelQueueCommon(q.cfg, &q.mu, func() baseLevelQueuePushPoper { return q.internal.Load() }) - return c.PushItem(ctx, data) -} - -func (q *baseLevelQueueUnique) PopItem(ctx context.Context) ([]byte, error) { - c := baseLevelQueueCommon(q.cfg, &q.mu, func() baseLevelQueuePushPoper { return q.internal.Load() }) - return c.PopItem(ctx) -} - func (q *baseLevelQueueUnique) HasItem(ctx context.Context, data []byte) (bool, error) { - q.mu.Lock() - defer q.mu.Unlock() + q.muBase.Lock() + defer q.muBase.Unlock() return q.internal.Load().Has(data) } func (q *baseLevelQueueUnique) Len(ctx context.Context) (int, error) { - q.mu.Lock() - defer q.mu.Unlock() + q.muBase.Lock() + defer q.muBase.Unlock() return int(q.internal.Load().Len()), nil } func (q *baseLevelQueueUnique) Close() error { - q.mu.Lock() - defer q.mu.Unlock() + q.muBase.Lock() + defer q.muBase.Unlock() err := q.internal.Load().Close() q.db = nil // the db is not managed by us, it's managed by the nosql manager _ = nosql.GetManager().CloseLevelDB(q.conn) @@ -73,8 +68,8 @@ func (q *baseLevelQueueUnique) Close() error { } func (q *baseLevelQueueUnique) RemoveAll(ctx context.Context) error { - q.mu.Lock() - defer q.mu.Unlock() + q.muBase.Lock() + defer q.muBase.Unlock() lqinternal.RemoveLevelQueueKeys(q.db, []byte(q.cfg.QueueFullName)) lqinternal.RemoveLevelQueueKeys(q.db, []byte(q.cfg.SetFullName)) lq, err := levelqueue.NewUniqueQueue(q.db, []byte(q.cfg.QueueFullName), []byte(q.cfg.SetFullName), false) diff --git a/modules/queue/base_redis.go b/modules/queue/base_redis.go index f9e1bf45801..b60f23288ba 100644 --- a/modules/queue/base_redis.go +++ b/modules/queue/base_redis.go @@ -16,6 +16,7 @@ import ( ) type baseRedis struct { + *baseQueueNotifiable client redis.UniversalClient isUnique bool cfg *BaseConfig @@ -23,7 +24,10 @@ type baseRedis struct { mu sync.Mutex // the old implementation is not thread-safe, the queue operation and set operation should be protected together } -var _ baseQueue = (*baseRedis)(nil) +var ( + _ baseQueue = (*baseRedis)(nil) + _ baseQueueNotifiableInterface = (*baseRedis)(nil) +) func newBaseRedisGeneric(cfg *BaseConfig, unique bool) (baseQueue, error) { client := nosql.GetManager().GetRedisClient(cfg.ConnStr) @@ -41,7 +45,7 @@ func newBaseRedisGeneric(cfg *BaseConfig, unique bool) (baseQueue, error) { return nil, err } - return &baseRedis{cfg: cfg, client: client, isUnique: unique}, nil + return &baseRedis{cfg: cfg, client: client, isUnique: unique, baseQueueNotifiable: newBaseQueueNotifiable()}, nil } func newBaseRedisSimple(cfg *BaseConfig) (baseQueue, error) { @@ -53,33 +57,38 @@ func newBaseRedisUnique(cfg *BaseConfig) (baseQueue, error) { } func (q *baseRedis) PushItem(ctx context.Context, data []byte) error { - return backoffErr(ctx, backoffBegin, backoffUpper, time.After(pushBlockTime), func() (retry bool, err error) { + _, err := backoffCall(ctx, backoffOptionsDefault(noNotifyChan, time.After(pushBlockTime)), func() (retry bool, ret any, err error) { q.mu.Lock() defer q.mu.Unlock() cnt, err := q.client.LLen(ctx, q.cfg.QueueFullName).Result() if err != nil { - return false, err + return false, nil, err } if int(cnt) >= q.cfg.Length { - return true, nil + return true, nil, nil } if q.isUnique { added, err := q.client.SAdd(ctx, q.cfg.SetFullName, data).Result() if err != nil { - return false, err + return false, nil, err } if added == 0 { - return false, ErrAlreadyInQueue + return false, nil, ErrAlreadyInQueue } } - return false, q.client.RPush(ctx, q.cfg.QueueFullName, data).Err() + retry, err = false, q.client.RPush(ctx, q.cfg.QueueFullName, data).Err() + if err == nil { + q.notifyPushItem() + } + return retry, nil, err }) + return err } func (q *baseRedis) PopItem(ctx context.Context) ([]byte, error) { - return backoffRetErr(ctx, backoffBegin, backoffUpper, infiniteTimerC, func() (retry bool, data []byte, err error) { + return backoffCall(ctx, backoffOptionsDefault(q.notifySignal, infiniteTimerC), func() (retry bool, data []byte, err error) { q.mu.Lock() defer q.mu.Unlock() diff --git a/modules/queue/base_redis_test.go b/modules/queue/base_redis_test.go index 329a40cfd57..14fb77c85c3 100644 --- a/modules/queue/base_redis_test.go +++ b/modules/queue/base_redis_test.go @@ -13,6 +13,8 @@ import ( func TestBaseRedis(t *testing.T) { redisConn := test.PrepareTestRedis(t) queueSetting := setting.QueueSettings{Length: 10, ConnStr: redisConn} - testQueueBasic(t, newBaseRedisSimple, toBaseConfig("baseRedis", queueSetting), false) - testQueueBasic(t, newBaseRedisUnique, toBaseConfig("baseRedisUnique", queueSetting), true) + optsSimple := testQueueBasicOptions{NotifiableQueue: true} + optsUnique := testQueueBasicOptions{UniqueQueue: true, NotifiableQueue: true} + testQueueBasic(t, newBaseRedisSimple, toBaseConfig("baseRedis", queueSetting), optsSimple) + testQueueBasic(t, newBaseRedisUnique, toBaseConfig("baseRedisUnique", queueSetting), optsUnique) } diff --git a/modules/queue/base_test.go b/modules/queue/base_test.go index 8e7c18d7407..bc9e6167a99 100644 --- a/modules/queue/base_test.go +++ b/modules/queue/base_test.go @@ -6,13 +6,20 @@ package queue import ( "context" "fmt" + "sync" "testing" "time" "github.com/stretchr/testify/assert" ) -func testQueueBasic(t *testing.T, newFn func(cfg *BaseConfig) (baseQueue, error), cfg *BaseConfig, isUnique bool) { +type testQueueBasicOptions struct { + UniqueQueue bool + NotifiableQueue bool +} + +func testQueueBasic(t *testing.T, newFn func(cfg *BaseConfig) (baseQueue, error), cfg *BaseConfig, opts testQueueBasicOptions) { + isUnique := opts.UniqueQueue t.Run(fmt.Sprintf("testQueueBasic-%s-unique:%v", cfg.ManagedName, isUnique), func(t *testing.T) { q, err := newFn(cfg) assert.NoError(t, err) @@ -85,6 +92,28 @@ func testQueueBasic(t *testing.T, newFn func(cfg *BaseConfig) (baseQueue, error) assert.ErrorIs(t, err, context.Canceled) assert.Nil(t, it) + t.Run("PushNotify", func(t *testing.T) { + defer mockBackoffDuration(5000 * time.Millisecond)() + // pop an empty queue, but it can be notified and pop the item immediately + wg := sync.WaitGroup{} + wg.Go(func() { + it, err := q.PopItem(ctx) // it should return immediately after PushItem, no "backoff" waiting + assert.NoError(t, err) + assert.Equal(t, "item-notify", string(it)) + }) + time.Sleep(10 * time.Millisecond) + err = q.PushItem(ctx, []byte("item-notify")) + wg.Wait() + if opts.NotifiableQueue { + v, _ := q.(baseQueueNotifiableInterface) + assert.Empty(t, v.getNotifySignalChan(), "notify signal should have been read") + assert.NoError(t, q.PushItem(ctx, []byte("item-dummy"))) + assert.Len(t, v.getNotifySignalChan(), 1, "notify signal should exist for newly pushed item") + _, err = q.PopItem(ctx) + assert.NoError(t, err) + } + }) + // test blocking push if queue is full for i := 0; i < cfg.Length; i++ { err = q.PushItem(ctx, fmt.Appendf(nil, "item-%d", i)) diff --git a/modules/queue/workergroup.go b/modules/queue/workergroup.go index f8f55330150..c7f734ffdd0 100644 --- a/modules/queue/workergroup.go +++ b/modules/queue/workergroup.go @@ -14,7 +14,9 @@ import ( ) var ( - infiniteTimerC = make(chan time.Time) + noNotifyChan chan struct{} + infiniteTimerC chan time.Time + batchDebounceDuration = 100 * time.Millisecond workerIdleDuration = 1 * time.Second shutdownDefaultTimeout = 2 * time.Second