refactor(bigquery/storage/managedwriter): flowcontroller byte reporting

This PR allows the flowcontroller to report bytes in flight for flow
controllers with a bounded byte definition.  The primary connection
load signals for a connection are the inserts/bytes in flight as
reported by the flow controller, and this makes the bytes in flight a
signal we can use.

Important note: an unbounded flow controller will not report any bytes
in flight.  This avoids introducing odd situations due to size
normalization where bytes tracked and the actual capacity of the
semaphore could get out of sync.

Towards: googleapis#7103

bigquery/storage/managedwriter/flow_controller.go

@@ -31,7 +31,8 @@ type flowController struct {
 	// Semaphores for governing pending inserts.
 	semInsertCount, semInsertBytes *semaphore.Weighted
 
-	countRemaining int64 // Atomic.
+	countTracked int64 // Atomic.
+	bytesTracked int64 // Atomic.  Only tracked if bytes are bounded.
 }
 
 func newFlowController(maxInserts, maxInsertBytes int) *flowController {
@@ -80,7 +81,8 @@ func (fc *flowController) acquire(ctx context.Context, sizeBytes int) error {
 			return err
 		}
 	}
-	atomic.AddInt64(&fc.countRemaining, 1)
+	atomic.AddInt64(&fc.bytesTracked, fc.bound(sizeBytes))
+	atomic.AddInt64(&fc.countTracked, 1)
 	return nil
 }
 
@@ -103,12 +105,14 @@ func (fc *flowController) tryAcquire(sizeBytes int) bool {
 			return false
 		}
 	}
-	atomic.AddInt64(&fc.countRemaining, 1)
+	atomic.AddInt64(&fc.bytesTracked, fc.bound(sizeBytes))
+	atomic.AddInt64(&fc.countTracked, 1)
 	return true
 }
 
 func (fc *flowController) release(sizeBytes int) {
-	atomic.AddInt64(&fc.countRemaining, -1)
+	atomic.AddInt64(&fc.countTracked, -1)
+	atomic.AddInt64(&fc.bytesTracked, (0 - fc.bound(sizeBytes)))
 	if fc.semInsertCount != nil {
 		fc.semInsertCount.Release(1)
 	}
@@ -126,5 +130,9 @@ func (fc *flowController) bound(sizeBytes int) int64 {
 }
 
 func (fc *flowController) count() int {
-	return int(atomic.LoadInt64(&fc.countRemaining))
+	return int(atomic.LoadInt64(&fc.countTracked))
+}
+
+func (fc *flowController) bytes() int {
+	return int(atomic.LoadInt64(&fc.bytesTracked))
 }

bigquery/storage/managedwriter/flow_controller_test.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -125,9 +124,6 @@ func TestFlowControllerSaturation(t *testing.T) {
 		},
 	} {
 		fc := newFlowController(maxCount, maxSize)
-		// Atomically track flow controller state.
-		// The flowController itself tracks count.
-		var curSize int64
 		success := errors.New("")
 		// Time out if wantSize or wantCount is never reached.
 		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
@@ -155,15 +151,15 @@ func TestFlowControllerSaturation(t *testing.T) {
 					if c == test.wantCount {
 						hitCount = true
 					}
-					s := atomic.AddInt64(&curSize, int64(test.acquireSize))
+					s := int64(fc.bytes())
 					if s > test.wantSize {
 						return fmt.Errorf("size %d exceeds want %d", s, test.wantSize)
 					}
 					if s == test.wantSize {
 						hitSize = true
 					}
 					time.Sleep(5 * time.Millisecond) // Let other goroutines make progress.
-					if atomic.AddInt64(&curSize, -int64(test.acquireSize)) < 0 {
+					if fc.bytes() < 0 {
 						return errors.New("negative size")
 					}
 					fc.release(test.acquireSize)
@@ -212,11 +208,20 @@ func TestFlowControllerUnboundedCount(t *testing.T) {
 	if !fc.tryAcquire(4) {
 		t.Error("got false, wanted true")
 	}
+	wantBytes := int64(8)
+	if gotB := int64(fc.bytes()); gotB != wantBytes {
+		t.Fatalf("got bytes %d, want %d", gotB, wantBytes)
+	}
 
 	// Fail to tryAcquire 3 bytes.
 	if fc.tryAcquire(3) {
 		t.Error("got true, wanted false")
 	}
+
+	if gotB := int64(fc.bytes()); gotB != wantBytes {
+		t.Fatalf("got bytes %d, want %d", gotB, wantBytes)
+	}
+
 }
 
 func TestFlowControllerUnboundedCount2(t *testing.T) {
@@ -227,14 +232,20 @@ func TestFlowControllerUnboundedCount2(t *testing.T) {
 	if err := fc.acquire(ctx, 4); err != nil {
 		t.Errorf("got %v, wanted no error", err)
 	}
+	wantBytes := int64(0)
+	if gotB := int64(fc.bytes()); gotB != wantBytes {
+		t.Fatalf("got bytes %d, want %d", gotB, wantBytes)
+	}
 	fc.release(1)
 	fc.release(1)
 	fc.release(1)
 	wantCount := int64(-2)
-	c := int64(fc.count())
-	if c != wantCount {
+	if c := int64(fc.count()); c != wantCount {
 		t.Fatalf("got count %d, want %d", c, wantCount)
 	}
+	if gotB := int64(fc.bytes()); gotB != wantBytes {
+		t.Fatalf("got bytes %d, want %d", gotB, wantBytes)
+	}
 }
 
 func TestFlowControllerUnboundedBytes(t *testing.T) {