github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/s3_table_reader_test.go (about) 1 // Copyright 2019 Dolthub, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 // This file incorporates work covered by the following copyright and 16 // permission notice: 17 // 18 // Copyright 2016 Attic Labs, Inc. All rights reserved. 19 // Licensed under the Apache License, version 2.0: 20 // http://www.apache.org/licenses/LICENSE-2.0 21 22 package nbs 23 24 import ( 25 "bytes" 26 "context" 27 "io/ioutil" 28 "net" 29 "os" 30 "syscall" 31 "testing" 32 33 "github.com/aws/aws-sdk-go/aws" 34 "github.com/aws/aws-sdk-go/aws/request" 35 "github.com/aws/aws-sdk-go/service/s3" 36 "github.com/stretchr/testify/assert" 37 "github.com/stretchr/testify/require" 38 ) 39 40 func TestS3TableReaderAt(t *testing.T) { 41 s3 := makeFakeS3(t) 42 43 chunks := [][]byte{ 44 []byte("hello2"), 45 []byte("goodbye2"), 46 []byte("badbye2"), 47 } 48 49 tableData, h, err := buildTable(chunks) 50 require.NoError(t, err) 51 s3.data[h.String()] = tableData 52 53 t.Run("TolerateFailingReads", func(t *testing.T) { 54 assert := assert.New(t) 55 56 baseline := s3.getCount 57 tra := &s3TableReaderAt{&s3ObjectReader{makeFlakyS3(s3), "bucket", nil, nil, ""}, h} 58 scratch := make([]byte, len(tableData)) 59 _, err := tra.ReadAtWithStats(context.Background(), scratch, 0, &Stats{}) 60 require.NoError(t, err) 61 // constructing the table reader should have resulted in 2 reads 62 assert.Equal(2, s3.getCount-baseline) 63 assert.Equal(tableData, scratch) 64 }) 65 66 t.Run("WithTableCache", func(t *testing.T) { 67 assert := assert.New(t) 68 dir := makeTempDir(t) 69 defer os.RemoveAll(dir) 70 stats := &Stats{} 71 72 tc, err := newFSTableCache(dir, uint64(2*len(tableData)), 4) 73 require.NoError(t, err) 74 tra := &s3TableReaderAt{&s3ObjectReader{s3, "bucket", nil, tc, ""}, h} 75 76 // First, read when table is not yet cached 77 scratch := make([]byte, len(tableData)) 78 baseline := s3.getCount 79 _, err = tra.ReadAtWithStats(context.Background(), scratch, 0, stats) 80 require.NoError(t, err) 81 assert.True(s3.getCount > baseline) 82 83 // Cache the table and read again 84 err = tc.store(h, bytes.NewReader(tableData), uint64(len(tableData))) 85 require.NoError(t, err) 86 baseline = s3.getCount 87 _, err = tra.ReadAtWithStats(context.Background(), scratch, 0, stats) 88 require.NoError(t, err) 89 assert.Zero(s3.getCount - baseline) 90 }) 91 } 92 93 func TestS3TableReaderAtNamespace(t *testing.T) { 94 assert := assert.New(t) 95 96 s3 := makeFakeS3(t) 97 98 chunks := [][]byte{ 99 []byte("hello2"), 100 []byte("goodbye2"), 101 []byte("badbye2"), 102 } 103 104 ns := "a-prefix-here" 105 106 tableData, h, err := buildTable(chunks) 107 require.NoError(t, err) 108 s3.data["a-prefix-here/"+h.String()] = tableData 109 110 tra := &s3TableReaderAt{&s3ObjectReader{s3, "bucket", nil, nil, ns}, h} 111 scratch := make([]byte, len(tableData)) 112 _, err = tra.ReadAtWithStats(context.Background(), scratch, 0, &Stats{}) 113 require.NoError(t, err) 114 assert.Equal(tableData, scratch) 115 } 116 117 type flakyS3 struct { 118 s3svc 119 alreadyFailed map[string]struct{} 120 } 121 122 func makeFlakyS3(svc s3svc) *flakyS3 { 123 return &flakyS3{svc, map[string]struct{}{}} 124 } 125 126 func (fs3 *flakyS3) GetObjectWithContext(ctx aws.Context, input *s3.GetObjectInput, opts ...request.Option) (*s3.GetObjectOutput, error) { 127 output, err := fs3.s3svc.GetObjectWithContext(ctx, input) 128 129 if err != nil { 130 return nil, err 131 } 132 133 if _, ok := fs3.alreadyFailed[*input.Key]; !ok { 134 fs3.alreadyFailed[*input.Key] = struct{}{} 135 output.Body = ioutil.NopCloser(resettingReader{}) 136 } 137 138 return output, nil 139 } 140 141 type resettingReader struct{} 142 143 func (rr resettingReader) Read(p []byte) (n int, err error) { 144 return 0, &net.OpError{Op: "read", Net: "tcp", Err: &os.SyscallError{Syscall: "read", Err: syscall.ECONNRESET}} 145 }