rclone/vendor/cloud.google.com/go/bigtable/retry_test.go
2018-03-19 15:51:38 +00:00

372 lines
11 KiB
Go

/*
Copyright 2016 Google Inc. All Rights Reserved.
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
http://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 bigtable
import (
"strings"
"testing"
"time"
"cloud.google.com/go/bigtable/bttest"
"cloud.google.com/go/internal/testutil"
"github.com/golang/protobuf/ptypes/wrappers"
"github.com/google/go-cmp/cmp"
"golang.org/x/net/context"
"google.golang.org/api/option"
btpb "google.golang.org/genproto/googleapis/bigtable/v2"
rpcpb "google.golang.org/genproto/googleapis/rpc/status"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
)
func setupFakeServer(opt ...grpc.ServerOption) (tbl *Table, cleanup func(), err error) {
srv, err := bttest.NewServer("localhost:0", opt...)
if err != nil {
return nil, nil, err
}
conn, err := grpc.Dial(srv.Addr, grpc.WithInsecure(), grpc.WithBlock())
if err != nil {
return nil, nil, err
}
client, err := NewClient(context.Background(), "client", "instance", option.WithGRPCConn(conn))
if err != nil {
return nil, nil, err
}
adminClient, err := NewAdminClient(context.Background(), "client", "instance", option.WithGRPCConn(conn))
if err != nil {
return nil, nil, err
}
if err := adminClient.CreateTable(context.Background(), "table"); err != nil {
return nil, nil, err
}
if err := adminClient.CreateColumnFamily(context.Background(), "table", "cf"); err != nil {
return nil, nil, err
}
t := client.Open("table")
cleanupFunc := func() {
adminClient.Close()
client.Close()
srv.Close()
}
return t, cleanupFunc, nil
}
func TestRetryApply(t *testing.T) {
ctx := context.Background()
errCount := 0
code := codes.Unavailable // Will be retried
// Intercept requests and return an error or defer to the underlying handler
errInjector := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
if strings.HasSuffix(info.FullMethod, "MutateRow") && errCount < 3 {
errCount++
return nil, grpc.Errorf(code, "")
}
return handler(ctx, req)
}
tbl, cleanup, err := setupFakeServer(grpc.UnaryInterceptor(errInjector))
if err != nil {
t.Fatalf("fake server setup: %v", err)
}
defer cleanup()
mut := NewMutation()
mut.Set("cf", "col", 1, []byte("val"))
if err := tbl.Apply(ctx, "row1", mut); err != nil {
t.Errorf("applying single mutation with retries: %v", err)
}
row, err := tbl.ReadRow(ctx, "row1")
if err != nil {
t.Errorf("reading single value with retries: %v", err)
}
if row == nil {
t.Errorf("applying single mutation with retries: could not read back row")
}
code = codes.FailedPrecondition // Won't be retried
errCount = 0
if err := tbl.Apply(ctx, "row", mut); err == nil {
t.Errorf("applying single mutation with no retries: no error")
}
// Check and mutate
mutTrue := NewMutation()
mutTrue.DeleteRow()
mutFalse := NewMutation()
mutFalse.Set("cf", "col", 1, []byte("val"))
condMut := NewCondMutation(ValueFilter("."), mutTrue, mutFalse)
errCount = 0
code = codes.Unavailable // Will be retried
if err := tbl.Apply(ctx, "row1", condMut); err != nil {
t.Errorf("conditionally mutating row with retries: %v", err)
}
row, err = tbl.ReadRow(ctx, "row1") // row1 already in the table
if err != nil {
t.Errorf("reading single value after conditional mutation: %v", err)
}
if row != nil {
t.Errorf("reading single value after conditional mutation: row not deleted")
}
errCount = 0
code = codes.FailedPrecondition // Won't be retried
if err := tbl.Apply(ctx, "row", condMut); err == nil {
t.Errorf("conditionally mutating row with no retries: no error")
}
}
func TestRetryApplyBulk(t *testing.T) {
ctx := context.Background()
// Intercept requests and delegate to an interceptor defined by the test case
errCount := 0
var f func(grpc.ServerStream) error
errInjector := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
if strings.HasSuffix(info.FullMethod, "MutateRows") {
return f(ss)
}
return handler(ctx, ss)
}
tbl, cleanup, err := setupFakeServer(grpc.StreamInterceptor(errInjector))
defer cleanup()
if err != nil {
t.Fatalf("fake server setup: %v", err)
}
errCount = 0
// Test overall request failure and retries
f = func(ss grpc.ServerStream) error {
if errCount < 3 {
errCount++
return grpc.Errorf(codes.Aborted, "")
}
return nil
}
mut := NewMutation()
mut.Set("cf", "col", 1, []byte{})
errors, err := tbl.ApplyBulk(ctx, []string{"row2"}, []*Mutation{mut})
if errors != nil || err != nil {
t.Errorf("bulk with request failure: got: %v, %v, want: nil", errors, err)
}
// Test failures and retries in one request
errCount = 0
m1 := NewMutation()
m1.Set("cf", "col", 1, []byte{})
m2 := NewMutation()
m2.Set("cf", "col2", 1, []byte{})
m3 := NewMutation()
m3.Set("cf", "col3", 1, []byte{})
f = func(ss grpc.ServerStream) error {
var err error
req := new(btpb.MutateRowsRequest)
ss.RecvMsg(req)
switch errCount {
case 0:
// Retryable request failure
err = grpc.Errorf(codes.Unavailable, "")
case 1:
// Two mutations fail
writeMutateRowsResponse(ss, codes.Unavailable, codes.OK, codes.Aborted)
err = nil
case 2:
// Two failures were retried. One will succeed.
if want, got := 2, len(req.Entries); want != got {
t.Errorf("2 bulk retries, got: %d, want %d", got, want)
}
writeMutateRowsResponse(ss, codes.OK, codes.Aborted)
err = nil
case 3:
// One failure was retried and will succeed.
if want, got := 1, len(req.Entries); want != got {
t.Errorf("1 bulk retry, got: %d, want %d", got, want)
}
writeMutateRowsResponse(ss, codes.OK)
err = nil
}
errCount++
return err
}
errors, err = tbl.ApplyBulk(ctx, []string{"row1", "row2", "row3"}, []*Mutation{m1, m2, m3})
if errors != nil || err != nil {
t.Errorf("bulk with retries: got: %v, %v, want: nil", errors, err)
}
// Test unretryable errors
niMut := NewMutation()
niMut.Set("cf", "col", ServerTime, []byte{}) // Non-idempotent
errCount = 0
f = func(ss grpc.ServerStream) error {
var err error
req := new(btpb.MutateRowsRequest)
ss.RecvMsg(req)
switch errCount {
case 0:
// Give non-idempotent mutation a retryable error code.
// Nothing should be retried.
writeMutateRowsResponse(ss, codes.FailedPrecondition, codes.Aborted)
err = nil
case 1:
t.Errorf("unretryable errors: got one retry, want no retries")
}
errCount++
return err
}
errors, err = tbl.ApplyBulk(ctx, []string{"row1", "row2"}, []*Mutation{m1, niMut})
if err != nil {
t.Errorf("unretryable errors: request failed %v", err)
}
want := []error{
grpc.Errorf(codes.FailedPrecondition, ""),
grpc.Errorf(codes.Aborted, ""),
}
if !testutil.Equal(want, errors) {
t.Errorf("unretryable errors: got: %v, want: %v", errors, want)
}
// Test individual errors and a deadline exceeded
f = func(ss grpc.ServerStream) error {
writeMutateRowsResponse(ss, codes.FailedPrecondition, codes.OK, codes.Aborted)
return nil
}
ctx, _ = context.WithTimeout(ctx, 100*time.Millisecond)
errors, err = tbl.ApplyBulk(ctx, []string{"row1", "row2", "row3"}, []*Mutation{m1, m2, m3})
wantErr := context.DeadlineExceeded
if wantErr != err {
t.Errorf("deadline exceeded error: got: %v, want: %v", err, wantErr)
}
if errors != nil {
t.Errorf("deadline exceeded errors: got: %v, want: nil", err)
}
}
func writeMutateRowsResponse(ss grpc.ServerStream, codes ...codes.Code) error {
res := &btpb.MutateRowsResponse{Entries: make([]*btpb.MutateRowsResponse_Entry, len(codes))}
for i, code := range codes {
res.Entries[i] = &btpb.MutateRowsResponse_Entry{
Index: int64(i),
Status: &rpcpb.Status{Code: int32(code), Message: ""},
}
}
return ss.SendMsg(res)
}
func TestRetainRowsAfter(t *testing.T) {
prevRowRange := NewRange("a", "z")
prevRowKey := "m"
want := NewRange("m\x00", "z")
got := prevRowRange.retainRowsAfter(prevRowKey)
if !testutil.Equal(want, got, cmp.AllowUnexported(RowRange{})) {
t.Errorf("range retry: got %v, want %v", got, want)
}
prevRowRangeList := RowRangeList{NewRange("a", "d"), NewRange("e", "g"), NewRange("h", "l")}
prevRowKey = "f"
wantRowRangeList := RowRangeList{NewRange("f\x00", "g"), NewRange("h", "l")}
got = prevRowRangeList.retainRowsAfter(prevRowKey)
if !testutil.Equal(wantRowRangeList, got, cmp.AllowUnexported(RowRange{})) {
t.Errorf("range list retry: got %v, want %v", got, wantRowRangeList)
}
prevRowList := RowList{"a", "b", "c", "d", "e", "f"}
prevRowKey = "b"
wantList := RowList{"c", "d", "e", "f"}
got = prevRowList.retainRowsAfter(prevRowKey)
if !testutil.Equal(wantList, got) {
t.Errorf("list retry: got %v, want %v", got, wantList)
}
}
func TestRetryReadRows(t *testing.T) {
ctx := context.Background()
// Intercept requests and delegate to an interceptor defined by the test case
errCount := 0
var f func(grpc.ServerStream) error
errInjector := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
if strings.HasSuffix(info.FullMethod, "ReadRows") {
return f(ss)
}
return handler(ctx, ss)
}
tbl, cleanup, err := setupFakeServer(grpc.StreamInterceptor(errInjector))
defer cleanup()
if err != nil {
t.Fatalf("fake server setup: %v", err)
}
errCount = 0
// Test overall request failure and retries
f = func(ss grpc.ServerStream) error {
var err error
req := new(btpb.ReadRowsRequest)
ss.RecvMsg(req)
switch errCount {
case 0:
// Retryable request failure
err = grpc.Errorf(codes.Unavailable, "")
case 1:
// Write two rows then error
if want, got := "a", string(req.Rows.RowRanges[0].GetStartKeyClosed()); want != got {
t.Errorf("first retry, no data received yet: got %q, want %q", got, want)
}
writeReadRowsResponse(ss, "a", "b")
err = grpc.Errorf(codes.Unavailable, "")
case 2:
// Retryable request failure
if want, got := "b\x00", string(req.Rows.RowRanges[0].GetStartKeyClosed()); want != got {
t.Errorf("2 range retries: got %q, want %q", got, want)
}
err = grpc.Errorf(codes.Unavailable, "")
case 3:
// Write two more rows
writeReadRowsResponse(ss, "c", "d")
err = nil
}
errCount++
return err
}
var got []string
tbl.ReadRows(ctx, NewRange("a", "z"), func(r Row) bool {
got = append(got, r.Key())
return true
})
want := []string{"a", "b", "c", "d"}
if !testutil.Equal(got, want) {
t.Errorf("retry range integration: got %v, want %v", got, want)
}
}
func writeReadRowsResponse(ss grpc.ServerStream, rowKeys ...string) error {
var chunks []*btpb.ReadRowsResponse_CellChunk
for _, key := range rowKeys {
chunks = append(chunks, &btpb.ReadRowsResponse_CellChunk{
RowKey: []byte(key),
FamilyName: &wrappers.StringValue{Value: "fm"},
Qualifier: &wrappers.BytesValue{Value: []byte("col")},
RowStatus: &btpb.ReadRowsResponse_CellChunk_CommitRow{CommitRow: true},
})
}
return ss.SendMsg(&btpb.ReadRowsResponse{Chunks: chunks})
}