feat(bigquery/storage/managedwriter): decouple connections and writers (#7314)

Author: shollyman

bigquery/storage/managedwriter/appendresult.go

@@ -16,13 +16,11 @@ package managedwriter
 
 import (
 	"context"
-	"fmt"
 
 	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
 	"github.com/googleapis/gax-go/v2/apierror"
 	grpcstatus "google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
-	"google.golang.org/protobuf/types/descriptorpb"
 )
 
 // NoStreamOffset is a sentinel value for signalling we're not tracking
@@ -31,9 +29,6 @@ const NoStreamOffset int64 = -1
 
 // AppendResult tracks the status of a batch of data rows.
 type AppendResult struct {
-	// rowData contains the serialized row data.
-	rowData [][]byte
-
 	ready chan struct{}
 
 	// if the append failed without a response, this will retain a reference to the error.
@@ -46,10 +41,9 @@ type AppendResult struct {
 	totalAttempts int
 }
 
-func newAppendResult(data [][]byte) *AppendResult {
+func newAppendResult() *AppendResult {
 	return &AppendResult{
-		ready:   make(chan struct{}),
-		rowData: data,
+		ready: make(chan struct{}),
 	}
 }
 
@@ -138,7 +132,7 @@ func (ar *AppendResult) offset(ctx context.Context) int64 {
 func (ar *AppendResult) UpdatedSchema(ctx context.Context) (*storagepb.TableSchema, error) {
 	select {
 	case <-ctx.Done():
-		return nil, fmt.Errorf("context done")
+		return nil, ctx.Err()
 	case <-ar.Ready():
 		if ar.response != nil {
 			if schema := ar.response.GetUpdatedSchema(); schema != nil {
@@ -155,7 +149,7 @@ func (ar *AppendResult) UpdatedSchema(ctx context.Context) (*storagepb.TableSche
 func (ar *AppendResult) TotalAttempts(ctx context.Context) (int, error) {
 	select {
 	case <-ctx.Done():
-		return 0, fmt.Errorf("context done")
+		return 0, ctx.Err()
 	case <-ar.Ready():
 		return ar.totalAttempts, nil
 	}
@@ -168,12 +162,18 @@ type pendingWrite struct {
 	// used is to inform routing decisions.
 	writer *ManagedStream
 
-	request *storagepb.AppendRowsRequest
-	// for schema evolution cases, accept a new schema
-	newSchema *descriptorpb.DescriptorProto
-	result    *AppendResult
+	// We store the request as it's simplex-optimized form, as statistically that's the most
+	// likely outcome when processing requests and it allows us to be efficient on send.
+	// We retain the additional information to build the complete request in the related fields.
+	req           *storagepb.AppendRowsRequest
+	descVersion   *descriptorVersion // schema at time of creation
+	traceID       string
+	writeStreamID string
+
+	// Reference to the AppendResult which is exposed to the user.
+	result *AppendResult
 
-	// this is used by the flow controller.
+	// Flow control is based on the unoptimized request size.
 	reqSize int
 
 	// retains the original request context, primarily for checking against
@@ -188,43 +188,66 @@ type pendingWrite struct {
 // to the pending results for later consumption.  The provided context is
 // embedded in the pending write, as the write may be retried and we want
 // to respect the original context for expiry/cancellation etc.
-func newPendingWrite(ctx context.Context, appends [][]byte) *pendingWrite {
+func newPendingWrite(ctx context.Context, src *ManagedStream, req *storagepb.AppendRowsRequest, curDescVersion *descriptorVersion, writeStreamID, traceID string) *pendingWrite {
 	pw := &pendingWrite{
-		request: &storagepb.AppendRowsRequest{
-			Rows: &storagepb.AppendRowsRequest_ProtoRows{
-				ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
-					Rows: &storagepb.ProtoRows{
-						SerializedRows: appends,
-					},
-				},
-			},
-		},
-		result: newAppendResult(appends),
+		writer: src,
+		result: newAppendResult(),
 		reqCtx: ctx,
+
+		req:           req,
+		descVersion:   curDescVersion,
+		writeStreamID: writeStreamID,
+		traceID:       traceID,
+	}
+	// Compute the approx size for flow control purposes.
+	pw.reqSize = proto.Size(pw.req) + len(writeStreamID) + len(traceID)
+	if pw.descVersion != nil {
+		pw.reqSize += proto.Size(pw.descVersion.descriptorProto)
 	}
-	// We compute the size now for flow controller purposes, though
-	// the actual request size may be slightly larger (e.g. the first
-	// request in a new stream bears schema and stream id).
-	pw.reqSize = proto.Size(pw.request)
 	return pw
 }
 
 // markDone propagates finalization of an append request to the associated
 // AppendResult.
-func (pw *pendingWrite) markDone(resp *storagepb.AppendRowsResponse, err error, fc *flowController) {
+func (pw *pendingWrite) markDone(resp *storagepb.AppendRowsResponse, err error) {
+	// First, propagate necessary state from the pendingWrite to the final result.
 	if resp != nil {
 		pw.result.response = resp
 	}
 	pw.result.err = err
-	// Record the final attempts in the result for the user.
 	pw.result.totalAttempts = pw.attemptCount
 
+	// Close the result's ready channel.
 	close(pw.result.ready)
-	// Clear the reference to the request.
-	pw.request = nil
-	// if there's a flow controller, signal release.  The only time this should be nil is when
-	// encountering issues with flow control during enqueuing the initial request.
-	if fc != nil {
-		fc.release(pw.reqSize)
+	// Cleanup references remaining on the write explicitly.
+	pw.req = nil
+	pw.descVersion = nil
+	pw.writer = nil
+	pw.reqCtx = nil
+}
+
+func (pw *pendingWrite) constructFullRequest(addTrace bool) *storagepb.AppendRowsRequest {
+	req := &storagepb.AppendRowsRequest{}
+	if pw.req != nil {
+		req = proto.Clone(pw.req).(*storagepb.AppendRowsRequest)
+	}
+	if addTrace {
+		req.TraceId = buildTraceID(&streamSettings{TraceID: pw.traceID})
+	}
+	req.WriteStream = pw.writeStreamID
+	if pw.descVersion != nil {
+		ps := &storagepb.ProtoSchema{
+			ProtoDescriptor: pw.descVersion.descriptorProto,
+		}
+		if pr := req.GetProtoRows(); pr != nil {
+			pr.WriterSchema = ps
+		} else {
+			req.Rows = &storagepb.AppendRowsRequest_ProtoRows{
+				ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
+					WriterSchema: ps,
+				},
+			}
+		}
 	}
+	return req
 }

bigquery/storage/managedwriter/appendresult_test.go

@@ -15,7 +15,6 @@
 package managedwriter
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"testing"
@@ -24,42 +23,36 @@ import (
 	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/testing/protocmp"
+	"google.golang.org/protobuf/types/descriptorpb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
-func TestAppendResult(t *testing.T) {
-
-	wantRowBytes := [][]byte{[]byte("rowdata")}
-
-	gotAR := newAppendResult(wantRowBytes)
-	if len(gotAR.rowData) != len(wantRowBytes) {
-		t.Fatalf("length mismatch, got %d want %d elements", len(gotAR.rowData), len(wantRowBytes))
-	}
-	for i := 0; i < len(gotAR.rowData); i++ {
-		if !bytes.Equal(gotAR.rowData[i], wantRowBytes[i]) {
-			t.Errorf("mismatch in row data %d, got %q want %q", i, gotAR.rowData, wantRowBytes)
-		}
-	}
-}
-
 func TestPendingWrite(t *testing.T) {
 	ctx := context.Background()
-	wantRowData := [][]byte{
-		[]byte("row1"),
-		[]byte("row2"),
-		[]byte("row3"),
+	wantReq := &storagepb.AppendRowsRequest{
+		Rows: &storagepb.AppendRowsRequest_ProtoRows{
+			ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
+				Rows: &storagepb.ProtoRows{
+					SerializedRows: [][]byte{
+						[]byte("row1"),
+						[]byte("row2"),
+						[]byte("row3"),
+					},
+				},
+			},
+		},
 	}
 
 	// verify no offset behavior
-	pending := newPendingWrite(ctx, wantRowData)
-	if pending.request.GetOffset() != nil {
-		t.Errorf("request should have no offset, but is present: %q", pending.request.GetOffset().GetValue())
+	pending := newPendingWrite(ctx, nil, wantReq, nil, "", "")
+	if pending.req.GetOffset() != nil {
+		t.Errorf("request should have no offset, but is present: %q", pending.req.GetOffset().GetValue())
 	}
 
-	gotRowCount := len(pending.request.GetProtoRows().GetRows().GetSerializedRows())
-	if gotRowCount != len(wantRowData) {
-		t.Errorf("pendingWrite request mismatch, got %d rows, want %d rows", gotRowCount, len(wantRowData))
+	if diff := cmp.Diff(pending.req, wantReq, protocmp.Transform()); diff != "" {
+		t.Errorf("request mismatch: -got, +want:\n%s", diff)
 	}
 
 	// Verify request is not acknowledged.
@@ -71,37 +64,28 @@ func TestPendingWrite(t *testing.T) {
 	}
 
 	// Mark completed, verify result.
-	pending.markDone(&storage.AppendRowsResponse{}, nil, nil)
+	pending.markDone(&storage.AppendRowsResponse{}, nil)
 	if gotOff := pending.result.offset(ctx); gotOff != NoStreamOffset {
 		t.Errorf("mismatch on completed AppendResult without offset: got %d want %d", gotOff, NoStreamOffset)
 	}
 	if pending.result.err != nil {
 		t.Errorf("mismatch in error on AppendResult, got %v want nil", pending.result.err)
 	}
-	gotData := pending.result.rowData
-	if len(gotData) != len(wantRowData) {
-		t.Errorf("length mismatch on appendresult, got %d, want %d", len(gotData), len(wantRowData))
-	}
-	for i := 0; i < len(gotData); i++ {
-		if !bytes.Equal(gotData[i], wantRowData[i]) {
-			t.Errorf("row %d mismatch in data: got %q want %q", i, gotData[i], wantRowData[i])
-		}
-	}
 
 	// Create new write to verify error result.
-	pending = newPendingWrite(ctx, wantRowData)
+	pending = newPendingWrite(ctx, nil, wantReq, nil, "", "")
 
 	// Manually invoke option to apply offset to request.
 	// This would normally be appied as part of the AppendRows() method on the managed stream.
 	wantOffset := int64(101)
 	f := WithOffset(wantOffset)
 	f(pending)
 
-	if pending.request.GetOffset() == nil {
+	if pending.req.GetOffset() == nil {
 		t.Errorf("expected offset, got none")
 	}
-	if pending.request.GetOffset().GetValue() != wantOffset {
-		t.Errorf("offset mismatch, got %d wanted %d", pending.request.GetOffset().GetValue(), wantOffset)
+	if pending.req.GetOffset().GetValue() != wantOffset {
+		t.Errorf("offset mismatch, got %d wanted %d", pending.req.GetOffset().GetValue(), wantOffset)
 	}
 
 	// Verify completion behavior with an error.
@@ -116,19 +100,10 @@ func TestPendingWrite(t *testing.T) {
 			},
 		},
 	}
-	pending.markDone(testResp, wantErr, nil)
+	pending.markDone(testResp, wantErr)
 
-	if pending.request != nil {
-		t.Errorf("expected request to be cleared, is present: %#v", pending.request)
-	}
-	gotData = pending.result.rowData
-	if len(gotData) != len(wantRowData) {
-		t.Errorf("length mismatch in data: got %d, want %d", len(gotData), len(wantRowData))
-	}
-	for i := 0; i < len(gotData); i++ {
-		if !bytes.Equal(gotData[i], wantRowData[i]) {
-			t.Errorf("row %d mismatch in data: got %q want %q", i, gotData[i], wantRowData[i])
-		}
+	if pending.req != nil {
+		t.Errorf("expected request to be cleared, is present: %#v", pending.req)
 	}
 
 	select {
@@ -153,3 +128,93 @@ func TestPendingWrite(t *testing.T) {
 		}
 	}
 }
+
+func TestPendingWrite_ConstructFullRequest(t *testing.T) {
+
+	testDP := &descriptorpb.DescriptorProto{Name: proto.String("foo")}
+	testDV := newDescriptorVersion(testDP)
+	testEmptyTraceID := buildTraceID(&streamSettings{})
+
+	for _, tc := range []struct {
+		desc     string
+		pw       *pendingWrite
+		addTrace bool
+		want     *storagepb.AppendRowsRequest
+	}{
+		{
+			desc: "nil request",
+			pw: &pendingWrite{
+				descVersion: testDV,
+			},
+			want: &storagepb.AppendRowsRequest{
+				Rows: &storagepb.AppendRowsRequest_ProtoRows{
+					ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
+						WriterSchema: &storagepb.ProtoSchema{
+							ProtoDescriptor: testDP,
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "empty req w/trace",
+			pw: &pendingWrite{
+				req:         &storagepb.AppendRowsRequest{},
+				descVersion: testDV,
+			},
+			addTrace: true,
+			want: &storagepb.AppendRowsRequest{
+				Rows: &storagepb.AppendRowsRequest_ProtoRows{
+					ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
+						WriterSchema: &storagepb.ProtoSchema{
+							ProtoDescriptor: testDP,
+						},
+					},
+				},
+				TraceId: testEmptyTraceID,
+			},
+		},
+		{
+			desc: "basic req",
+			pw: &pendingWrite{
+				req:         &storagepb.AppendRowsRequest{},
+				descVersion: testDV,
+			},
+			want: &storagepb.AppendRowsRequest{
+				Rows: &storagepb.AppendRowsRequest_ProtoRows{
+					ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
+						WriterSchema: &storagepb.ProtoSchema{
+							ProtoDescriptor: testDP,
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "everything w/trace",
+			pw: &pendingWrite{
+				req:           &storagepb.AppendRowsRequest{},
+				descVersion:   testDV,
+				traceID:       "foo",
+				writeStreamID: "streamid",
+			},
+			addTrace: true,
+			want: &storagepb.AppendRowsRequest{
+				WriteStream: "streamid",
+				Rows: &storagepb.AppendRowsRequest_ProtoRows{
+					ProtoRows: &storagepb.AppendRowsRequest_ProtoData{
+						WriterSchema: &storagepb.ProtoSchema{
+							ProtoDescriptor: testDP,
+						},
+					},
+				},
+				TraceId: buildTraceID(&streamSettings{TraceID: "foo"}),
+			},
+		},
+	} {
+		got := tc.pw.constructFullRequest(tc.addTrace)
+		if diff := cmp.Diff(got, tc.want, protocmp.Transform()); diff != "" {
+			t.Errorf("%s diff: %s", tc.desc, diff)
+		}
+	}
+}