bigquery/storage/managedwriter/managed_stream_test.go - gocloud - Git at Google

 // Copyright 2021 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"
 	"errors"
 	"io"
 	"runtime"
 	"sync"
 	"testing"
 	"time"

 	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
 	"github.com/googleapis/gax-go/v2"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/descriptorpb"
 )

 type testRecvResponse struct {
 	resp *storagepb.AppendRowsResponse
 	err  error
 }

 type testAppendRowsClient struct {
 	storagepb.BigQueryWrite_AppendRowsClient
 	openCount int
 	requests  []*storagepb.AppendRowsRequest
 	responses []*testRecvResponse
 	sendF     func(*storagepb.AppendRowsRequest) error
 	recvF     func() (*storagepb.AppendRowsResponse, error)
 	closeF    func() error
 }

 func (tarc *testAppendRowsClient) Send(req *storagepb.AppendRowsRequest) error {
 	return tarc.sendF(req)
 }

 func (tarc *testAppendRowsClient) Recv() (*storagepb.AppendRowsResponse, error) {
 	return tarc.recvF()
 }

 func (tarc *testAppendRowsClient) CloseSend() error {
 	return tarc.closeF()
 }

 // openTestArc handles wiring in a test AppendRowsClient into a managedstream by providing the open function.
 func openTestArc(testARC *testAppendRowsClient, sendF func(req *storagepb.AppendRowsRequest) error, recvF func() (*storagepb.AppendRowsResponse, error)) func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 	sF := func(req *storagepb.AppendRowsRequest) error {
 		testARC.requests = append(testARC.requests, req)
 		return nil
 	}
 	if sendF != nil {
 		sF = sendF
 	}
 	rF := func() (*storagepb.AppendRowsResponse, error) {
 		return &storagepb.AppendRowsResponse{
 			Response: &storagepb.AppendRowsResponse_AppendResult_{},
 		}, nil
 	}
 	if recvF != nil {
 		rF = recvF
 	}
 	testARC.sendF = sF
 	testARC.recvF = rF
 	testARC.closeF = func() error {
 		return nil
 	}
 	return func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 		testARC.openCount = testARC.openCount + 1
 		// Simulate grpc finalizer goroutine
 		go func() {
 			<-ctx.Done()
 		}()
 		return testARC, nil
 	}
 }

 func TestManagedStream_RequestOptimization(t *testing.T) {

 	ctx := context.Background()
 	testARC := &testAppendRowsClient{}
 	pool := &connectionPool{
 		ctx:                ctx,
 		open:               openTestArc(testARC, nil, nil),
 		baseFlowController: newFlowController(0, 0),
 	}
 	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}
 	ms := &ManagedStream{
 		id:             "foo",
 		ctx:            ctx,
 		streamSettings: defaultStreamSettings(),
 	}
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}
 	ms.streamSettings.streamID = "FOO"
 	ms.streamSettings.TraceID = "TRACE"
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))

 	fakeData := [][]byte{
 		[]byte("foo"),
 		[]byte("bar"),
 	}

 	wantReqs := 3

 	for i := 0; i < wantReqs; i++ {
 		_, err := ms.AppendRows(ctx, fakeData, WithOffset(int64(i)))
 		if err != nil {
 			t.Errorf("AppendRows; %v", err)
 		}
 	}

 	if testARC.openCount != 1 {
 		t.Errorf("expected a single open, got %d", testARC.openCount)
 	}

 	if len(testARC.requests) != wantReqs {
 		t.Errorf("expected %d requests, got %d", wantReqs, len(testARC.requests))
 	}

 	for k, v := range testARC.requests {
 		if v == nil {
 			t.Errorf("request %d was nil", k)
 		}
 		if v.GetOffset() == nil {
 			t.Errorf("request %d had no offset", k)
 		} else {
 			gotOffset := v.GetOffset().GetValue()
 			if gotOffset != int64(k) {
 				t.Errorf("request %d wanted offset %d, got %d", k, k, gotOffset)
 			}
 		}
 		if k == 0 {
 			if v.GetTraceId() == "" {
 				t.Errorf("expected TraceId on first request, was empty")
 			}
 			if v.GetWriteStream() == "" {
 				t.Errorf("expected WriteStream on first request, was empty")
 			}
 			if v.GetProtoRows().GetWriterSchema().GetProtoDescriptor() == nil {
 				t.Errorf("expected WriterSchema on first request, was empty")
 			}

 		} else {
 			// TODO: add validation to ensure we're optimizing requests on the wire.
 			// Sending consecutive requests with same dest/schema we should redact.
 		}
 	}
 }

 func TestManagedStream_FlowControllerFailure(t *testing.T) {

 	ctx := context.Background()

 	pool := &connectionPool{
 		ctx:                ctx,
 		open:               openTestArc(&testAppendRowsClient{}, nil, nil),
 		baseFlowController: newFlowController(1, 0),
 	}
 	router := newSimpleRouter("")
 	if err := pool.activateRouter(router); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	ms := &ManagedStream{
 		id:             "foo",
 		ctx:            ctx,
 		streamSettings: defaultStreamSettings(),
 	}
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWritre: %v", err)
 	}

 	// Exhaust inflight requests on the single connection.
 	router.conn.fc = newFlowController(1, 0)
 	router.conn.fc.acquire(ctx, 0)

 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))

 	fakeData := [][]byte{
 		[]byte("foo"),
 		[]byte("bar"),
 	}

 	// Create a context that will expire during the append.
 	// This is expected to surface a flowcontroller error, as there's no
 	// capacity.
 	expireCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
 	defer cancel()
 	_, err := ms.AppendRows(expireCtx, fakeData)
 	if err == nil {
 		t.Errorf("expected AppendRows to error, but it succeeded")
 	}
 }

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

 	pool := &connectionPool{
 		ctx:                ctx,
 		baseFlowController: newFlowController(0, 0),
 		open: openTestArc(&testAppendRowsClient{},
 			func(req *storagepb.AppendRowsRequest) error {
 				// Append is intentionally slow.
 				time.Sleep(200 * time.Millisecond)
 				return nil
 			}, nil),
 	}
 	router := newSimpleRouter("")
 	if err := pool.activateRouter(router); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	ms := &ManagedStream{
 		id:             "foo",
 		ctx:            ctx,
 		streamSettings: defaultStreamSettings(),
 	}
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}
 	conn := router.conn
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))

 	fakeData := [][]byte{
 		[]byte("foo"),
 	}

 	wantCount := 0
 	if ct := conn.fc.count(); ct != wantCount {
 		t.Errorf("flowcontroller count mismatch, got %d want %d", ct, wantCount)
 	}

 	// Create a context that will expire during the append, to verify the passed in
 	// context expires.
 	expireCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
 	defer cancel()

 	_, err := ms.AppendRows(expireCtx, fakeData)
 	if err == nil {
 		t.Errorf("expected AppendRows to error, but it succeeded")
 	}

 	// We expect the flowcontroller count to still be occupied, as the Send is slow.
 	wantCount = 1
 	if ct := conn.fc.count(); ct != wantCount {
 		t.Errorf("flowcontroller post-append count mismatch, got %d want %d", ct, wantCount)
 	}

 	// Wait for the append to finish, then check again.
 	time.Sleep(300 * time.Millisecond)
 	wantCount = 0
 	if ct := conn.fc.count(); ct != wantCount {
 		t.Errorf("flowcontroller post-append count mismatch, got %d want %d", ct, wantCount)
 	}
 }

 func TestManagedStream_ContextExpiry(t *testing.T) {
 	// Issue: retaining error from append as stream error
 	// https://github.com/googleapis/google-cloud-go/issues/6657
 	ctx := context.Background()

 	pool := &connectionPool{
 		ctx:                ctx,
 		baseFlowController: newFlowController(0, 0),
 		open: openTestArc(&testAppendRowsClient{},
 			func(req *storagepb.AppendRowsRequest) error {
 				return nil
 			}, nil),
 	}
 	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	ms := &ManagedStream{
 		id:             "foo",
 		ctx:            ctx,
 		streamSettings: defaultStreamSettings(),
 	}
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}

 	fakeData := [][]byte{
 		[]byte("foo"),
 	}
 	fakeReq := &storagepb.AppendRowsRequest{
 		Rows: &storagepb.AppendRowsRequest_ProtoRows{
 			ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
 				Rows: &storagepb.ProtoRows{
 					SerializedRows: fakeData,
 				},
 			},
 		},
 	}

 	// Create a context and immediately cancel it.
 	cancelCtx, cancel := context.WithCancel(ctx)
 	cancel()

 	// First, append with an invalid context.
 	pw := newPendingWrite(cancelCtx, ms, fakeReq, ms.curTemplate, "", "")
 	err := ms.appendWithRetry(pw)
 	if err != context.Canceled {
 		t.Errorf("expected cancelled context error, got: %v", err)
 	}

 	// a second append with a valid context should succeed
 	_, err = ms.AppendRows(ctx, fakeData)
 	if err != nil {
 		t.Errorf("expected second append to succeed, but failed: %v", err)
 	}
 }

 func TestManagedStream_AppendDeadlocks(t *testing.T) {
 	// Ensure we don't deadlock by issing two appends.
 	testCases := []struct {
 		desc       string
 		openErrors []error
 		ctx        context.Context
 		respErr    error
 	}{
 		{
 			desc:       "no errors",
 			openErrors: []error{nil, nil},
 			ctx:        context.Background(),
 			respErr:    nil,
 		},
 		{
 			desc:       "cancelled caller context",
 			openErrors: []error{nil, nil},
 			ctx: func() context.Context {
 				cctx, cancel := context.WithCancel(context.Background())
 				cancel()
 				return cctx
 			}(),
 			respErr: context.Canceled,
 		},
 		{
 			desc:       "expired caller context",
 			openErrors: []error{nil, nil},
 			ctx: func() context.Context {
 				cctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
 				defer cancel()
 				time.Sleep(2 * time.Millisecond)
 				return cctx
 			}(),
 			respErr: context.DeadlineExceeded,
 		},
 		{
 			desc:       "errored getstream",
 			openErrors: []error{status.Errorf(codes.ResourceExhausted, "some error"), status.Errorf(codes.ResourceExhausted, "some error")},
 			ctx:        context.Background(),
 			respErr:    status.Errorf(codes.ResourceExhausted, "some error"),
 		},
 	}

 	for _, tc := range testCases {
 		ctx := context.Background()
 		openF := openTestArc(&testAppendRowsClient{}, nil, nil)
 		pool := &connectionPool{
 			ctx: ctx,
 			open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
 				if len(tc.openErrors) == 0 {
 					panic("out of open errors")
 				}
 				curErr := tc.openErrors[0]
 				tc.openErrors = tc.openErrors[1:]
 				if curErr == nil {
 					return openF(ctx, opts...)
 				}
 				return nil, curErr
 			},
 		}
 		router := newSimpleRouter("")
 		if err := pool.activateRouter(router); err != nil {
 			t.Errorf("activateRouter: %v", err)
 		}
 		ms := &ManagedStream{
 			id: "foo",
 			streamSettings: &streamSettings{
 				streamID: "foo",
 			},
 		}
 		ms.ctx, ms.cancel = context.WithCancel(pool.ctx)
 		if err := pool.addWriter(ms); err != nil {
 			t.Errorf("addWriter: %v", err)
 		}

 		testReq := ms.buildRequest([][]byte{[]byte("foo")})
 		// first append
 		pw := newPendingWrite(tc.ctx, ms, testReq, nil, "", "")
 		gotErr := ms.appendWithRetry(pw)
 		if !errors.Is(gotErr, tc.respErr) {
 			t.Errorf("%s first response: got %v, want %v", tc.desc, gotErr, tc.respErr)
 		}
 		// second append
 		pw = newPendingWrite(tc.ctx, ms, testReq, nil, "", "")
 		gotErr = ms.appendWithRetry(pw)
 		if !errors.Is(gotErr, tc.respErr) {
 			t.Errorf("%s second response: got %v, want %v", tc.desc, gotErr, tc.respErr)
 		}

 		// Issue two closes, to ensure we're not deadlocking there either.
 		ms.Close()
 		ms.Close()

 		// Issue two more appends, ensure we're not deadlocked as the writer is closed.
 		gotErr = ms.appendWithRetry(pw)
 		if !errors.Is(gotErr, io.EOF) {
 			t.Errorf("expected io.EOF, got %v", gotErr)
 		}
 		gotErr = ms.appendWithRetry(pw)
 		if !errors.Is(gotErr, io.EOF) {
 			t.Errorf("expected io.EOF, got %v", gotErr)
 		}

 	}

 }

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

 	pool := &connectionPool{
 		ctx: ctx,
 		open: openTestArc(&testAppendRowsClient{},
 			func(req *storagepb.AppendRowsRequest) error {
 				// Append is intentionally slower than context to cause pressure.
 				time.Sleep(40 * time.Millisecond)
 				return nil
 			}, nil),
 		baseFlowController: newFlowController(10, 0),
 	}
 	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}
 	ms := &ManagedStream{
 		id:             "foo",
 		ctx:            ctx,
 		streamSettings: defaultStreamSettings(),
 	}
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}

 	fakeData := [][]byte{
 		[]byte("foo"),
 	}

 	threshold := runtime.NumGoroutine() + 20

 	// Send a bunch of appends that expire quicker than response, and monitor that
 	// goroutine growth stays within bounded threshold.
 	for i := 0; i < 250; i++ {
 		expireCtx, cancel := context.WithTimeout(ctx, 25*time.Millisecond)
 		defer cancel()
 		ms.AppendRows(expireCtx, fakeData)
 		if i%50 == 0 {
 			if current := runtime.NumGoroutine(); current > threshold {
 				t.Errorf("potential goroutine leak, append %d: current %d, threshold %d", i, current, threshold)
 			}
 		}
 	}
 }

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

 	reqCount := 0
 	testArc := &testAppendRowsClient{}
 	pool := &connectionPool{
 		ctx: ctx,
 		open: openTestArc(testArc,
 			func(req *storagepb.AppendRowsRequest) error {
 				reqCount++
 				if reqCount%2 == 1 {
 					return io.EOF
 				}
 				return nil
 			}, nil),
 		baseFlowController: newFlowController(1000, 0),
 	}
 	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}
 	ms := &ManagedStream{
 		id:             "foo",
 		ctx:            ctx,
 		streamSettings: defaultStreamSettings(),
 		retry:          newStatelessRetryer(),
 	}
 	ms.retry.maxAttempts = 4
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}

 	fakeData := [][]byte{
 		[]byte("foo"),
 	}

 	threshold := runtime.NumGoroutine() + 5

 	// Send a bunch of appends that will trigger reconnects and monitor that
 	// goroutine growth stays within bounded threshold.
 	for i := 0; i < 30; i++ {
 		writeCtx := context.Background()
 		r, err := ms.AppendRows(writeCtx, fakeData)
 		if err != nil {
 			t.Fatalf("failed to append row: %v", err)
 		}
 		_, err = r.GetResult(context.Background())
 		if err != nil {
 			t.Fatalf("failed to get result: %v", err)
 		}
 		if r.totalAttempts != 2 {
 			t.Fatalf("should trigger a retry, but found: %d attempts", r.totalAttempts)
 		}
 		if testArc.openCount != i+2 {
 			t.Errorf("should trigger a reconnect, but found openCount %d", testArc.openCount)
 		}
 		if i%10 == 0 {
 			if current := runtime.NumGoroutine(); current > threshold {
 				t.Errorf("potential goroutine leak, append %d: current %d, threshold %d", i, current, threshold)
 			}
 		}
 	}
 }

 func TestManagedWriter_CancellationDuringRetry(t *testing.T) {
 	// Issue: double close of pending write.
 	// https://github.com/googleapis/google-cloud-go/issues/7380
 	ctx, cancel := context.WithCancel(context.Background())
 	pool := &connectionPool{
 		ctx: ctx,
 		open: openTestArc(&testAppendRowsClient{},
 			func(req *storagepb.AppendRowsRequest) error {
 				// Append doesn't error, but is slow.
 				time.Sleep(time.Second)
 				return nil
 			},
 			func() (*storagepb.AppendRowsResponse, error) {
 				// Response is slow and always returns a retriable error.
 				time.Sleep(2 * time.Second)
 				return nil, io.EOF
 			}),
 		baseFlowController: newFlowController(10, 0),
 	}
 	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	ms := &ManagedStream{
 		id:             "foo",
 		ctx:            ctx,
 		streamSettings: defaultStreamSettings(),
 		retry:          newStatelessRetryer(),
 	}
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter: %v", err)
 	}

 	fakeData := [][]byte{
 		[]byte("foo"),
 	}

 	res, err := ms.AppendRows(context.Background(), fakeData)
 	if err != nil {
 		t.Errorf("AppendRows send err: %v", err)
 	}
 	cancel()

 	select {

 	case <-res.Ready():
 		if _, err := res.GetResult(context.Background()); err == nil {
 			t.Errorf("expected failure, got success")
 		}

 	case <-time.After(5 * time.Second):
 		t.Errorf("result was not ready in expected time")
 	}
 }

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

 	pool := &connectionPool{
 		ctx:                ctx,
 		cancel:             cancel,
 		baseFlowController: newFlowController(0, 0),
 		open: openTestArc(&testAppendRowsClient{},
 			func(req *storagepb.AppendRowsRequest) error {
 				return nil
 			}, nil),
 	}
 	router := newSimpleRouter("")
 	if err := pool.activateRouter(router); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	ms := &ManagedStream{
 		id:             "foo",
 		streamSettings: defaultStreamSettings(),
 	}
 	ms.ctx, ms.cancel = context.WithCancel(pool.ctx)
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter A: %v", err)
 	}

 	if router.conn == nil {
 		t.Errorf("expected non-nil connection")
 	}

 	if err := ms.Close(); err != io.EOF {
 		t.Errorf("msB.Close, want %v got %v", io.EOF, err)
 	}
 	if router.conn != nil {
 		t.Errorf("expected nil connection")
 	}
 	if ms.ctx.Err() == nil {
 		t.Errorf("expected writer ctx to be dead, is alive")
 	}
 }

 // This test exists to try to surface data races by sharing
 // a single writer with multiple goroutines.  It doesn't assert
 // anything about the behavior of the system.
 func TestManagedStream_RaceFinder(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())

 	var totalsMu sync.Mutex
 	totalSends := 0
 	totalRecvs := 0
 	pool := &connectionPool{
 		ctx:                ctx,
 		cancel:             cancel,
 		baseFlowController: newFlowController(0, 0),
 		open: openTestArc(&testAppendRowsClient{},
 			func(req *storagepb.AppendRowsRequest) error {
 				totalsMu.Lock()
 				totalSends = totalSends + 1
 				curSends := totalSends
 				totalsMu.Unlock()
 				if curSends%25 == 0 {
 					//time.Sleep(10 * time.Millisecond)
 					return io.EOF
 				}
 				return nil
 			},
 			func() (*storagepb.AppendRowsResponse, error) {
 				totalsMu.Lock()
 				totalRecvs = totalRecvs + 1
 				curRecvs := totalRecvs
 				totalsMu.Unlock()
 				if curRecvs%15 == 0 {
 					return nil, io.EOF
 				}
 				return &storagepb.AppendRowsResponse{}, nil
 			}),
 	}
 	router := newSimpleRouter("")
 	if err := pool.activateRouter(router); err != nil {
 		t.Errorf("activateRouter: %v", err)
 	}

 	ms := &ManagedStream{
 		id:             "foo",
 		streamSettings: defaultStreamSettings(),
 		retry:          newStatelessRetryer(),
 	}
 	ms.retry.maxAttempts = 4
 	ms.ctx, ms.cancel = context.WithCancel(pool.ctx)
 	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
 	if err := pool.addWriter(ms); err != nil {
 		t.Errorf("addWriter A: %v", err)
 	}

 	if router.conn == nil {
 		t.Errorf("expected non-nil connection")
 	}

 	numWriters := 5
 	numWrites := 15

 	var wg sync.WaitGroup
 	wg.Add(numWriters)
 	for i := 0; i < numWriters; i++ {
 		go func() {
 			for j := 0; j < numWrites; j++ {
 				result, err := ms.AppendRows(ctx, [][]byte{[]byte("foo")})
 				if err != nil {
 					continue
 				}
 				_, err = result.GetResult(ctx)
 				if err != nil {
 					continue
 				}
 			}
 			wg.Done()
 		}()
 	}
 	wg.Wait()
 	cancel()
 }
	// Copyright 2021 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"
	"errors"
	"io"
	"runtime"
	"sync"
	"testing"
	"time"

	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
	"github.com/googleapis/gax-go/v2"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/descriptorpb"
	)

	type testRecvResponse struct {
	resp *storagepb.AppendRowsResponse
	err error
	}

	type testAppendRowsClient struct {
	storagepb.BigQueryWrite_AppendRowsClient
	openCount int
	requests []*storagepb.AppendRowsRequest
	responses []*testRecvResponse
	sendF func(*storagepb.AppendRowsRequest) error
	recvF func() (*storagepb.AppendRowsResponse, error)
	closeF func() error
	}

	func (tarc testAppendRowsClient) Send(req storagepb.AppendRowsRequest) error {
	return tarc.sendF(req)
	}

	func (tarc testAppendRowsClient) Recv() (storagepb.AppendRowsResponse, error) {
	return tarc.recvF()
	}

	func (tarc *testAppendRowsClient) CloseSend() error {
	return tarc.closeF()
	}

	// openTestArc handles wiring in a test AppendRowsClient into a managedstream by providing the open function.
	func openTestArc(testARC testAppendRowsClient, sendF func(req storagepb.AppendRowsRequest) error, recvF func() (*storagepb.AppendRowsResponse, error)) func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	sF := func(req *storagepb.AppendRowsRequest) error {
	testARC.requests = append(testARC.requests, req)
	return nil
	}
	if sendF != nil {
	sF = sendF
	}
	rF := func() (*storagepb.AppendRowsResponse, error) {
	return &storagepb.AppendRowsResponse{
	Response: &storagepb.AppendRowsResponse_AppendResult_{},
	}, nil
	}
	if recvF != nil {
	rF = recvF
	}
	testARC.sendF = sF
	testARC.recvF = rF
	testARC.closeF = func() error {
	return nil
	}
	return func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	testARC.openCount = testARC.openCount + 1
	// Simulate grpc finalizer goroutine
	go func() {
	<-ctx.Done()
	}()
	return testARC, nil
	}
	}

	func TestManagedStream_RequestOptimization(t *testing.T) {

	ctx := context.Background()
	testARC := &testAppendRowsClient{}
	pool := &connectionPool{
	ctx: ctx,
	open: openTestArc(testARC, nil, nil),
	baseFlowController: newFlowController(0, 0),
	}
	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
	t.Errorf("activateRouter: %v", err)
	}
	ms := &ManagedStream{
	id: "foo",
	ctx: ctx,
	streamSettings: defaultStreamSettings(),
	}
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}
	ms.streamSettings.streamID = "FOO"
	ms.streamSettings.TraceID = "TRACE"
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))

	fakeData := [][]byte{
	[]byte("foo"),
	[]byte("bar"),
	}

	wantReqs := 3

	for i := 0; i < wantReqs; i++ {
	_, err := ms.AppendRows(ctx, fakeData, WithOffset(int64(i)))
	if err != nil {
	t.Errorf("AppendRows; %v", err)
	}
	}

	if testARC.openCount != 1 {
	t.Errorf("expected a single open, got %d", testARC.openCount)
	}

	if len(testARC.requests) != wantReqs {
	t.Errorf("expected %d requests, got %d", wantReqs, len(testARC.requests))
	}

	for k, v := range testARC.requests {
	if v == nil {
	t.Errorf("request %d was nil", k)
	}
	if v.GetOffset() == nil {
	t.Errorf("request %d had no offset", k)
	} else {
	gotOffset := v.GetOffset().GetValue()
	if gotOffset != int64(k) {
	t.Errorf("request %d wanted offset %d, got %d", k, k, gotOffset)
	}
	}
	if k == 0 {
	if v.GetTraceId() == "" {
	t.Errorf("expected TraceId on first request, was empty")
	}
	if v.GetWriteStream() == "" {
	t.Errorf("expected WriteStream on first request, was empty")
	}
	if v.GetProtoRows().GetWriterSchema().GetProtoDescriptor() == nil {
	t.Errorf("expected WriterSchema on first request, was empty")
	}

	} else {
	// TODO: add validation to ensure we're optimizing requests on the wire.
	// Sending consecutive requests with same dest/schema we should redact.
	}
	}
	}

	func TestManagedStream_FlowControllerFailure(t *testing.T) {

	ctx := context.Background()

	pool := &connectionPool{
	ctx: ctx,
	open: openTestArc(&testAppendRowsClient{}, nil, nil),
	baseFlowController: newFlowController(1, 0),
	}
	router := newSimpleRouter("")
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	ms := &ManagedStream{
	id: "foo",
	ctx: ctx,
	streamSettings: defaultStreamSettings(),
	}
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWritre: %v", err)
	}

	// Exhaust inflight requests on the single connection.
	router.conn.fc = newFlowController(1, 0)
	router.conn.fc.acquire(ctx, 0)

	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))

	fakeData := [][]byte{
	[]byte("foo"),
	[]byte("bar"),
	}

	// Create a context that will expire during the append.
	// This is expected to surface a flowcontroller error, as there's no
	// capacity.
	expireCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
	defer cancel()
	_, err := ms.AppendRows(expireCtx, fakeData)
	if err == nil {
	t.Errorf("expected AppendRows to error, but it succeeded")
	}
	}

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

	pool := &connectionPool{
	ctx: ctx,
	baseFlowController: newFlowController(0, 0),
	open: openTestArc(&testAppendRowsClient{},
	func(req *storagepb.AppendRowsRequest) error {
	// Append is intentionally slow.
	time.Sleep(200 * time.Millisecond)
	return nil
	}, nil),
	}
	router := newSimpleRouter("")
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	ms := &ManagedStream{
	id: "foo",
	ctx: ctx,
	streamSettings: defaultStreamSettings(),
	}
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}
	conn := router.conn
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))

	fakeData := [][]byte{
	[]byte("foo"),
	}

	wantCount := 0
	if ct := conn.fc.count(); ct != wantCount {
	t.Errorf("flowcontroller count mismatch, got %d want %d", ct, wantCount)
	}

	// Create a context that will expire during the append, to verify the passed in
	// context expires.
	expireCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
	defer cancel()

	_, err := ms.AppendRows(expireCtx, fakeData)
	if err == nil {
	t.Errorf("expected AppendRows to error, but it succeeded")
	}

	// We expect the flowcontroller count to still be occupied, as the Send is slow.
	wantCount = 1
	if ct := conn.fc.count(); ct != wantCount {
	t.Errorf("flowcontroller post-append count mismatch, got %d want %d", ct, wantCount)
	}

	// Wait for the append to finish, then check again.
	time.Sleep(300 * time.Millisecond)
	wantCount = 0
	if ct := conn.fc.count(); ct != wantCount {
	t.Errorf("flowcontroller post-append count mismatch, got %d want %d", ct, wantCount)
	}
	}

	func TestManagedStream_ContextExpiry(t *testing.T) {
	// Issue: retaining error from append as stream error
	// https://github.com/googleapis/google-cloud-go/issues/6657
	ctx := context.Background()

	pool := &connectionPool{
	ctx: ctx,
	baseFlowController: newFlowController(0, 0),
	open: openTestArc(&testAppendRowsClient{},
	func(req *storagepb.AppendRowsRequest) error {
	return nil
	}, nil),
	}
	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	ms := &ManagedStream{
	id: "foo",
	ctx: ctx,
	streamSettings: defaultStreamSettings(),
	}
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}

	fakeData := [][]byte{
	[]byte("foo"),
	}
	fakeReq := &storagepb.AppendRowsRequest{
	Rows: &storagepb.AppendRowsRequest_ProtoRows{
	ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
	Rows: &storagepb.ProtoRows{
	SerializedRows: fakeData,
	},
	},
	},
	}

	// Create a context and immediately cancel it.
	cancelCtx, cancel := context.WithCancel(ctx)
	cancel()

	// First, append with an invalid context.
	pw := newPendingWrite(cancelCtx, ms, fakeReq, ms.curTemplate, "", "")
	err := ms.appendWithRetry(pw)
	if err != context.Canceled {
	t.Errorf("expected cancelled context error, got: %v", err)
	}

	// a second append with a valid context should succeed
	_, err = ms.AppendRows(ctx, fakeData)
	if err != nil {
	t.Errorf("expected second append to succeed, but failed: %v", err)
	}
	}

	func TestManagedStream_AppendDeadlocks(t *testing.T) {
	// Ensure we don't deadlock by issing two appends.
	testCases := []struct {
	desc string
	openErrors []error
	ctx context.Context
	respErr error
	}{
	{
	desc: "no errors",
	openErrors: []error{nil, nil},
	ctx: context.Background(),
	respErr: nil,
	},
	{
	desc: "cancelled caller context",
	openErrors: []error{nil, nil},
	ctx: func() context.Context {
	cctx, cancel := context.WithCancel(context.Background())
	cancel()
	return cctx
	}(),
	respErr: context.Canceled,
	},
	{
	desc: "expired caller context",
	openErrors: []error{nil, nil},
	ctx: func() context.Context {
	cctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
	defer cancel()
	time.Sleep(2 * time.Millisecond)
	return cctx
	}(),
	respErr: context.DeadlineExceeded,
	},
	{
	desc: "errored getstream",
	openErrors: []error{status.Errorf(codes.ResourceExhausted, "some error"), status.Errorf(codes.ResourceExhausted, "some error")},
	ctx: context.Background(),
	respErr: status.Errorf(codes.ResourceExhausted, "some error"),
	},
	}

	for _, tc := range testCases {
	ctx := context.Background()
	openF := openTestArc(&testAppendRowsClient{}, nil, nil)
	pool := &connectionPool{
	ctx: ctx,
	open: func(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
	if len(tc.openErrors) == 0 {
	panic("out of open errors")
	}
	curErr := tc.openErrors[0]
	tc.openErrors = tc.openErrors[1:]
	if curErr == nil {
	return openF(ctx, opts...)
	}
	return nil, curErr
	},
	}
	router := newSimpleRouter("")
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}
	ms := &ManagedStream{
	id: "foo",
	streamSettings: &streamSettings{
	streamID: "foo",
	},
	}
	ms.ctx, ms.cancel = context.WithCancel(pool.ctx)
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}

	testReq := ms.buildRequest([][]byte{[]byte("foo")})
	// first append
	pw := newPendingWrite(tc.ctx, ms, testReq, nil, "", "")
	gotErr := ms.appendWithRetry(pw)
	if !errors.Is(gotErr, tc.respErr) {
	t.Errorf("%s first response: got %v, want %v", tc.desc, gotErr, tc.respErr)
	}
	// second append
	pw = newPendingWrite(tc.ctx, ms, testReq, nil, "", "")
	gotErr = ms.appendWithRetry(pw)
	if !errors.Is(gotErr, tc.respErr) {
	t.Errorf("%s second response: got %v, want %v", tc.desc, gotErr, tc.respErr)
	}

	// Issue two closes, to ensure we're not deadlocking there either.
	ms.Close()
	ms.Close()

	// Issue two more appends, ensure we're not deadlocked as the writer is closed.
	gotErr = ms.appendWithRetry(pw)
	if !errors.Is(gotErr, io.EOF) {
	t.Errorf("expected io.EOF, got %v", gotErr)
	}
	gotErr = ms.appendWithRetry(pw)
	if !errors.Is(gotErr, io.EOF) {
	t.Errorf("expected io.EOF, got %v", gotErr)
	}

	}

	}

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

	pool := &connectionPool{
	ctx: ctx,
	open: openTestArc(&testAppendRowsClient{},
	func(req *storagepb.AppendRowsRequest) error {
	// Append is intentionally slower than context to cause pressure.
	time.Sleep(40 * time.Millisecond)
	return nil
	}, nil),
	baseFlowController: newFlowController(10, 0),
	}
	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
	t.Errorf("activateRouter: %v", err)
	}
	ms := &ManagedStream{
	id: "foo",
	ctx: ctx,
	streamSettings: defaultStreamSettings(),
	}
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}

	fakeData := [][]byte{
	[]byte("foo"),
	}

	threshold := runtime.NumGoroutine() + 20

	// Send a bunch of appends that expire quicker than response, and monitor that
	// goroutine growth stays within bounded threshold.
	for i := 0; i < 250; i++ {
	expireCtx, cancel := context.WithTimeout(ctx, 25*time.Millisecond)
	defer cancel()
	ms.AppendRows(expireCtx, fakeData)
	if i%50 == 0 {
	if current := runtime.NumGoroutine(); current > threshold {
	t.Errorf("potential goroutine leak, append %d: current %d, threshold %d", i, current, threshold)
	}
	}
	}
	}

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

	reqCount := 0
	testArc := &testAppendRowsClient{}
	pool := &connectionPool{
	ctx: ctx,
	open: openTestArc(testArc,
	func(req *storagepb.AppendRowsRequest) error {
	reqCount++
	if reqCount%2 == 1 {
	return io.EOF
	}
	return nil
	}, nil),
	baseFlowController: newFlowController(1000, 0),
	}
	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
	t.Errorf("activateRouter: %v", err)
	}
	ms := &ManagedStream{
	id: "foo",
	ctx: ctx,
	streamSettings: defaultStreamSettings(),
	retry: newStatelessRetryer(),
	}
	ms.retry.maxAttempts = 4
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}

	fakeData := [][]byte{
	[]byte("foo"),
	}

	threshold := runtime.NumGoroutine() + 5

	// Send a bunch of appends that will trigger reconnects and monitor that
	// goroutine growth stays within bounded threshold.
	for i := 0; i < 30; i++ {
	writeCtx := context.Background()
	r, err := ms.AppendRows(writeCtx, fakeData)
	if err != nil {
	t.Fatalf("failed to append row: %v", err)
	}
	_, err = r.GetResult(context.Background())
	if err != nil {
	t.Fatalf("failed to get result: %v", err)
	}
	if r.totalAttempts != 2 {
	t.Fatalf("should trigger a retry, but found: %d attempts", r.totalAttempts)
	}
	if testArc.openCount != i+2 {
	t.Errorf("should trigger a reconnect, but found openCount %d", testArc.openCount)
	}
	if i%10 == 0 {
	if current := runtime.NumGoroutine(); current > threshold {
	t.Errorf("potential goroutine leak, append %d: current %d, threshold %d", i, current, threshold)
	}
	}
	}
	}

	func TestManagedWriter_CancellationDuringRetry(t *testing.T) {
	// Issue: double close of pending write.
	// https://github.com/googleapis/google-cloud-go/issues/7380
	ctx, cancel := context.WithCancel(context.Background())
	pool := &connectionPool{
	ctx: ctx,
	open: openTestArc(&testAppendRowsClient{},
	func(req *storagepb.AppendRowsRequest) error {
	// Append doesn't error, but is slow.
	time.Sleep(time.Second)
	return nil
	},
	func() (*storagepb.AppendRowsResponse, error) {
	// Response is slow and always returns a retriable error.
	time.Sleep(2 * time.Second)
	return nil, io.EOF
	}),
	baseFlowController: newFlowController(10, 0),
	}
	if err := pool.activateRouter(newSimpleRouter("")); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	ms := &ManagedStream{
	id: "foo",
	ctx: ctx,
	streamSettings: defaultStreamSettings(),
	retry: newStatelessRetryer(),
	}
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter: %v", err)
	}

	fakeData := [][]byte{
	[]byte("foo"),
	}

	res, err := ms.AppendRows(context.Background(), fakeData)
	if err != nil {
	t.Errorf("AppendRows send err: %v", err)
	}
	cancel()

	select {

	case <-res.Ready():
	if _, err := res.GetResult(context.Background()); err == nil {
	t.Errorf("expected failure, got success")
	}

	case <-time.After(5 * time.Second):
	t.Errorf("result was not ready in expected time")
	}
	}

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

	pool := &connectionPool{
	ctx: ctx,
	cancel: cancel,
	baseFlowController: newFlowController(0, 0),
	open: openTestArc(&testAppendRowsClient{},
	func(req *storagepb.AppendRowsRequest) error {
	return nil
	}, nil),
	}
	router := newSimpleRouter("")
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	ms := &ManagedStream{
	id: "foo",
	streamSettings: defaultStreamSettings(),
	}
	ms.ctx, ms.cancel = context.WithCancel(pool.ctx)
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter A: %v", err)
	}

	if router.conn == nil {
	t.Errorf("expected non-nil connection")
	}

	if err := ms.Close(); err != io.EOF {
	t.Errorf("msB.Close, want %v got %v", io.EOF, err)
	}
	if router.conn != nil {
	t.Errorf("expected nil connection")
	}
	if ms.ctx.Err() == nil {
	t.Errorf("expected writer ctx to be dead, is alive")
	}
	}

	// This test exists to try to surface data races by sharing
	// a single writer with multiple goroutines. It doesn't assert
	// anything about the behavior of the system.
	func TestManagedStream_RaceFinder(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	var totalsMu sync.Mutex
	totalSends := 0
	totalRecvs := 0
	pool := &connectionPool{
	ctx: ctx,
	cancel: cancel,
	baseFlowController: newFlowController(0, 0),
	open: openTestArc(&testAppendRowsClient{},
	func(req *storagepb.AppendRowsRequest) error {
	totalsMu.Lock()
	totalSends = totalSends + 1
	curSends := totalSends
	totalsMu.Unlock()
	if curSends%25 == 0 {
	//time.Sleep(10 * time.Millisecond)
	return io.EOF
	}
	return nil
	},
	func() (*storagepb.AppendRowsResponse, error) {
	totalsMu.Lock()
	totalRecvs = totalRecvs + 1
	curRecvs := totalRecvs
	totalsMu.Unlock()
	if curRecvs%15 == 0 {
	return nil, io.EOF
	}
	return &storagepb.AppendRowsResponse{}, nil
	}),
	}
	router := newSimpleRouter("")
	if err := pool.activateRouter(router); err != nil {
	t.Errorf("activateRouter: %v", err)
	}

	ms := &ManagedStream{
	id: "foo",
	streamSettings: defaultStreamSettings(),
	retry: newStatelessRetryer(),
	}
	ms.retry.maxAttempts = 4
	ms.ctx, ms.cancel = context.WithCancel(pool.ctx)
	ms.curTemplate = newVersionedTemplate().revise(reviseProtoSchema(&descriptorpb.DescriptorProto{}))
	if err := pool.addWriter(ms); err != nil {
	t.Errorf("addWriter A: %v", err)
	}

	if router.conn == nil {
	t.Errorf("expected non-nil connection")
	}

	numWriters := 5
	numWrites := 15

	var wg sync.WaitGroup
	wg.Add(numWriters)
	for i := 0; i < numWriters; i++ {
	go func() {
	for j := 0; j < numWrites; j++ {
	result, err := ms.AppendRows(ctx, [][]byte{[]byte("foo")})
	if err != nil {
	continue
	}
	_, err = result.GetResult(ctx)
	if err != nil {
	continue
	}
	}
	wg.Done()
	}()
	}
	wg.Wait()
	cancel()
	}