Add capability for memory and persistent queue to block when add items

Signed-off-by: Bogdan Drutu <[email protected]>

exporter/exporterqueue/bounded_memory_queue.go

@@ -23,17 +23,18 @@ type boundedMemoryQueue[T any] struct {
 type memoryQueueSettings[T any] struct {
 	sizer    sizer[T]
 	capacity int64
+	blocking bool
 }
 
 // newBoundedMemoryQueue constructs the new queue of specified capacity, and with an optional
 // callback for dropped items (e.g. useful to emit metrics).
 func newBoundedMemoryQueue[T any](set memoryQueueSettings[T]) Queue[T] {
 	return &boundedMemoryQueue[T]{
-		sizedQueue: newSizedQueue[T](set.capacity, set.sizer),
+		sizedQueue: newSizedQueue[T](set.capacity, set.sizer, set.blocking),
 	}
 }
 
-func (q *boundedMemoryQueue[T]) Read(_ context.Context) (uint64, context.Context, T, bool) {
+func (q *boundedMemoryQueue[T]) Read(context.Context) (uint64, context.Context, T, bool) {
 	ctx, req, ok := q.sizedQueue.pop()
 	return 0, ctx, req, ok
 }

exporter/exporterqueue/bounded_memory_queue_test.go

@@ -133,6 +133,48 @@ func TestQueueUsage(t *testing.T) {
 	}
 }
 
+func TestBlockingQueueUsage(t *testing.T) {
+	tests := []struct {
+		name  string
+		sizer sizer[ptrace.Traces]
+	}{
+		{
+			name:  "requests_based",
+			sizer: &requestSizer[ptrace.Traces]{},
+		},
+		{
+			name:  "items_based",
+			sizer: &itemsSizer{},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			q := newBoundedMemoryQueue[ptrace.Traces](memoryQueueSettings[ptrace.Traces]{sizer: tt.sizer, capacity: int64(100), blocking: true})
+			consumed := &atomic.Int64{}
+			require.NoError(t, q.Start(context.Background(), componenttest.NewNopHost()))
+			ac := newAsyncConsumer(q, 10, func(context.Context, ptrace.Traces) error {
+				consumed.Add(1)
+				return nil
+			})
+			td := testdata.GenerateTraces(10)
+			wg := &sync.WaitGroup{}
+			for i := 0; i < 10; i++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					for j := 0; j < 100_000; j++ {
+						require.NoError(t, q.Offer(context.Background(), td))
+					}
+				}()
+			}
+			wg.Wait()
+			assert.NoError(t, q.Shutdown(context.Background()))
+			assert.NoError(t, ac.Shutdown(context.Background()))
+			assert.Equal(t, int64(1_000_000), consumed.Load())
+		})
+	}
+}
+
 func TestZeroSizeNoConsumers(t *testing.T) {
 	q := newBoundedMemoryQueue[string](memoryQueueSettings[string]{sizer: &requestSizer[string]{}, capacity: 0})
 
@@ -149,8 +191,7 @@ func consume[T any](q Queue[T], consumeFunc func(context.Context, T) error) bool
 	if !ok {
 		return false
 	}
-	consumeErr := consumeFunc(ctx, req)
-	q.OnProcessingFinished(index, consumeErr)
+	q.OnProcessingFinished(index, consumeFunc(ctx, req))
 	return true
 }
 

exporter/exporterqueue/cond.go

@@ -0,0 +1,59 @@
+package exporterqueue
+
+import (
+	"context"
+	"sync"
+)
+
+// cond is equivalent with sync.Cond, but context.Context aware. Which means Wait() will return if context is done.
+// Also, it requires the caller to hold the c.L during all calls.
+type cond struct {
+	L       sync.Locker
+	ch      chan struct{}
+	waiting int64
+}
+
+func newCond(l sync.Locker) *cond {
+	return &cond{L: l, ch: make(chan struct{}, 1)}
+}
+
+// Signal wakes one goroutine waiting on c, if there is any.
+// It requires for the caller to hold c.L during the call.
+func (c *cond) Signal() {
+	if c.waiting == 0 {
+		return
+	}
+	c.waiting--
+	c.ch <- struct{}{}
+}
+
+// Broadcast wakes all goroutines waiting on c.
+// It requires for the caller to hold c.L during the call.
+func (c *cond) Broadcast() {
+	for ; c.waiting > 0; c.waiting-- {
+		c.ch <- struct{}{}
+	}
+}
+
+// Wait atomically unlocks c.L and suspends execution of the calling goroutine. After later resuming execution, Wait locks c.L before returning.
+func (c *cond) Wait(ctx context.Context) error {
+	c.waiting++
+	c.L.Unlock()
+	select {
+	case <-ctx.Done():
+		c.L.Lock()
+		if c.waiting == 0 {
+			// If waiting is 0, it means that there was a signal sent and nobody else waits for it.
+			// Consume it, so that we don't unblock other consumer unnecessary,
+			// or we don't block the producer because the channel buffer is full.
+			<-c.ch
+		} else {
+			// Decrease the number of waiting routines.
+			c.waiting--
+		}
+		return ctx.Err()
+	case <-c.ch:
+		c.L.Lock()
+		return nil
+	}
+}

