Merge pull request #1612 from aaijazi/aaijazi_reject_tx_buffering

Return an error when tx buffer is full

Author: aaijazi

go/vt/vtgate/discoverygateway.go

@@ -143,7 +143,9 @@ func (dg *discoveryGateway) Begin(ctx context.Context, keyspace string, shard st
 	err = dg.withRetry(ctx, keyspace, shard, tabletType, func(conn tabletconn.TabletConn) error {
 		var innerErr error
 		// Potentially buffer this transaction.
-		txbuffer.FakeBuffer(keyspace, shard, attemptNumber)
+		if bufferErr := txbuffer.FakeBuffer(keyspace, shard, attemptNumber); bufferErr != nil {
+			return bufferErr
+		}
 		transactionID, innerErr = conn.Begin(ctx)
 		attemptNumber++
 		return innerErr

go/vt/vtgate/shard_conn.go

@@ -170,7 +170,9 @@ func (sdc *ShardConn) Begin(ctx context.Context) (transactionID int64, err error
 	err = sdc.withRetry(ctx, func(conn tabletconn.TabletConn) error {
 		var innerErr error
 		// Potentially buffer this transaction.
-		txbuffer.FakeBuffer(sdc.keyspace, sdc.shard, attemptNumber)
+		if bufferErr := txbuffer.FakeBuffer(sdc.keyspace, sdc.shard, attemptNumber); bufferErr != nil {
+			return bufferErr
+		}
 		transactionID, innerErr = conn.Begin(ctx)
 		attemptNumber++
 		return innerErr

go/vt/vtgate/txbuffer/tx_buffer.go

@@ -15,11 +15,14 @@ but will not return transient errors during the buffering time.
 package txbuffer
 
 import (
+	"errors"
 	"flag"
 	"sync"
 	"time"
 
 	"github.com/youtube/vitess/go/stats"
+	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
+	"github.com/youtube/vitess/go/vt/vterrors"
 )
 
 var (
@@ -39,26 +42,32 @@ var (
 // timeSleep can be mocked out in unit tests
 var timeSleep = time.Sleep
 
+// errBufferFull is the error returned a buffer request is rejected because the buffer is full.
+var errBufferFull = vterrors.FromError(
+	vtrpcpb.ErrorCode_TRANSIENT_ERROR,
+	errors.New("transaction buffer full, rejecting request"),
+)
+
 // FakeBuffer will pretend to buffer new transactions in VTGate.
 // Transactions *will NOT actually be buffered*, they will just have a delayed start time.
 // This can be useful to understand what the impact of trasaction buffering will be
 // on upstream callers. Once the impact is measured, it can be used to tweak parameter values
 // for the best behavior.
 // FakeBuffer should be called before the VtTablet Begin, otherwise it will increase transaction times.
-func FakeBuffer(keyspace, shard string, attemptNumber int) {
+func FakeBuffer(keyspace, shard string, attemptNumber int) error {
 	// Only buffer on the first Begin attempt, not on possible retries.
 	if !*enableFakeTxBuffer || attemptNumber != 0 {
-		return
+		return nil
 	}
 	if keyspace != *bufferKeyspace || shard != *bufferShard {
-		return
+		return nil
 	}
 	bufferedTransactionsAttempted.Add(1)
 
 	bufferMu.Lock()
 	if int(bufferedTransactions.Get()) >= *maxBufferSize {
 		bufferMu.Unlock()
-		return
+		return errBufferFull
 	}
 	bufferedTransactions.Add(1)
 	bufferMu.Unlock()
@@ -67,4 +76,5 @@ func FakeBuffer(keyspace, shard string, attemptNumber int) {
 	timeSleep(*fakeBufferDelay)
 	// Don't need to lock for this, as there's no race when decrementing the count
 	bufferedTransactions.Add(-1)
+	return nil
 }

go/vt/vtgate/txbuffer/tx_buffer_test.go

@@ -58,6 +58,7 @@ func TestFakeBuffer(t *testing.T) {
 		wantCalled bool
 		// expected value of BufferedTransactionAttempts
 		wantAttempted int
+		wantErr       error
 	}{
 		{
 			desc:             "enableFakeBuffer=False",
@@ -88,6 +89,7 @@ func TestFakeBuffer(t *testing.T) {
 			bufferedTransactions: *maxBufferSize,
 			// When the buffer is full, bufferedTransactionsAttempted should still be incremented
 			wantAttempted: 1,
+			wantErr:       errBufferFull,
 		},
 		{
 			desc:             "buffered successful",
@@ -107,7 +109,11 @@ func TestFakeBuffer(t *testing.T) {
 
 		*enableFakeTxBuffer = test.enableFakeBuffer
 
-		FakeBuffer(test.keyspace, test.shard, test.attemptNumber)
+		gotErr := FakeBuffer(test.keyspace, test.shard, test.attemptNumber)
+
+		if gotErr != test.wantErr {
+			t.Errorf("With %v, FakeBuffer() => %v; want: %v", test.desc, gotErr, test.wantErr)
+		}
 
 		if controller.called != test.wantCalled {
 			t.Errorf("With %v, FakeBuffer() => timeSleep.called: %v; want: %v",
@@ -162,23 +168,27 @@ func TestParallelFakeBuffer(t *testing.T) {
 
 		wg.Add(1)
 		finished := make(chan struct{})
+		var gotErr error
 		go func() {
 			defer wg.Done()
-			FakeBuffer(*bufferKeyspace, *bufferShard, 0)
+			gotErr = FakeBuffer(*bufferKeyspace, *bufferShard, 0)
 			close(finished)
 		}()
 
-		if wantFakeSleepCalled {
-			// the first maxBufferSize calls to FakeBuffer
-			// should call sleep, wait until they do
-			<-controller.blocked
-		} else {
-			// the rest should not block, wait until they're done
-			<-finished
+		// wait until either the gorouotine is blocked (because it's buffering) or until
+		// it's finished (because it shouldn't be buffered).
+		select {
+		case <-controller.blocked:
+		case <-finished:
 		}
 
 		if controller.called {
 			controllers = append(controllers, controller)
+		} else {
+			// if we didn't call fakeSleep, the buffer is full and should return an error saying so.
+			if gotErr != errBufferFull {
+				t.Errorf("On iteration %v, FakeBuffer() => %v; want: %v", i, gotErr, errBufferFull)
+			}
 		}
 
 		if controller.called != wantFakeSleepCalled {