Skip to content

Commit a3e97a6

Browse files
shollymanshollyman
authored
feat(bigquery/storage/managedwriter): switch to opt-in retry (#6765)
This PR changes the write option from a disable (opt-out) to an enable (opt-in) for performing retries of appends. We add additional guidance to the option indicating its usage. We also expand the existing large insert integration test to validate the expected behavior both with and without retries enabled.
1 parent 76b06c8 commit a3e97a6

File tree

5 files changed

+104
-13
lines changed

5 files changed

+104
-13
lines changed

‎bigquery/storage/managedwriter/client.go‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,8 +117,6 @@ func (c *Client) buildManagedStream(ctx context.Context, streamFunc streamClient
117117
gax.WithGRPCOptions(grpc.MaxCallRecvMsgSize(10 * 1024 * 1024)),
118118
},
119119
open: createOpenF(ctx, streamFunc),
120-
// We add the new retryer by default, and add a new option to disable it.
121-
retry: newStatelessRetryer(),
122120
}
123121

124122
// apply writer options

‎bigquery/storage/managedwriter/integration_test.go‎

Lines changed: 79 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -159,8 +159,11 @@ func TestIntegration_ManagedWriter(t *testing.T) {
159159
// Don't run this in parallel, we only want to collect stats from this subtest.
160160
testInstrumentation(ctx, t, mwClient, bqClient, dataset)
161161
})
162-
t.Run("TestLargeInsert", func(t *testing.T) {
163-
testLargeInsert(ctx, t, mwClient, bqClient, dataset)
162+
t.Run("TestLargeInsertNoRetry", func(t *testing.T) {
163+
testLargeInsertNoRetry(ctx, t, mwClient, bqClient, dataset)
164+
})
165+
t.Run("TestLargeInsertWithRetry", func(t *testing.T) {
166+
testLargeInsertWithRetry(ctx, t, mwClient, bqClient, dataset)
164167
})
165168
})
166169
}
@@ -596,7 +599,73 @@ func testPendingStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
596599
withExactRowCount(int64(len(testSimpleData))))
597600
}
598601