exporter/exporterqueue/persistent_queue.go

@@ -11,7 +11,6 @@ import (
 	"strconv"
 	"sync"
 
-	"go.uber.org/multierr"
 	"go.uber.org/zap"
 
 	"go.opentelemetry.io/collector/component"
@@ -42,6 +41,7 @@ var (
 type persistentQueueSettings[T any] struct {
 	sizer       sizer[T]
 	capacity    int64
+	blocking    bool
 	signal      pipeline.Signal
 	storageID   component.ID
 	marshaler   Marshaler[T]
@@ -78,10 +78,12 @@ type persistentQueue[T any] struct {
 
 	// isRequestSized indicates whether the queue is sized by the number of requests.
 	isRequestSized bool
+	blocking       bool
 
 	// mu guards everything declared below.
 	mu                       sync.Mutex
-	hasElements              *sync.Cond
+	hasMoreElements          *cond
+	hasMoreSpace             *cond
 	readIndex                uint64
 	writeIndex               uint64
 	currentlyDispatchedItems []uint64
@@ -94,11 +96,13 @@ type persistentQueue[T any] struct {
 func newPersistentQueue[T any](set persistentQueueSettings[T]) Queue[T] {
 	_, isRequestSized := set.sizer.(*requestSizer[T])
 	pq := &persistentQueue[T]{
-		set:            set,
 		logger:         set.set.Logger,
+		set:            set,
 		isRequestSized: isRequestSized,
+		blocking:       set.blocking,
 	}
-	pq.hasElements = sync.NewCond(&pq.mu)
+	pq.hasMoreElements = newCond(&pq.mu)
+	pq.hasMoreSpace = newCond(&pq.mu)
 	return pq
 }
 
@@ -194,8 +198,8 @@ func (pq *persistentQueue[T]) Shutdown(ctx context.Context) error {
 	backupErr := pq.backupQueueSize(ctx)
 	// Mark this queue as stopped, so consumer don't start any more work.
 	pq.stopped = true
-	pq.hasElements.Broadcast()
-	return multierr.Combine(backupErr, pq.unrefClient(ctx))
+	pq.hasMoreElements.Broadcast()
+	return errors.Join(backupErr, pq.unrefClient(ctx))
 }
 
 // backupQueueSize writes the current queue size to storage. The value is used to recover the queue size
@@ -233,8 +237,13 @@ func (pq *persistentQueue[T]) Offer(ctx context.Context, req T) error {
 // putInternal is the internal version that requires caller to hold the mutex lock.
 func (pq *persistentQueue[T]) putInternal(ctx context.Context, req T) error {
 	reqSize := pq.set.sizer.Sizeof(req)
-	if pq.queueSize+reqSize > pq.set.capacity {
-		return ErrQueueIsFull
+	for pq.queueSize+reqSize > pq.set.capacity {
+		if !pq.blocking {
+			return ErrQueueIsFull
+		}
+		if err := pq.hasMoreSpace.Wait(ctx); err != nil {
+			return err
+		}
 	}
 
 	reqBuf, err := pq.set.marshaler(req)
@@ -253,7 +262,7 @@ func (pq *persistentQueue[T]) putInternal(ctx context.Context, req T) error {
 
 	pq.writeIndex++
 	pq.queueSize += reqSize
-	pq.hasElements.Signal()
+	pq.hasMoreElements.Signal()
 
 	// Back up the queue size to storage every 10 writes. The stored value is used to recover the queue size
 	// in case if the collector is killed. The recovered queue size is allowed to be inaccurate.
@@ -269,32 +278,38 @@ func (pq *persistentQueue[T]) putInternal(ctx context.Context, req T) error {
 func (pq *persistentQueue[T]) Read(ctx context.Context) (uint64, context.Context, T, bool) {
 	pq.mu.Lock()
 	defer pq.mu.Unlock()
-	for {
-		if pq.stopped {
-			var req T
-			return 0, context.Background(), req, false
-		}
-
-		// If queue is empty, wait until more elements and restart.
-		if pq.readIndex == pq.writeIndex {
-			pq.hasElements.Wait()
-			continue
-		}
 
-		index, req, consumed := pq.getNextItem(ctx)
-		if consumed {
-			pq.queueSize -= pq.set.sizer.Sizeof(req)
-			// The size might be not in sync with the queue in case it's restored from the disk
-			// because we don't flush the current queue size on the disk on every read/write.
-			// In that case we need to make sure it doesn't go below 0.
-			if pq.queueSize < 0 {
+	for {
+		// Read until either a successful retrieved element or no more elements in the storage.
+		for pq.readIndex != pq.writeIndex {
+			index, req, consumed := pq.getNextItem(ctx)
+			// Ensure the used size and the channel size are in sync.
+			if pq.readIndex == pq.writeIndex {
 				pq.queueSize = 0
+				pq.hasMoreSpace.Signal()
 			}
+			if consumed {
+				pq.queueSize -= pq.set.sizer.Sizeof(req)
+				// The size might be not in sync with the queue in case it's restored from the disk
+				// because we don't flush the current queue size on the disk on every read/write.
+				// In that case we need to make sure it doesn't go below 0.
+				if pq.queueSize < 0 {
+					pq.queueSize = 0
+				}
+				pq.hasMoreSpace.Signal()
+
+				return index, context.Background(), req, true
+			}
+		}
 
-			return index, context.Background(), req, true
+		if pq.stopped {
+			var req T
+			return 0, context.Background(), req, false
 		}
 
-		// If we did not consume any element retry from the beginning.
+		// TODO: Change the Queue interface to return an error to allow distinguish between shutdown and context canceled.
+		// Ok to ignore the error, since the context.Background() will never be done.
+		_ = pq.hasMoreElements.Wait(context.Background())
 	}
 }
 
@@ -363,11 +378,6 @@ func (pq *persistentQueue[T]) OnProcessingFinished(index uint64, consumeErr erro
 			pq.logger.Error("Error writing queue size to storage", zap.Error(qsErr))
 		}
 	}
-
-	// Ensure the used size and the channel size are in sync.
-	if pq.readIndex == pq.writeIndex {
-		pq.queueSize = 0
-	}
 }
 
 // retrieveAndEnqueueNotDispatchedReqs gets the items for which sending was not finished, cleans the storage

exporter/exporterqueue/persistent_queue_test.go

@@ -370,12 +370,10 @@ func TestPersistentQueue_ConsumersProducers(t *testing.T) {
 		t.Run(fmt.Sprintf("#messages: %d #consumers: %d", c.numMessagesProduced, c.numConsumers), func(t *testing.T) {
 			req := newTracesRequest(1, 10)
 
-			numMessagesConsumed := &atomic.Int32{}
+			numMessagesConsumed := &atomic.Int64{}
 			pq := createAndStartTestPersistentQueue(t, &requestSizer[ptrace.Traces]{}, 1000, c.numConsumers,
-				func(context.Context,
-					ptrace.Traces,
-				) error {
-					numMessagesConsumed.Add(int32(1))
+				func(context.Context, ptrace.Traces) error {
+					numMessagesConsumed.Add(int64(1))
 					return nil
 				})
 
@@ -390,6 +388,63 @@ func TestPersistentQueue_ConsumersProducers(t *testing.T) {
 	}
 }
 
+func TestPersistentBlockingQueue(t *testing.T) {
+	tests := []struct {
+		name  string
+		sizer sizer[ptrace.Traces]
+	}{
+		{
+			name:  "requests_based",
+			sizer: &requestSizer[ptrace.Traces]{},
+		},
+		{
+			name:  "items_based",
+			sizer: &itemsSizer{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			pq := newPersistentQueue[ptrace.Traces](persistentQueueSettings[ptrace.Traces]{
+				sizer:       tt.sizer,
+				capacity:    100,
+				blocking:    true,
+				signal:      pipeline.SignalTraces,
+				storageID:   component.ID{},
+				marshaler:   marshalTracesRequest,
+				unmarshaler: unmarshalTracesRequest,
+				set:         exportertest.NewNopSettings(),
+			})
+			host := &mockHost{ext: map[component.ID]component.Component{
+				{}: storagetest.NewMockStorageExtension(nil),
+			}}
+			require.NoError(t, pq.Start(context.Background(), host))
+			consumed := &atomic.Int64{}
+			ac := newAsyncConsumer(pq, 10, func(context.Context, ptrace.Traces) error {
+				consumed.Add(int64(1))
+				return nil
+			})
+
+			td := newTracesRequest(1, 10)
+			wg := &sync.WaitGroup{}
+			for i := 0; i < 10; i++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					for j := 0; j < 100_000; j++ {
+						require.NoError(t, pq.Offer(context.Background(), td))
+					}
+				}()
+			}
+			wg.Wait()
+			assert.NoError(t, pq.Shutdown(context.Background()))
+			assert.NoError(t, ac.Shutdown(context.Background()))
+			assert.Equal(t, int64(1_000_000), consumed.Load())
+
+		})
+	}
+}
+
 func newTracesRequest(numTraces int, numSpans int) ptrace.Traces {
 	traces := ptrace.NewTraces()
 	batch := traces.ResourceSpans().AppendEmpty()