Merge pull request #1724 from ydb-platform/allow-writer-queue-overflow

Allow writer queue overflow

Author: rekby

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Allowed overflow queue limit for one goroutine at time for topic writer
 * Removed delay before send commit in sync mode of a topic reader
 
 ## v3.105.2

internal/topic/topicwriterinternal/writer_reconnector.go

@@ -13,7 +13,6 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/jonboulle/clockwork"
-	"golang.org/x/sync/semaphore"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/credentials"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/background"
@@ -36,8 +35,8 @@ var (
 	errNonZeroSeqNo                                = xerrors.Wrap(errors.New("ydb: non zero seqno for auto set seqno mode"))                         //nolint:lll
 	errNonZeroCreatedAt                            = xerrors.Wrap(errors.New("ydb: non zero Message.CreatedAt and set auto fill created at option")) //nolint:lll
 	errNoAllowedCodecs                             = xerrors.Wrap(errors.New("ydb: no allowed codecs for write to topic"))
-	errLargeMessage                                = xerrors.Wrap(errors.New("ydb: message uncompressed size more, then limit")) //nolint:lll
-	PublicErrQueueIsFull                           = xerrors.Wrap(errors.New("ydb: queue is full"))
+	errLargeMessage                                = xerrors.Wrap(errors.New("ydb: message uncompressed size more, then limit"))                                                                                                                                                                                             //nolint:lll
+	PublicErrQueueIsFull                           = xerrors.Wrap(errors.New("ydb: queue is full"))                                                                                                                                                                                                                          // Deprecated.
 	PublicErrMessagesPutToInternalQueueBeforeError = xerrors.Wrap(errors.New("ydb: the messages was put to internal buffer before the error happened. It mean about the messages can be delivered to the server"))                                                                                                           //nolint:lll
 	errDiffetentTransactions                       = xerrors.Wrap(errors.New("ydb: internal writer has messages from different trasactions. It is internal logic error, write issue please: https://github.com/ydb-platform/ydb-go-sdk/issues/new?assignees=&labels=bug&projects=&template=01_BUG_REPORT.md&title=bug%3A+")) //nolint:lll
 
@@ -122,7 +121,7 @@ type WriterReconnector struct {
 	background                     background.Worker
 	retrySettings                  topic.RetrySettings
 	writerInstanceID               string
-	semaphore                      *semaphore.Weighted
+	semaphore                      *xsync.SoftWeightedSemaphore
 	firstInitResponseProcessedChan empty.Chan
 	lastSeqNo                      int64
 	encodersMap                    *MultiEncoder
@@ -153,7 +152,7 @@ func newWriterReconnectorStopped(
 	writerInstanceID, _ := rand.Int(rand.Reader, big.NewInt(math.MaxInt64))
 	res := &WriterReconnector{
 		cfg:                            cfg,
-		semaphore:                      semaphore.NewWeighted(int64(cfg.MaxQueueLen)),
+		semaphore:                      xsync.NewSoftWeightedSemaphore(int64(cfg.MaxQueueLen)),
 		queue:                          newMessageQueue(),
 		lastSeqNo:                      -1,
 		firstInitResponseProcessedChan: make(empty.Chan),
@@ -223,21 +222,10 @@ func (w *WriterReconnector) Write(ctx context.Context, messages []PublicMessage)
 	}
 
 	semaphoreWeight := int64(len(messages))
-	if semaphoreWeight > int64(w.cfg.MaxQueueLen) {
-		return xerrors.WithStackTrace(fmt.Errorf(
-			"ydb: add more messages, then max queue limit. max queue: %v, try to add: %v: %w",
-			w.cfg.MaxQueueLen,
-			semaphoreWeight,
-			PublicErrQueueIsFull,
-		))
-	}
 	if err := w.semaphore.Acquire(ctx, semaphoreWeight); err != nil {
 		return xerrors.WithStackTrace(
-			fmt.Errorf("ydb: add new messages exceed max queue size limit. Add count: %v, max size: %v: %w",
-				semaphoreWeight,
-				w.cfg.MaxQueueLen,
-				PublicErrQueueIsFull,
-			))
+			fmt.Errorf("ydb: timeout waiting for queue space to become available: %w", err),
+		)
 	}
 	defer func() {
 		w.semaphore.Release(semaphoreWeight)

internal/topic/topicwriterinternal/writer_reconnector_test.go

@@ -283,48 +283,147 @@ func TestWriterImpl_WriteCodecs(t *testing.T) {
 func TestWriterReconnector_Write_QueueLimit(t *testing.T) {
 	xtest.TestManyTimes(t, func(t testing.TB) {
 		ctx := xtest.Context(t)
+		maxQueueLen := int64(2)
 		w := newWriterReconnectorStopped(NewWriterReconnectorConfig(
 			WithAutoSetSeqNo(false),
-			WithMaxQueueLen(2),
+			WithMaxQueueLen(int(maxQueueLen)),
 		))
 		w.firstConnectionHandled.Store(true)
 
-		waitStartQueueWait := func(targetWaiters int) {
+		waitSemaphoreFull := func() {
 			xtest.SpinWaitCondition(t, nil, func() bool {
-				res := getWaitersCount(w.semaphore) == targetWaiters
-
-				return res
+				// Semaphore is fully acquired when no tokens are available
+				return w.semaphore.TryAcquire(1) == false
 			})
 		}
 
+		// Test normal case within queue limit
 		err := w.Write(ctx, newTestMessages(1, 2))
 		require.NoError(t, err)
+		waitSemaphoreFull()
 
+		// Test queue overflow with context cancellation
 		ctxNoQueueSpace, ctxNoQueueSpaceCancel := xcontext.WithCancel(ctx)
-
 		go func() {
-			waitStartQueueWait(1)
+			waitSemaphoreFull()
 			ctxNoQueueSpaceCancel()
 		}()
 		err = w.Write(ctxNoQueueSpace, newTestMessages(3))
 		require.Error(t, err)
 		require.NotErrorIs(t, err, PublicErrMessagesPutToInternalQueueBeforeError)
 
+		// Test queue space becomes available after ack
 		go func() {
-			waitStartQueueWait(1)
+			waitSemaphoreFull()
 			ackErr := w.queue.AcksReceived([]rawtopicwriter.WriteAck{
 				{
 					SeqNo: 1,
 				},
 			})
 			require.NoError(t, ackErr)
 		}()
-
 		err = w.Write(ctx, newTestMessages(3))
 		require.NoError(t, err)
 	})
 }
 
+// TestWriterReconnector_Write_SoftQueueLimit tests the soft queue limit functionality
+func TestWriterReconnector_Write_SoftQueueLimit(t *testing.T) {
+	xtest.TestManyTimes(t, func(t testing.TB) {
+		ctx := xtest.Context(t)
+		maxQueueLen := int64(2)
+		w := newWriterReconnectorStopped(NewWriterReconnectorConfig(
+			WithAutoSetSeqNo(false),
+			WithMaxQueueLen(int(maxQueueLen)),
+		))
+		w.firstConnectionHandled.Store(true)
+
+		// First write should succeed even if it exceeds queue limit
+		// because semaphore is completely free
+		err := w.Write(ctx, newTestMessages(1, 2, 3)) // 3 messages > maxQueueLen (2)
+		require.NoError(t, err, "first write with overflow should succeed")
+
+		// Second write should fail with timeout because semaphore is already acquired
+		ctxTimeout, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+		defer cancel()
+		err = w.Write(ctxTimeout, newTestMessages(4))
+		require.Error(t, err, "second write should fail when queue is full")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+
+		// Release messages and wait for semaphore to be released
+		ackErr := w.queue.AcksReceived([]rawtopicwriter.WriteAck{
+			{SeqNo: 1},
+			{SeqNo: 2},
+			{SeqNo: 3},
+		})
+		require.NoError(t, ackErr)
+
+		// Now we should be able to write up to maxQueueLen messages
+		err = w.Write(ctx, newTestMessages(4, 5))
+		require.NoError(t, err, "write should succeed after release")
+
+		// But we still can't write more than maxQueueLen when semaphore is not empty
+		ctxTimeout2, cancel2 := context.WithTimeout(ctx, 100*time.Millisecond)
+		defer cancel2()
+		err = w.Write(ctxTimeout2, newTestMessages(6, 7, 8))
+		require.Error(t, err, "write should fail when exceeding queue limit")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+
+		// Cleanup: release remaining messages
+		ackErr = w.queue.AcksReceived([]rawtopicwriter.WriteAck{
+			{SeqNo: 4},
+			{SeqNo: 5},
+		})
+		require.NoError(t, ackErr)
+	})
+}
+
+// TestWriterReconnector_Write_SoftQueueLimitPartialRelease tests partial release behavior
+func TestWriterReconnector_Write_SoftQueueLimitPartialRelease(t *testing.T) {
+	xtest.TestManyTimes(t, func(t testing.TB) {
+		ctx := xtest.Context(t)
+		maxQueueLen := int64(2)
+		w := newWriterReconnectorStopped(NewWriterReconnectorConfig(
+			WithAutoSetSeqNo(false),
+			WithMaxQueueLen(int(maxQueueLen)),
+		))
+		w.firstConnectionHandled.Store(true)
+
+		// Write more than queue limit when semaphore is free
+		err := w.Write(ctx, newTestMessages(1, 2, 3)) // 3 messages > maxQueueLen (2)
+		require.NoError(t, err, "first write with overflow should succeed")
+
+		// Release two messages, one place is enough to write one message
+		ackErr := w.queue.AcksReceived([]rawtopicwriter.WriteAck{
+			{SeqNo: 1},
+			{SeqNo: 2},
+		})
+		require.NoError(t, ackErr)
+
+		// We should be able to write one message now
+		err = w.Write(ctx, newTestMessages(4))
+		require.NoError(t, err, "write should succeed after partial release")
+
+		// But we can't write two messages due to remaining tokens
+		ctxTimeout, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+		defer cancel()
+		err = w.Write(ctxTimeout, newTestMessages(5, 6))
+		require.Error(t, err, "write should fail when exceeding available space")
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+
+		// Release all remaining messages
+		ackErr = w.queue.AcksReceived([]rawtopicwriter.WriteAck{
+			{SeqNo: 3},
+			{SeqNo: 4},
+		})
+		require.NoError(t, ackErr)
+
+		// Now we can write up to maxQueueLen again
+		err = w.Write(ctx, newTestMessages(5, 6))
+		require.NoError(t, err, "write should succeed after full release")
+	})
+}
+
 func TestMessagesPutToInternalQueueBeforeError(t *testing.T) {
 	ctx := xtest.Context(t)
 	w := newWriterReconnectorStopped(NewWriterReconnectorConfig(

internal/topic/topicwriterinternal/writer_reconnector_unsafe_test.go

@@ -1,29 +1,25 @@
 package topicwriterinternal
 
 import (
-	"container/list"
 	"reflect"
 	"runtime"
-	"sync"
 	"unsafe"
 
-	"golang.org/x/sync/semaphore"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 )
 
-func getWaitersCount(sem *semaphore.Weighted) int {
+// getWaitersCount returns number of goroutines waiting for semaphore acquisition
+// by checking overflow state. This is unsafe function which uses reflection
+// and unsafe pointers to access internal fields. It should be used only in tests.
+func getWaitersCount(sem *xsync.SoftWeightedSemaphore) int64 {
+	// Prevent garbage collection of the semaphore while we work with its fields
 	defer runtime.KeepAlive(sem)
 
+	// Get access to overflow field through reflection
 	semVal := reflect.ValueOf(sem).Elem()
-	mutexField := semVal.FieldByName("mu")
+	overflowField := semVal.FieldByName("overflow")
+	overflowAddr := unsafe.Pointer(overflowField.UnsafeAddr())
+	overflow := (*int64)(overflowAddr)
 
-	mutexAddr := unsafe.Pointer(mutexField.UnsafeAddr())
-	mutex := (*sync.Mutex)(mutexAddr)
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	waitersField := semVal.FieldByName("waiters")
-	waitersPointer := unsafe.Pointer(waitersField.UnsafeAddr())
-	waiters := (*list.List)(waitersPointer)
-
-	return waiters.Len()
+	return *overflow
 }

internal/xsync/soft_semaphore.go

@@ -0,0 +1,82 @@
+package xsync
+
+import (
+	"context"
+
+	"golang.org/x/sync/semaphore"
+)
+
+// SoftWeightedSemaphore extends semaphore.Weighted with ability to acquire
+// one request over capacity if semaphore is completely free
+type SoftWeightedSemaphore struct {
+	sem      *semaphore.Weighted
+	capacity int64
+	mu       Mutex
+	overflow int64 // Number of tokens in overflow state
+}
+
+// NewSoftWeightedSemaphore creates new SoftWeightedSemaphore with given capacity
+func NewSoftWeightedSemaphore(n int64) *SoftWeightedSemaphore {
+	return &SoftWeightedSemaphore{
+		sem:      semaphore.NewWeighted(n),
+		capacity: n,
+	}
+}
+
+// Acquire acquires the semaphore with a weight of n.
+// If the semaphore is completely free, the acquisition will succeed regardless of weight.
+func (s *SoftWeightedSemaphore) Acquire(ctx context.Context, n int64) error {
+	// If request doesn't exceed capacity, use normal path
+	if n <= s.capacity {
+		return s.sem.Acquire(ctx, n)
+	}
+
+	// For large requests, try to acquire entire semaphore
+	if err := s.sem.Acquire(ctx, s.capacity); err != nil {
+		return err
+	}
+
+	s.setOverflow(n)
+
+	return nil
+}
+
+// Release releases n tokens back to the semaphore.
+func (s *SoftWeightedSemaphore) Release(n int64) {
+	s.mu.WithLock(func() {
+		if n >= s.overflow {
+			n -= s.overflow
+			s.overflow = 0
+			if n > 0 {
+				s.sem.Release(n)
+			}
+		} else {
+			s.overflow -= n
+		}
+	})
+}
+
+// TryAcquire tries to acquire the semaphore with a weight of n without blocking.
+// If the semaphore is completely free, the acquisition will succeed regardless of weight.
+func (s *SoftWeightedSemaphore) TryAcquire(n int64) bool {
+	// If request doesn't exceed capacity, use normal path
+	if n <= s.capacity {
+		return s.sem.TryAcquire(n)
+	}
+
+	// For large requests, try to acquire entire semaphore
+	if !s.sem.TryAcquire(s.capacity) {
+		return false
+	}
+
+	s.setOverflow(n)
+
+	return true
+}
+
+// setOverflow sets the overflow value when entire semaphore is successfully acquired
+func (s *SoftWeightedSemaphore) setOverflow(n int64) {
+	s.mu.WithLock(func() {
+		s.overflow = n - s.capacity
+	})
+}