599-
func testLargeInsert(ctx context.Context, t *testing.T, mwClient *Client, bqClient *bigquery.Client, dataset *bigquery.Dataset) {
602+
func testLargeInsertNoRetry(ctx context.Context, t *testing.T, mwClient *Client, bqClient *bigquery.Client, dataset *bigquery.Dataset) {
603+
testTable := dataset.Table(tableIDs.New())
604+
if err := testTable.Create(ctx, &bigquery.TableMetadata{Schema: testdata.SimpleMessageSchema}); err != nil {
605+
t.Fatalf("failed to create test table %s: %v", testTable.FullyQualifiedName(), err)
606+
}
607+
608+
m := &testdata.SimpleMessageProto2{}
609+
descriptorProto := protodesc.ToDescriptorProto(m.ProtoReflect().Descriptor())
610+
611+
ms, err := mwClient.NewManagedStream(ctx,
612+
WithDestinationTable(TableParentFromParts(testTable.ProjectID, testTable.DatasetID, testTable.TableID)),
613+
WithType(CommittedStream),
614+
WithSchemaDescriptor(descriptorProto),
615+
)
616+
if err != nil {
617+
t.Fatalf("NewManagedStream: %v", err)
618+
}
619+
validateTableConstraints(ctx, t, bqClient, testTable, "before send",
620+
withExactRowCount(0))
621+
622+
// Construct a Very Large request.
623+
var data [][]byte
624+
targetSize := 11 * 1024 * 1024 // 11 MB
625+
b, err := proto.Marshal(testSimpleData[0])
626+
if err != nil {
627+
t.Errorf("failed to marshal message: %v", err)
628+
}
629+
630+
numRows := targetSize / len(b)
631+
data = make([][]byte, numRows)
632+
633+
for i := 0; i < numRows; i++ {
634+
data[i] = b
635+
}
636+
637+
result, err := ms.AppendRows(ctx, data, WithOffset(0))
638+
if err != nil {
639+
t.Errorf("single append failed: %v", err)
640+
}
641+
_, err = result.GetResult(ctx)
642+
if err != nil {
643+
apiErr, ok := apierror.FromError(err)
644+
if !ok {
645+
t.Errorf("GetResult error was not an instance of ApiError")
646+
}
647+
status := apiErr.GRPCStatus()
648+
if status.Code() != codes.InvalidArgument {
649+
t.Errorf("expected InvalidArgument status, got %v", status)
650+
}
651+
}
652+
// our next append should fail (we don't have retries enabled).
653+
if _, err = ms.AppendRows(ctx, [][]byte{b}); err == nil {
654+
t.Fatalf("expected second append to fail, got success: %v", err)
655+
}
656+
657+
// The send failure triggers reconnect, so an additional append will succeed.
658+
result, err = ms.AppendRows(ctx, [][]byte{b})
659+
if err != nil {
660+
t.Fatalf("third append expected to succeed, got error: %v", err)
661+
}
662+
_, err = result.GetResult(ctx)
663+
if err != nil {
664+
t.Errorf("failure result from third append: %v", err)
665+
}
666+
}
667+
668+
func testLargeInsertWithRetry(ctx context.Context, t *testing.T, mwClient *Client, bqClient *bigquery.Client, dataset *bigquery.Dataset) {
600669
testTable := dataset.Table(tableIDs.New())
601670
if err := testTable.Create(ctx, &bigquery.TableMetadata{Schema: testdata.SimpleMessageSchema}); err != nil {
602671
t.Fatalf("failed to create test table %s: %v", testTable.FullyQualifiedName(), err)
@@ -609,6 +678,7 @@ func testLargeInsert(ctx context.Context, t *testing.T, mwClient *Client, bqClie
609678
WithDestinationTable(TableParentFromParts(testTable.ProjectID, testTable.DatasetID, testTable.TableID)),
610679
WithType(CommittedStream),
611680
WithSchemaDescriptor(descriptorProto),
681+
EnableWriteRetries(true),
612682
)
613683
if err != nil {
614684
t.Fatalf("NewManagedStream: %v", err)
@@ -646,15 +716,19 @@ func testLargeInsert(ctx context.Context, t *testing.T, mwClient *Client, bqClie
646716
t.Errorf("expected InvalidArgument status, got %v", status)
647717
}
648718
}
649-
// send a subsequent append as verification we can proceed.
719+
720+
// The second append will succeed, but internally will show a retry.
650721
result, err = ms.AppendRows(ctx, [][]byte{b})
651722
if err != nil {
652-
t.Fatalf("subsequent append failed: %v", err)
723+
t.Fatalf("second append expected to succeed, got error: %v", err)
653724
}
654725
_, err = result.GetResult(ctx)
655726
if err != nil {
656727
t.Errorf("failure result from second append: %v", err)
657728
}
729+
if attempts, _ := result.TotalAttempts(ctx); attempts != 2 {
730+
t.Errorf("expected 2 attempts, got %d", attempts)
731+
}
658732
}
659733

660734
func testInstrumentation(ctx context.Context, t *testing.T, mwClient *Client, bqClient *bigquery.Client, dataset *bigquery.Dataset) {

‎bigquery/storage/managedwriter/managed_stream.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -527,7 +527,7 @@ func recvProcessor(ms *ManagedStream, arc storagepb.BigQueryWrite_AppendRowsClie
527527
func (ms *ManagedStream) processRetry(pw *pendingWrite, appendResp *storagepb.AppendRowsResponse, initialErr error) {
528528
err := initialErr
529529
for {
530-
pause, shouldRetry := ms.retry.Retry(err, pw.attemptCount)
530+
pause, shouldRetry := ms.statelessRetryer().Retry(err, pw.attemptCount)
531531
if !shouldRetry {
532532
// Should not attempt to re-append.
533533
pw.markDone(appendResp, err, ms.fc)

‎bigquery/storage/managedwriter/options.go‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -98,11 +98,15 @@ func WithAppendRowsCallOption(o gax.CallOption) WriterOption {
9898
}
9999
}
100100

101-
// DisableWriteRetries disables the logic for automatically re-enqueuing failed writes.
102-
func DisableWriteRetries(disable bool) WriterOption {
101+
// EnableWriteRetries enables ManagedStream to automatically retry failed appends.
102+
//
103+
// Enabling retries is best suited for cases where users want to achieve at-least-once
104+
// append semantics. Use of automatic retries may complicate patterns where the user
105+
// is designing for exactly-once append semantics.
106+
func EnableWriteRetries(enable bool) WriterOption {
103107
return func(ms *ManagedStream) {
104-
if disable {
105-
ms.retry = nil
108+
if enable {
109+
ms.retry = newStatelessRetryer()
106110
}
107111
}
108112
}

‎bigquery/storage/managedwriter/options_test.go‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121

2222
"cloud.google.com/go/bigquery/internal"
2323
"github.com/google/go-cmp/cmp"
24+
"github.com/google/go-cmp/cmp/cmpopts"
2425
"github.com/googleapis/gax-go/v2"
2526
"google.golang.org/grpc"
2627
)
@@ -122,12 +123,24 @@ func TestWriterOptions(t *testing.T) {
122123
return ms
123124
}(),
124125
},
126+
{
127+
desc: "EnableRetries",
128+
options: []WriterOption{EnableWriteRetries(true)},
129+
want: func() *ManagedStream {
130+
ms := &ManagedStream{
131+
streamSettings: defaultStreamSettings(),
132+
}
133+
ms.retry = newStatelessRetryer()
134+
return ms
135+
}(),
136+
},
125137
{
126138
desc: "multiple",
127139
options: []WriterOption{
128140
WithType(PendingStream),
129141
WithMaxInflightBytes(5),
130142
WithTraceID("traceid"),
143+
EnableWriteRetries(true),
131144
},
132145
want: func() *ManagedStream {
133146
ms := &ManagedStream{
@@ -136,6 +149,7 @@ func TestWriterOptions(t *testing.T) {
136149
ms.streamSettings.MaxInflightBytes = 5
137150
ms.streamSettings.streamType = PendingStream
138151
ms.streamSettings.TraceID = fmt.Sprintf("go-managedwriter:%s traceid", internal.Version)
152+
ms.retry = newStatelessRetryer()
139153
return ms
140154
}(),
141155
},
@@ -151,7 +165,8 @@ func TestWriterOptions(t *testing.T) {
151165

152166
if diff := cmp.Diff(got, tc.want,
153167
cmp.AllowUnexported(ManagedStream{}, streamSettings{}),
154-
cmp.AllowUnexported(sync.Mutex{})); diff != "" {
168+
cmp.AllowUnexported(sync.Mutex{}),
169+
cmpopts.IgnoreUnexported(statelessRetryer{})); diff != "" {
155170
t.Errorf("diff in case (%s):\n%v", tc.desc, diff)
156171
}
157172
}

0 commit comments

Comments
 (0)