fix: make local queue PopItem can be notified (#39011)

modules/queue/backoff.go

@@ -13,10 +13,18 @@ var (
 	backoffUpper = 2 * time.Second
 )
 
-type (
+type backoffFunc[T any] func() (retry bool, ret T, err error)
-	backoffFuncRetErr[T any] func() (retry bool, ret T, err error)
+
-	backoffFuncErr           func() (retry bool, 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) {
+func backoffCall[T any](ctx context.Context, opts backoffOptions, fn backoffFunc[T]) (ret T, err error) {
-	d := begin
+	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 {
+			if d > opts.upper {
-				d = 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
-}

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

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, newBaseChannelSimple, &BaseConfig{ManagedName: "baseChannel", Length: 10}, testQueueBasicOptions{})
-	testQueueBasic(t, newBaseChannelUnique, &BaseConfig{ManagedName: "baseChannel", Length: 10}, true)
+	testQueueBasic(t, newBaseChannelUnique, &BaseConfig{ManagedName: "baseChannel", Length: 10}, testQueueBasicOptions{UniqueQueue: true})
 }

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
 }

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) {
+	_, err := backoffCall(ctx, backoffOptionsDefault(noNotifyChan, time.After(pushBlockTime)), func() (retry bool, ret any, err error) {
-		if q.mu != nil {
+		if q.muCommon != nil {
-			q.mu.Lock()
+			q.muCommon.Lock()
-			defer q.mu.Unlock()
+			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) {
+	return backoffCall(ctx, backoffOptionsDefault(q.notifySignal, infiniteTimerC), func() (retry bool, data []byte, err error) {
-		if q.mu != nil {
+		if q.muCommon != nil {
-			q.mu.Lock()
+			q.muCommon.Lock()
-			defer q.mu.Unlock()
+			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) {

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)
+	optsSimple := testQueueBasicOptions{NotifiableQueue: true}
-	testQueueBasic(t, newBaseLevelQueueUnique, toBaseConfig("baseLevelQueueUnique", setting.QueueSettings{ConnStr: "leveldb://" + t.TempDir() + "/queue-test", Length: 10}), 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) {

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()
+	q.muBase.Lock()
-	defer q.mu.Unlock()
+	defer q.muBase.Unlock()
 	return q.internal.Load().Has(data)
 }
 
 func (q *baseLevelQueueUnique) Len(ctx context.Context) (int, error) {
-	q.mu.Lock()
+	q.muBase.Lock()
-	defer q.mu.Unlock()
+	defer q.muBase.Unlock()
 	return int(q.internal.Load().Len()), nil
 }
 
 func (q *baseLevelQueueUnique) Close() error {
-	q.mu.Lock()
+	q.muBase.Lock()
-	defer q.mu.Unlock()
+	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()
+	q.muBase.Lock()
-	defer q.mu.Unlock()
+	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)

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()
 

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)
+	optsSimple := testQueueBasicOptions{NotifiableQueue: true}
-	testQueueBasic(t, newBaseRedisUnique, toBaseConfig("baseRedisUnique", queueSetting), true)
+	optsUnique := testQueueBasicOptions{UniqueQueue: true, NotifiableQueue: true}
+	testQueueBasic(t, newBaseRedisSimple, toBaseConfig("baseRedis", queueSetting), optsSimple)
+	testQueueBasic(t, newBaseRedisUnique, toBaseConfig("baseRedisUnique", queueSetting), optsUnique)
 }

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

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