github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/aws_chunk_source.go (about)

     1  // Copyright 2019-2021 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  	"context"
    26  	"errors"
    27  	"sync"
    28  	"time"
    29  )
    30  
    31  type indexParserF func([]byte) (tableIndex, error)
    32  
    33  func newAWSChunkSource(ctx context.Context, ddb *ddbTableStore, s3 *s3ObjectReader, al awsLimits, name addr, chunkCount uint32, indexCache *indexCache, stats *Stats, parseIndex indexParserF) (cs chunkSource, err error) {
    34  	if indexCache != nil {
    35  		indexCache.lockEntry(name)
    36  		defer func() {
    37  			unlockErr := indexCache.unlockEntry(name)
    38  
    39  			if err == nil {
    40  				err = unlockErr
    41  			}
    42  		}()
    43  
    44  		if index, found := indexCache.get(name); found {
    45  			tra := &awsTableReaderAt{al: al, ddb: ddb, s3: s3, name: name, chunkCount: chunkCount}
    46  			return &chunkSourceAdapter{newTableReader(index, tra, s3BlockSize), name}, nil
    47  		}
    48  	}
    49  
    50  	t1 := time.Now()
    51  	indexBytes, tra, err := func() ([]byte, tableReaderAt, error) {
    52  		if al.tableMayBeInDynamo(chunkCount) {
    53  			data, err := ddb.ReadTable(ctx, name, stats)
    54  
    55  			if data == nil && err == nil { // There MUST be either data or an error
    56  				return nil, &dynamoTableReaderAt{}, errors.New("no data available")
    57  			}
    58  
    59  			if data != nil {
    60  				return data, &dynamoTableReaderAt{ddb: ddb, h: name}, nil
    61  			}
    62  
    63  			if _, ok := err.(tableNotInDynamoErr); !ok {
    64  				return nil, &dynamoTableReaderAt{}, err
    65  			}
    66  		}
    67  
    68  		size := indexSize(chunkCount) + footerSize
    69  		buff := make([]byte, size)
    70  
    71  		n, _, err := s3.ReadFromEnd(ctx, name, buff, stats)
    72  
    73  		if err != nil {
    74  			return nil, &dynamoTableReaderAt{}, err
    75  		}
    76  
    77  		if size != uint64(n) {
    78  			return nil, &dynamoTableReaderAt{}, errors.New("failed to read all data")
    79  		}
    80  
    81  		return buff, &s3TableReaderAt{s3: s3, h: name}, nil
    82  	}()
    83  
    84  	if err != nil {
    85  		return &chunkSourceAdapter{}, err
    86  	}
    87  
    88  	stats.IndexBytesPerRead.Sample(uint64(len(indexBytes)))
    89  	stats.IndexReadLatency.SampleTimeSince(t1)
    90  
    91  	index, err := parseIndex(indexBytes)
    92  
    93  	if err != nil {
    94  		return emptyChunkSource{}, err
    95  	}
    96  
    97  	if ohi, ok := index.(onHeapTableIndex); indexCache != nil && ok {
    98  		indexCache.put(name, ohi)
    99  	}
   100  
   101  	return &chunkSourceAdapter{newTableReader(index, tra, s3BlockSize), name}, nil
   102  }
   103  
   104  type awsTableReaderAt struct {
   105  	once     sync.Once
   106  	getTRErr error
   107  	tra      tableReaderAt
   108  
   109  	al  awsLimits
   110  	ddb *ddbTableStore
   111  	s3  *s3ObjectReader
   112  
   113  	name       addr
   114  	chunkCount uint32
   115  }
   116  
   117  func (atra *awsTableReaderAt) ReadAtWithStats(ctx context.Context, p []byte, off int64, stats *Stats) (int, error) {
   118  	atra.once.Do(func() {
   119  		atra.tra, atra.getTRErr = atra.getTableReaderAt(ctx, stats)
   120  	})
   121  
   122  	if atra.getTRErr != nil {
   123  		return 0, atra.getTRErr
   124  	}
   125  
   126  	return atra.tra.ReadAtWithStats(ctx, p, off, stats)
   127  }
   128  
   129  func (atra *awsTableReaderAt) getTableReaderAt(ctx context.Context, stats *Stats) (tableReaderAt, error) {
   130  	if atra.al.tableMayBeInDynamo(atra.chunkCount) {
   131  		data, err := atra.ddb.ReadTable(ctx, atra.name, stats)
   132  
   133  		if data == nil && err == nil { // There MUST be either data or an error
   134  			return &dynamoTableReaderAt{}, errors.New("no data available")
   135  		}
   136  
   137  		if data != nil {
   138  			return &dynamoTableReaderAt{ddb: atra.ddb, h: atra.name}, nil
   139  		}
   140  
   141  		if _, ok := err.(tableNotInDynamoErr); !ok {
   142  			return &dynamoTableReaderAt{}, err
   143  		}
   144  	}
   145  
   146  	return &s3TableReaderAt{s3: atra.s3, h: atra.name}, nil
   147  }