bigquery/storage/managedwriter/routers_test.go - gocloud.git - Git at Google

 // Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     https://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package managedwriter

 import (
 	"context"
 	"fmt"
 	"math/rand"
 	"testing"
 	"time"

 	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
 	"github.com/googleapis/gax-go/v2"
 )

 func TestSimpleRouter(t *testing.T) {
 	ctx := context.Background()

 	pool := &connectionPool{
 		ctx: ctx,
 		open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 			return &testAppendRowsClient{}, nil
 		},
 	}

 	router := newSimpleRouter("")
 	if err := pool.activateRouter(router); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	ms := &ManagedStream{
 		ctx:   ctx,
 		retry: newStatelessRetryer(),
 	}

 	pw := newPendingWrite(ctx, ms, &storagepb.AppendRowsRequest{}, nil, "", "")

 	// picking before attaching should yield error
 	if _, err := pool.router.pickConnection(pw); err == nil {
 		t.Errorf("pickConnection: expected error, got success")
 	}
 	writer := &ManagedStream{
 		id: "writer",
 	}
 	if err := pool.addWriter(writer); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}
 	if _, err := pool.router.pickConnection(pw); err != nil {
 		t.Errorf("pickConnection error: %v", err)
 	}
 	if err := pool.removeWriter(writer); err != nil {
 		t.Errorf("disconnectWriter: %v", err)
 	}
 	if _, err := pool.router.pickConnection(pw); err == nil {
 		t.Errorf("pickConnection: expected error, got success")
 	}
 }

 func TestSharedRouter_Basic(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())

 	pool := &connectionPool{
 		ctx:    ctx,
 		cancel: cancel,
 		open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 			return &testAppendRowsClient{}, nil
 		},
 	}

 	router := newSharedRouter(false, 0)
 	if err := pool.activateRouter(router); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}
 	if gotConns := len(router.exclusiveConns); gotConns != 0 {
 		t.Errorf("expected zero connections are start, got %d", gotConns)
 	}

 	ms := &ManagedStream{
 		ctx:   ctx,
 		retry: newStatelessRetryer(),
 	}
 	pw := newPendingWrite(ctx, ms, &storagepb.AppendRowsRequest{}, nil, "", "")
 	// picking before attaching should yield error
 	if _, err := pool.router.pickConnection(pw); err == nil {
 		t.Errorf("pickConnection: expected error, got success")
 	}
 	// attaching a writer without an ID should error.
 	if err := pool.addWriter(ms); err == nil {
 		t.Errorf("expected id-less addWriter to fail")
 	}
 	ms.id = "writer"
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}

 	if _, err := pool.router.pickConnection(pw); err != nil {
 		t.Errorf("pickConnection error: %v", err)
 	}
 	if err := pool.removeWriter(ms); err != nil {
 		t.Errorf("disconnectWriter: %v", err)
 	}
 	if _, err := pool.router.pickConnection(pw); err == nil {
 		t.Errorf("pickConnection: expected error, got success")
 	}
 }

 func TestSharedRouter_Multiplex(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())

 	pool := &connectionPool{
 		id:     newUUID(poolIDPrefix),
 		ctx:    ctx,
 		cancel: cancel,
 		open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 			return &testAppendRowsClient{}, nil
 		},
 		baseFlowController: newFlowController(2, 10),
 	}
 	defer pool.Close()

 	router := newSharedRouter(true, 3)
 	if err := pool.activateRouter(router); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	wantConnCount := 0
 	if got := len(router.multiConns); wantConnCount != got {
 		t.Errorf("wanted %d conns, got %d", wantConnCount, got)
 	}

 	writerA := &ManagedStream{
 		id:             newUUID(writerIDPrefix),
 		streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},
 		ctx:            ctx,
 		cancel:         cancel,
 	}
 	if err := pool.router.writerAttach(writerA); err != nil {
 		t.Fatalf("writerA attach: %v", err)
 	}

 	// after a writer attached, we expect one conn.
 	wantConnCount = 1
 	if got := len(router.multiConns); wantConnCount != got {
 		t.Errorf("wanted %d conns, got %d", wantConnCount, got)
 	}

 	writerB := &ManagedStream{
 		id:             newUUID(writerIDPrefix),
 		streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},
 		ctx:            ctx,
 		cancel:         cancel,
 	}
 	if err := pool.router.writerAttach(writerB); err != nil {
 		t.Fatalf("writerA attach: %v", err)
 	}
 	writerC := &ManagedStream{
 		id:             newUUID(writerIDPrefix),
 		streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},
 		ctx:            ctx,
 		cancel:         cancel,
 	}
 	if err := pool.router.writerAttach(writerC); err != nil {
 		t.Fatalf("writerA attach: %v", err)
 	}

 	wantConnCount = 1
 	if got := len(router.multiConns); wantConnCount != got {
 		t.Fatalf("wanted %d conns, got %d", wantConnCount, got)
 	}

 	pw := newPendingWrite(ctx, writerA, &storagepb.AppendRowsRequest{}, nil, "", "")
 	conn, err := router.pickConnection(pw)
 	if err != nil {
 		t.Fatalf("pickConnection writerA: %v", err)
 	}
 	// generate fake load on the conn associated with writer A
 	conn.fc.acquire(ctx, 1)
 	conn.fc.acquire(ctx, 1)

 	if !conn.isLoaded() {
 		t.Errorf("expected conn to be loaded, was not")
 	}
 	// wait for a watchdog interval
 	time.Sleep(watchDogInterval * 2)

 	wantConnCount = 2
 	// grab read lock so we can assert internal state of the router
 	router.multiMu.RLock()
 	defer router.multiMu.RUnlock()
 	if got := len(router.multiConns); wantConnCount != got {
 		t.Fatalf("wanted %d conns, got %d", wantConnCount, got)
 	}
 	gotLoad0 := router.multiConns[0].curLoad()
 	gotLoad1 := router.multiConns[1].curLoad()
 	if gotLoad0 > gotLoad1 {
 		t.Errorf("expected connections to be ordered by load, got %f, %f", gotLoad0, gotLoad1)
 	}
 	// verify that rebalance occurred
 	connsWithWriters := 0
 	for _, v := range router.invertedMultiMap {
 		if len(v) > 0 {
 			connsWithWriters++
 		}
 	}
 	if connsWithWriters < wantConnCount {
 		t.Errorf("wanted at least %d connections to have writers attached, got %d", wantConnCount, connsWithWriters)
 	}

 }

 func BenchmarkRoutingParallel(b *testing.B) {

 	for _, bm := range []struct {
 		desc              string
 		router            poolRouter
 		numWriters        int
 		numDefaultWriters int
 	}{
 		{
 			desc:              "SimpleRouter",
 			router:            newSimpleRouter(""),
 			numWriters:        1,
 			numDefaultWriters: 1,
 		},
 		{
 			desc:              "SimpleRouter",
 			router:            newSimpleRouter(""),
 			numWriters:        10,
 			numDefaultWriters: 10,
 		},
 		{
 			desc:              "SharedRouter_NoMultiplex",
 			router:            newSharedRouter(false, 0),
 			numWriters:        1,
 			numDefaultWriters: 1,
 		},
 		{
 			desc:              "SharedRouter_NoMultiplex",
 			router:            newSharedRouter(false, 0),
 			numWriters:        10,
 			numDefaultWriters: 10,
 		},
 		{
 			desc:              "SharedRouter_Multiplex1conn",
 			router:            newSharedRouter(true, 1),
 			numWriters:        1,
 			numDefaultWriters: 1,
 		},
 		{
 			desc:              "SharedRouterMultiplex1conn",
 			router:            newSharedRouter(true, 1),
 			numWriters:        10,
 			numDefaultWriters: 10,
 		},
 		{
 			desc:              "SharedRouterMultiplex1conn",
 			router:            newSharedRouter(true, 1),
 			numWriters:        50,
 			numDefaultWriters: 50,
 		},
 		{
 			desc:              "SharedRouterMultiplex10conn",
 			router:            newSharedRouter(true, 10),
 			numWriters:        50,
 			numDefaultWriters: 50,
 		},
 	} {

 		ctx, cancel := context.WithCancel(context.Background())
 		pool := &connectionPool{
 			ctx:    ctx,
 			cancel: cancel,
 			open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 				return &testAppendRowsClient{}, nil
 			},
 		}
 		if err := pool.activateRouter(bm.router); err != nil {
 			b.Errorf("%q: activateRouter: %v", bm.desc, err)
 		}

 		// setup both explicit and default stream writers.
 		var explicitWriters []*ManagedStream
 		var defaultWriters []*ManagedStream

 		for i := 0; i < bm.numWriters; i++ {
 			wCtx, wCancel := context.WithCancel(ctx)
 			writer := &ManagedStream{
 				id:             newUUID(writerIDPrefix),
 				streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/abc123"},
 				ctx:            wCtx,
 				cancel:         wCancel,
 				retry:          newStatelessRetryer(),
 			}
 			explicitWriters = append(explicitWriters, writer)
 		}
 		for i := 0; i < bm.numDefaultWriters; i++ {
 			wCtx, wCancel := context.WithCancel(ctx)
 			writer := &ManagedStream{
 				id:             newUUID(writerIDPrefix),
 				streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},

 				ctx:    wCtx,
 				cancel: wCancel,
 				retry:  newStatelessRetryer(),
 			}
 			defaultWriters = append(defaultWriters, writer)
 		}

 		// attach all writers to router.
 		for k, writer := range explicitWriters {
 			if err := pool.addWriter(writer); err != nil {
 				b.Errorf("addWriter %d: %v", k, err)
 			}
 		}
 		for k, writer := range defaultWriters {
 			if err := pool.addWriter(writer); err != nil {
 				b.Errorf("addWriter %d: %v", k, err)
 			}
 		}

 		baseBenchName := fmt.Sprintf("%s_%dexwriters_%dmpwriters", bm.desc, bm.numWriters, bm.numDefaultWriters)

 		// Benchmark routing for explicit writers.
 		if bm.numWriters > 0 {
 			benchName := fmt.Sprintf("%s_explicitwriters", baseBenchName)

 			b.Run(benchName, func(b *testing.B) {
 				r := rand.New(rand.NewSource(1))
 				b.ResetTimer()
 				for i := 0; i < b.N; i++ {
 					// pick a random explicit writer each time.
 					writer := explicitWriters[r.Intn(bm.numWriters)]
 					pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
 					if _, err := bm.router.pickConnection(pw); err != nil {
 						b.Errorf("pickConnection: %v", err)
 					}
 				}
 			})
 		}

 		// Benchmark concurrent routing for explicit writers.
 		if bm.numWriters > 0 {
 			benchName := fmt.Sprintf("%s_explicitwriters_concurrent", baseBenchName)
 			b.Run(benchName, func(b *testing.B) {
 				b.RunParallel(func(pb *testing.PB) {
 					r := rand.New(rand.NewSource(1))
 					for pb.Next() {
 						writer := explicitWriters[r.Intn(bm.numWriters)]
 						pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
 						if _, err := bm.router.pickConnection(pw); err != nil {
 							b.Errorf("pickConnection: %v", err)
 						}
 					}
 				})
 			})
 		}

 		// Benchmark routing for default writers.
 		if bm.numDefaultWriters > 0 {
 			benchName := fmt.Sprintf("%s_defaultwriters", baseBenchName)

 			b.Run(benchName, func(b *testing.B) {
 				r := rand.New(rand.NewSource(1))
 				b.ResetTimer()
 				for i := 0; i < b.N; i++ {
 					// pick a random default writer each time.
 					writer := defaultWriters[r.Intn(bm.numDefaultWriters)]
 					pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
 					if _, err := bm.router.pickConnection(pw); err != nil {
 						b.Errorf("pickConnection: %v", err)
 					}
 				}
 			})
 		}

 		// Benchmark concurrent routing for default writers.
 		if bm.numDefaultWriters > 0 {
 			benchName := fmt.Sprintf("%s_defaultwriters_concurrent", baseBenchName)

 			b.Run(benchName, func(b *testing.B) {
 				b.RunParallel(func(pb *testing.PB) {
 					r := rand.New(rand.NewSource(1))
 					for pb.Next() {
 						writer := defaultWriters[r.Intn(bm.numDefaultWriters)]
 						pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
 						if _, err := bm.router.pickConnection(pw); err != nil {
 							b.Errorf("pickConnection: %v", err)
 						}
 					}
 				})
 			})
 		}

 		for _, writer := range explicitWriters {
 			writer.Close()
 		}
 		for _, writer := range defaultWriters {
 			writer.Close()
 		}

 		pool.Close()

 	}

 }

 func BenchmarkWatchdogPulse(b *testing.B) {
 	maxFlowInserts := 100
 	maxFlowBytes := 1024
 	for _, numWriters := range []int{1, 2, 5, 10, 50, 100, 250} {
 		for _, numConnections := range []int{1, 2, 4} {

 			ctx, cancel := context.WithCancel(context.Background())
 			// we build the router manually so we can control the watchdog for this benchmark.
 			router := newSharedRouter(false, numConnections)

 			pool := &connectionPool{
 				ctx:    ctx,
 				cancel: cancel,
 				open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 					return &testAppendRowsClient{}, nil
 				},
 				baseFlowController: newFlowController(maxFlowInserts, maxFlowBytes),
 			}
 			if err := pool.activateRouter(router); err != nil {
 				b.Fatalf("(@%d-@%d): activateRouter: %v", numWriters, numConnections, err)
 			}
 			// now, set router as multiplex.  We do this to avoid router activation starting the watchdog
 			// in a seperate goroutine.
 			router.multiplex = true

 			var writers []*ManagedStream

 			for i := 0; i < numWriters; i++ {
 				wCtx, wCancel := context.WithCancel(ctx)
 				writer := &ManagedStream{
 					id:             newUUID(writerIDPrefix),
 					streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},

 					ctx:    wCtx,
 					cancel: wCancel,
 					retry:  newStatelessRetryer(),
 				}
 				writers = append(writers, writer)
 				if err := pool.addWriter(writer); err != nil {
 					b.Fatalf("addWriter %d (@%d-@%d): %v", i, numWriters, numConnections, err)
 				}
 			}

 			// Generate fake load for all connections.
 			r := rand.New(rand.NewSource(time.Now().UnixNano()))
 			countLoad := make([]int, numConnections)
 			byteLoad := make([]int, numConnections)
 			for i := 0; i < numConnections; i++ {
 				countLoad[i] = r.Intn(maxFlowInserts)
 				byteLoad[i] = r.Intn(maxFlowBytes)
 			}

 			benchName := fmt.Sprintf("%dwriters_%dconns", numWriters, numConnections)
 			b.Run(benchName, func(b *testing.B) {
 				if b.N > 9999 {
 					b.Skip("benchmark unstable, only run with -benchtime=NNNNx")
 				}
 				for i := 0; i < b.N; i++ {
 					b.StopTimer()
 					// Each iteration, we reset the loads to the predetermined values, and repoint
 					// all writers to the first connection.
 					for c := 0; c < len(router.multiConns); c++ {
 						router.multiConns[c].fc.countTracked = int64(countLoad[c])
 						router.multiConns[c].fc.bytesTracked = int64(byteLoad[c])
 					}
 					for k := range router.multiMap {
 						router.multiMap[k] = router.multiConns[0]
 					}
 					router.invertedMultiMap = make(map[string][]*ManagedStream)
 					writerSlice := make([]*ManagedStream, len(writers))
 					copy(writerSlice, writers)
 					router.invertedMultiMap[router.multiConns[0].id] = writerSlice
 					b.StartTimer()
 					router.watchdogPulse()
 				}
 			})

 			for _, writer := range writers {
 				writer.Close()
 			}

 			pool.Close()

 		}
 	}

 }
	// Copyright 2023 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package managedwriter

	import (
	"context"
	"fmt"
	"math/rand"
	"testing"
	"time"

	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
	"github.com/googleapis/gax-go/v2"
	)

	func TestSimpleRouter(t *testing.T) {
	ctx := context.Background()

	pool := &connectionPool{
	ctx: ctx,
	open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	return &testAppendRowsClient{}, nil
	},
	}

	router := newSimpleRouter("")
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	ms := &ManagedStream{
	ctx: ctx,
	retry: newStatelessRetryer(),
	}

	pw := newPendingWrite(ctx, ms, &storagepb.AppendRowsRequest{}, nil, "", "")

	// picking before attaching should yield error
	if _, err := pool.router.pickConnection(pw); err == nil {
	t.Errorf("pickConnection: expected error, got success")
	}
	writer := &ManagedStream{
	id: "writer",
	}
	if err := pool.addWriter(writer); err != nil {
	t.Errorf("addWriter: %v", err)
	}
	if _, err := pool.router.pickConnection(pw); err != nil {
	t.Errorf("pickConnection error: %v", err)
	}
	if err := pool.removeWriter(writer); err != nil {
	t.Errorf("disconnectWriter: %v", err)
	}
	if _, err := pool.router.pickConnection(pw); err == nil {
	t.Errorf("pickConnection: expected error, got success")
	}
	}

	func TestSharedRouter_Basic(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	pool := &connectionPool{
	ctx: ctx,
	cancel: cancel,
	open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	return &testAppendRowsClient{}, nil
	},
	}

	router := newSharedRouter(false, 0)
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}
	if gotConns := len(router.exclusiveConns); gotConns != 0 {
	t.Errorf("expected zero connections are start, got %d", gotConns)
	}

	ms := &ManagedStream{
	ctx: ctx,
	retry: newStatelessRetryer(),
	}
	pw := newPendingWrite(ctx, ms, &storagepb.AppendRowsRequest{}, nil, "", "")
	// picking before attaching should yield error
	if _, err := pool.router.pickConnection(pw); err == nil {
	t.Errorf("pickConnection: expected error, got success")
	}
	// attaching a writer without an ID should error.
	if err := pool.addWriter(ms); err == nil {
	t.Errorf("expected id-less addWriter to fail")
	}
	ms.id = "writer"
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}

	if _, err := pool.router.pickConnection(pw); err != nil {
	t.Errorf("pickConnection error: %v", err)
	}
	if err := pool.removeWriter(ms); err != nil {
	t.Errorf("disconnectWriter: %v", err)
	}
	if _, err := pool.router.pickConnection(pw); err == nil {
	t.Errorf("pickConnection: expected error, got success")
	}
	}

	func TestSharedRouter_Multiplex(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	pool := &connectionPool{
	id: newUUID(poolIDPrefix),
	ctx: ctx,
	cancel: cancel,
	open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	return &testAppendRowsClient{}, nil
	},
	baseFlowController: newFlowController(2, 10),
	}
	defer pool.Close()

	router := newSharedRouter(true, 3)
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	wantConnCount := 0
	if got := len(router.multiConns); wantConnCount != got {
	t.Errorf("wanted %d conns, got %d", wantConnCount, got)
	}

	writerA := &ManagedStream{
	id: newUUID(writerIDPrefix),
	streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},
	ctx: ctx,
	cancel: cancel,
	}
	if err := pool.router.writerAttach(writerA); err != nil {
	t.Fatalf("writerA attach: %v", err)
	}

	// after a writer attached, we expect one conn.
	wantConnCount = 1
	if got := len(router.multiConns); wantConnCount != got {
	t.Errorf("wanted %d conns, got %d", wantConnCount, got)
	}

	writerB := &ManagedStream{
	id: newUUID(writerIDPrefix),
	streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},
	ctx: ctx,
	cancel: cancel,
	}
	if err := pool.router.writerAttach(writerB); err != nil {
	t.Fatalf("writerA attach: %v", err)
	}
	writerC := &ManagedStream{
	id: newUUID(writerIDPrefix),
	streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},
	ctx: ctx,
	cancel: cancel,
	}
	if err := pool.router.writerAttach(writerC); err != nil {
	t.Fatalf("writerA attach: %v", err)
	}

	wantConnCount = 1
	if got := len(router.multiConns); wantConnCount != got {
	t.Fatalf("wanted %d conns, got %d", wantConnCount, got)
	}

	pw := newPendingWrite(ctx, writerA, &storagepb.AppendRowsRequest{}, nil, "", "")
	conn, err := router.pickConnection(pw)
	if err != nil {
	t.Fatalf("pickConnection writerA: %v", err)
	}
	// generate fake load on the conn associated with writer A
	conn.fc.acquire(ctx, 1)
	conn.fc.acquire(ctx, 1)

	if !conn.isLoaded() {
	t.Errorf("expected conn to be loaded, was not")
	}
	// wait for a watchdog interval
	time.Sleep(watchDogInterval * 2)

	wantConnCount = 2
	// grab read lock so we can assert internal state of the router
	router.multiMu.RLock()
	defer router.multiMu.RUnlock()
	if got := len(router.multiConns); wantConnCount != got {
	t.Fatalf("wanted %d conns, got %d", wantConnCount, got)
	}
	gotLoad0 := router.multiConns[0].curLoad()
	gotLoad1 := router.multiConns[1].curLoad()
	if gotLoad0 > gotLoad1 {
	t.Errorf("expected connections to be ordered by load, got %f, %f", gotLoad0, gotLoad1)
	}
	// verify that rebalance occurred
	connsWithWriters := 0
	for _, v := range router.invertedMultiMap {
	if len(v) > 0 {
	connsWithWriters++
	}
	}
	if connsWithWriters < wantConnCount {
	t.Errorf("wanted at least %d connections to have writers attached, got %d", wantConnCount, connsWithWriters)
	}

	}

	func BenchmarkRoutingParallel(b *testing.B) {

	for _, bm := range []struct {
	desc string
	router poolRouter
	numWriters int
	numDefaultWriters int
	}{
	{
	desc: "SimpleRouter",
	router: newSimpleRouter(""),
	numWriters: 1,
	numDefaultWriters: 1,
	},
	{
	desc: "SimpleRouter",
	router: newSimpleRouter(""),
	numWriters: 10,
	numDefaultWriters: 10,
	},
	{
	desc: "SharedRouter_NoMultiplex",
	router: newSharedRouter(false, 0),
	numWriters: 1,
	numDefaultWriters: 1,
	},
	{
	desc: "SharedRouter_NoMultiplex",
	router: newSharedRouter(false, 0),
	numWriters: 10,
	numDefaultWriters: 10,
	},
	{
	desc: "SharedRouter_Multiplex1conn",
	router: newSharedRouter(true, 1),
	numWriters: 1,
	numDefaultWriters: 1,
	},
	{
	desc: "SharedRouterMultiplex1conn",
	router: newSharedRouter(true, 1),
	numWriters: 10,
	numDefaultWriters: 10,
	},
	{
	desc: "SharedRouterMultiplex1conn",
	router: newSharedRouter(true, 1),
	numWriters: 50,
	numDefaultWriters: 50,
	},
	{
	desc: "SharedRouterMultiplex10conn",
	router: newSharedRouter(true, 10),
	numWriters: 50,
	numDefaultWriters: 50,
	},
	} {

	ctx, cancel := context.WithCancel(context.Background())
	pool := &connectionPool{
	ctx: ctx,
	cancel: cancel,
	open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	return &testAppendRowsClient{}, nil
	},
	}
	if err := pool.activateRouter(bm.router); err != nil {
	b.Errorf("%q: activateRouter: %v", bm.desc, err)
	}

	// setup both explicit and default stream writers.
	var explicitWriters []*ManagedStream
	var defaultWriters []*ManagedStream

	for i := 0; i < bm.numWriters; i++ {
	wCtx, wCancel := context.WithCancel(ctx)
	writer := &ManagedStream{
	id: newUUID(writerIDPrefix),
	streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/abc123"},
	ctx: wCtx,
	cancel: wCancel,
	retry: newStatelessRetryer(),
	}
	explicitWriters = append(explicitWriters, writer)
	}
	for i := 0; i < bm.numDefaultWriters; i++ {
	wCtx, wCancel := context.WithCancel(ctx)
	writer := &ManagedStream{
	id: newUUID(writerIDPrefix),
	streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},

	ctx: wCtx,
	cancel: wCancel,
	retry: newStatelessRetryer(),
	}
	defaultWriters = append(defaultWriters, writer)
	}

	// attach all writers to router.
	for k, writer := range explicitWriters {
	if err := pool.addWriter(writer); err != nil {
	b.Errorf("addWriter %d: %v", k, err)
	}
	}
	for k, writer := range defaultWriters {
	if err := pool.addWriter(writer); err != nil {
	b.Errorf("addWriter %d: %v", k, err)
	}
	}

	baseBenchName := fmt.Sprintf("%s_%dexwriters_%dmpwriters", bm.desc, bm.numWriters, bm.numDefaultWriters)

	// Benchmark routing for explicit writers.
	if bm.numWriters > 0 {
	benchName := fmt.Sprintf("%s_explicitwriters", baseBenchName)

	b.Run(benchName, func(b *testing.B) {
	r := rand.New(rand.NewSource(1))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
	// pick a random explicit writer each time.
	writer := explicitWriters[r.Intn(bm.numWriters)]
	pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
	if _, err := bm.router.pickConnection(pw); err != nil {
	b.Errorf("pickConnection: %v", err)
	}
	}
	})
	}

	// Benchmark concurrent routing for explicit writers.
	if bm.numWriters > 0 {
	benchName := fmt.Sprintf("%s_explicitwriters_concurrent", baseBenchName)
	b.Run(benchName, func(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
	r := rand.New(rand.NewSource(1))
	for pb.Next() {
	writer := explicitWriters[r.Intn(bm.numWriters)]
	pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
	if _, err := bm.router.pickConnection(pw); err != nil {
	b.Errorf("pickConnection: %v", err)
	}
	}
	})
	})
	}

	// Benchmark routing for default writers.
	if bm.numDefaultWriters > 0 {
	benchName := fmt.Sprintf("%s_defaultwriters", baseBenchName)

	b.Run(benchName, func(b *testing.B) {
	r := rand.New(rand.NewSource(1))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
	// pick a random default writer each time.
	writer := defaultWriters[r.Intn(bm.numDefaultWriters)]
	pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
	if _, err := bm.router.pickConnection(pw); err != nil {
	b.Errorf("pickConnection: %v", err)
	}
	}
	})
	}

	// Benchmark concurrent routing for default writers.
	if bm.numDefaultWriters > 0 {
	benchName := fmt.Sprintf("%s_defaultwriters_concurrent", baseBenchName)

	b.Run(benchName, func(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
	r := rand.New(rand.NewSource(1))
	for pb.Next() {
	writer := defaultWriters[r.Intn(bm.numDefaultWriters)]
	pw := newPendingWrite(context.Background(), writer, &storagepb.AppendRowsRequest{}, nil, "", "")
	if _, err := bm.router.pickConnection(pw); err != nil {
	b.Errorf("pickConnection: %v", err)
	}
	}
	})
	})
	}

	for _, writer := range explicitWriters {
	writer.Close()
	}
	for _, writer := range defaultWriters {
	writer.Close()
	}

	pool.Close()

	}

	}

	func BenchmarkWatchdogPulse(b *testing.B) {
	maxFlowInserts := 100
	maxFlowBytes := 1024
	for _, numWriters := range []int{1, 2, 5, 10, 50, 100, 250} {
	for _, numConnections := range []int{1, 2, 4} {

	ctx, cancel := context.WithCancel(context.Background())
	// we build the router manually so we can control the watchdog for this benchmark.
	router := newSharedRouter(false, numConnections)

	pool := &connectionPool{
	ctx: ctx,
	cancel: cancel,
	open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	return &testAppendRowsClient{}, nil
	},
	baseFlowController: newFlowController(maxFlowInserts, maxFlowBytes),
	}
	if err := pool.activateRouter(router); err != nil {
	b.Fatalf("(@%d-@%d): activateRouter: %v", numWriters, numConnections, err)
	}
	// now, set router as multiplex. We do this to avoid router activation starting the watchdog
	// in a seperate goroutine.
	router.multiplex = true

	var writers []*ManagedStream

	for i := 0; i < numWriters; i++ {
	wCtx, wCancel := context.WithCancel(ctx)
	writer := &ManagedStream{
	id: newUUID(writerIDPrefix),
	streamSettings: &streamSettings{streamID: "projects/foo/datasets/bar/tables/baz/streams/_default"},

	ctx: wCtx,
	cancel: wCancel,
	retry: newStatelessRetryer(),
	}
	writers = append(writers, writer)
	if err := pool.addWriter(writer); err != nil {
	b.Fatalf("addWriter %d (@%d-@%d): %v", i, numWriters, numConnections, err)
	}
	}

	// Generate fake load for all connections.
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	countLoad := make([]int, numConnections)
	byteLoad := make([]int, numConnections)
	for i := 0; i < numConnections; i++ {
	countLoad[i] = r.Intn(maxFlowInserts)
	byteLoad[i] = r.Intn(maxFlowBytes)
	}

	benchName := fmt.Sprintf("%dwriters_%dconns", numWriters, numConnections)
	b.Run(benchName, func(b *testing.B) {
	if b.N > 9999 {
	b.Skip("benchmark unstable, only run with -benchtime=NNNNx")
	}
	for i := 0; i < b.N; i++ {
	b.StopTimer()
	// Each iteration, we reset the loads to the predetermined values, and repoint
	// all writers to the first connection.
	for c := 0; c < len(router.multiConns); c++ {
	router.multiConns[c].fc.countTracked = int64(countLoad[c])
	router.multiConns[c].fc.bytesTracked = int64(byteLoad[c])
	}
	for k := range router.multiMap {
	router.multiMap[k] = router.multiConns[0]
	}
	router.invertedMultiMap = make(map[string][]*ManagedStream)
	writerSlice := make([]*ManagedStream, len(writers))
	copy(writerSlice, writers)
	router.invertedMultiMap[router.multiConns[0].id] = writerSlice
	b.StartTimer()
	router.watchdogPulse()
	}
	})

	for _, writer := range writers {
	writer.Close()
	}

	pool.Close()

	}
	}

	}