Reduce writer lock contention + unify the writer/medium queue

channel_medium.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/centrifugal/centrifuge/internal/queue"
 	"github.com/centrifugal/centrifuge/internal/timers"
 )
 
@@ -51,7 +52,6 @@ func (o ChannelMediumOptions) isMediumEnabled() bool {
 	return o.SharedPositionSync || o.KeepLatestPublication || o.enableQueue || o.broadcastDelay > 0
 }
 
-
 // channelMedium is initialized when first subscriber comes into channel, and dropped as soon as last
 // subscriber leaves the channel on the Node.
 type channelMedium struct {
@@ -63,7 +63,7 @@ type channelMedium struct {
 	mu      sync.RWMutex
 	closeCh chan struct{}
 	// optional queue for publications.
-	messages *publicationQueue
+	messages *queue.Queue[queuedPublication]
 	// We must synchronize broadcast method between general publications and insufficient state notifications.
 	// Only used when queue is disabled.
 	broadcastMu sync.Mutex
@@ -96,7 +96,7 @@ func newChannelMedium(channel string, node nodeSubset, options ChannelMediumOpti
 	}
 	c.positionCheckTime = c.nowFn().UnixNano()
 	if options.enableQueue {
-		c.messages = newPublicationQueue(2)
+		c.messages = queue.New(2, queuedPublicationSize)
 		go c.writer()
 	}
 	return c, nil
@@ -282,12 +282,13 @@ func (c *channelMedium) checkPositionOnce(historyMetaTTL time.Duration, clientPo
 
 func (c *channelMedium) close() {
 	close(c.closeCh)
-	// Unblock the writer goroutine. publicationQueue.Wait sleeps on a
-	// sync.Cond and is woken only by an Add (cnt > 0) or Close (broadcast).
-	// closeCh is checked only inside the broadcastDelay timer branch, which
-	// is unreachable until Wait returns — so without closing the queue the
-	// writer goroutine sits forever on an empty channel and leaks for every
-	// channelMedium that ever existed.
+	// Unblock the writer goroutine. queue.Queue.Wait parks on its own
+	// internal closeCh + notify chan and is woken only by an Add
+	// (cnt > 0) or by the queue's Close. channelMedium's own closeCh
+	// is checked only inside the broadcastDelay timer branch, which is
+	// unreachable until Wait returns — so without closing the queue
+	// the writer goroutine sits forever on an empty queue and leaks
+	// for every channelMedium that ever existed.
 	if c.messages != nil {
 		c.messages.Close()
 	}
@@ -297,148 +298,12 @@ type queuedPublication struct {
 	Publication queuedPub
 }
 
-// publicationQueue is an unbounded queue of queuedPublication.
-// The queue is goroutine safe.
-// Inspired by http://blog.dubbelboer.com/2015/04/25/go-faster-queue.html (MIT)
-type publicationQueue struct {
-	mu      sync.RWMutex
-	cond    *sync.Cond
-	nodes   []queuedPublication
-	head    int
-	tail    int
-	cnt     int
-	size    int
-	closed  bool
-	initCap int
-}
-
-// newPublicationQueue returns a new queuedPublication queue with initial capacity.
-func newPublicationQueue(initialCapacity int) *publicationQueue {
-	sq := &publicationQueue{
-		initCap: initialCapacity,
-		nodes:   make([]queuedPublication, initialCapacity),
-	}
-	sq.cond = sync.NewCond(&sq.mu)
-	return sq
-}
-
-// Mutex must be held when calling.
-func (q *publicationQueue) resize(n int) {
-	nodes := make([]queuedPublication, n)
-	if q.head < q.tail {
-		copy(nodes, q.nodes[q.head:q.tail])
-	} else {
-		copy(nodes, q.nodes[q.head:])
-		copy(nodes[len(q.nodes)-q.head:], q.nodes[:q.tail])
-	}
-
-	q.tail = q.cnt % n
-	q.head = 0
-	q.nodes = nodes
-}
-
-// Add an queuedPublication to the back of the queue
-// will return false if the queue is closed.
-// In that case the queuedPublication is dropped.
-func (q *publicationQueue) Add(i queuedPublication) bool {
-	q.mu.Lock()
-	if q.closed {
-		q.mu.Unlock()
-		return false
-	}
-	if q.cnt == len(q.nodes) {
-		// Also tested a growth rate of 1.5, see: http://stackoverflow.com/questions/2269063/buffer-growth-strategy
-		// In Go this resulted in a higher memory usage.
-		q.resize(q.cnt * 2)
-	}
-	q.nodes[q.tail] = i
-	q.tail = (q.tail + 1) % len(q.nodes)
-	if i.Publication.pub != nil {
-		q.size += len(i.Publication.pub.Data)
-	}
-	q.cnt++
-	q.cond.Signal()
-	q.mu.Unlock()
-	return true
-}
-
-// Close the queue and discard all entries in the queue
-// all goroutines in wait() will return
-func (q *publicationQueue) Close() {
-	q.mu.Lock()
-	defer q.mu.Unlock()
-	q.closed = true
-	q.cnt = 0
-	q.nodes = nil
-	q.size = 0
-	q.cond.Broadcast()
-}
-
-// Closed returns true if the queue has been closed
-// The call cannot guarantee that the queue hasn't been
-// closed while the function returns, so only "true" has a definite meaning.
-func (q *publicationQueue) Closed() bool {
-	q.mu.RLock()
-	c := q.closed
-	q.mu.RUnlock()
-	return c
-}
-
-// Wait for a message to be added.
-// If there are items on the queue will return immediately.
-// Will return false if the queue is closed.
-// Otherwise, returns true.
-func (q *publicationQueue) Wait() bool {
-	q.mu.Lock()
-	if q.closed {
-		q.mu.Unlock()
-		return false
+// queuedPublicationSize is the byte-size accounting used by the
+// channelMedium's publication queue. Empty publications (e.g. the
+// insufficient-state marker) contribute zero.
+func queuedPublicationSize(qp queuedPublication) int {
+	if qp.Publication.pub == nil {
+		return 0
 	}
-	if q.cnt != 0 {
-		q.mu.Unlock()
-		return true
-	}
-	q.cond.Wait()
-	q.mu.Unlock()
-	return true
-}
-
-// Remove will remove an queuedPublication from the queue.
-// If false is returned, it either means 1) there were no items on the queue
-// or 2) the queue is closed.
-func (q *publicationQueue) Remove() (queuedPublication, bool) {
-	q.mu.Lock()
-	if q.cnt == 0 {
-		q.mu.Unlock()
-		return queuedPublication{}, false
-	}
-	i := q.nodes[q.head]
-	q.head = (q.head + 1) % len(q.nodes)
-	q.cnt--
-	if i.Publication.pub != nil {
-		q.size -= len(i.Publication.pub.Data)
-	}
-
-	if n := len(q.nodes) / 2; n >= q.initCap && q.cnt <= n {
-		q.resize(n)
-	}
-
-	q.mu.Unlock()
-	return i, true
-}
-
-// Len returns the current length of the queue.
-func (q *publicationQueue) Len() int {
-	q.mu.RLock()
-	l := q.cnt
-	q.mu.RUnlock()
-	return l
-}
-
-// Size returns the current size of the queue.
-func (q *publicationQueue) Size() int {
-	q.mu.RLock()
-	s := q.size
-	q.mu.RUnlock()
-	return s
+	return len(qp.Publication.pub.Data)
 }

channel_medium_test.go

@@ -6,9 +6,11 @@ import (
 	"runtime"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/centrifugal/centrifuge/internal/queue"
 	"github.com/stretchr/testify/require"
 )
 
@@ -326,13 +328,15 @@ func TestChannelMediumNoLocalPrevPubForMapSubs(t *testing.T) {
 // channel medium teardown.
 func TestPublicationQueueCloseAndClosed(t *testing.T) {
 	t.Parallel()
-	q := newPublicationQueue(2)
+	q := queue.New(2, queuedPublicationSize)
 	require.False(t, q.Closed())
-	require.True(t, q.Add(queuedPublication{}))
+	_, ok := q.Add(queuedPublication{})
+	require.True(t, ok)
 	q.Close()
 	require.True(t, q.Closed())
 	// Adding after Close must report failure and not panic.
-	require.False(t, q.Add(queuedPublication{}))
+	_, ok = q.Add(queuedPublication{})
+	require.False(t, ok)
 	// Wait must report false on a closed queue without blocking.
 	require.False(t, q.Wait())
 }
@@ -605,3 +609,84 @@ func TestChannelMediumKeepLatestPublicationDelta(t *testing.T) {
 	require.NotNil(t, captured.second, "second broadcast should see prior pub as localPrev")
 	require.Equal(t, []byte("first"), captured.second.Data)
 }
+
+// TestPublicationQueueSizeAccounting verifies that queuedPublicationSize
+// is wired correctly into the generic queue's Size() — including the
+// nil-pub case (insufficient-state markers contribute 0).
+func TestPublicationQueueSizeAccounting(t *testing.T) {
+	t.Parallel()
+	q := queue.New(2, queuedPublicationSize)
+	require.Equal(t, 0, q.Size())
+
+	q.Add(queuedPublication{Publication: queuedPub{pub: &Publication{Data: []byte("hello")}}})
+	q.Add(queuedPublication{Publication: queuedPub{pub: &Publication{Data: []byte("world!")}}})
+	// Insufficient-state marker has nil pub → 0 size contribution.
+	q.Add(queuedPublication{Publication: queuedPub{isInsufficientState: true}})
+	require.Equal(t, len("hello")+len("world!"), q.Size())
+	require.Equal(t, 3, q.Len())
+
+	_, _ = q.Remove()
+	require.Equal(t, len("world!"), q.Size())
+	_, _ = q.Remove()
+	require.Equal(t, 0, q.Size())
+	_, _ = q.Remove()
+	require.Equal(t, 0, q.Size())
+	require.Equal(t, 0, q.Len())
+}
+
+// TestPublicationQueueStress runs the same MPSC stress shape as the
+// internal/queue test battery, but against Queue[queuedPublication] —
+// proving the generic implementation behaves correctly for the
+// channel-medium use case under -race.
+func TestPublicationQueueStress(t *testing.T) {
+	t.Parallel()
+	const (
+		nProducers  = 8
+		perProducer = 1000
+	)
+	q := queue.New(2, queuedPublicationSize)
+
+	var consumed atomic.Int64
+	var consumedBytes atomic.Int64
+	done := make(chan struct{})
+	go func() {
+		for {
+			if !q.Wait() {
+				close(done)
+				return
+			}
+			item, ok := q.Remove()
+			if !ok {
+				continue
+			}
+			consumed.Add(1)
+			consumedBytes.Add(int64(queuedPublicationSize(item)))
+		}
+	}()
+
+	var wg sync.WaitGroup
+	var expectedBytes atomic.Int64
+	for p := 0; p < nProducers; p++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			for i := 0; i < perProducer; i++ {
+				data := []byte("pub-" + strconv.Itoa(id*perProducer+i))
+				q.Add(queuedPublication{Publication: queuedPub{pub: &Publication{Data: data}}})
+				expectedBytes.Add(int64(len(data)))
+			}
+		}(p)
+	}
+	wg.Wait()
+
+	for _, it := range q.CloseRemaining() {
+		consumed.Add(1)
+		consumedBytes.Add(int64(queuedPublicationSize(it)))
+	}
+	<-done
+
+	require.Equal(t, int64(nProducers*perProducer), consumed.Load(),
+		"items lost or duplicated")
+	require.Equal(t, expectedBytes.Load(), consumedBytes.Load(),
+		"byte count mismatch — sizeFn wiring is wrong")
+}