github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/dynamo_table_reader.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 2017 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  	"fmt"
    27  	"io"
    28  	"time"
    29  
    30  	"github.com/aws/aws-sdk-go/aws"
    31  	"github.com/aws/aws-sdk-go/service/dynamodb"
    32  
    33  	"github.com/dolthub/dolt/go/store/util/sizecache"
    34  	"github.com/dolthub/dolt/go/store/util/verbose"
    35  )
    36  
    37  const (
    38  	dataAttr    = "data"
    39  	tablePrefix = "*" // I want to use NBS table names as keys when they are written to DynamoDB, but a bare table name is a legal Noms Database name as well. To avoid collisions, dynamoTableReader prepends this prefix (which is not a legal character in a Noms Database name).
    40  )
    41  
    42  // dynamoTableReaderAt assumes the existence of a DynamoDB table whose primary partition key is in String format and named `db`.
    43  type dynamoTableReaderAt struct {
    44  	ddb *ddbTableStore
    45  	h   addr
    46  }
    47  
    48  type tableNotInDynamoErr struct {
    49  	nbs, dynamo string
    50  }
    51  
    52  func (t tableNotInDynamoErr) Error() string {
    53  	return fmt.Sprintf("NBS table %s not present in DynamoDB table %s", t.nbs, t.dynamo)
    54  }
    55  
    56  func (dtra *dynamoTableReaderAt) ReadAtWithStats(ctx context.Context, p []byte, off int64, stats *Stats) (n int, err error) {
    57  	data, err := dtra.ddb.ReadTable(ctx, dtra.h, stats)
    58  
    59  	if err != nil {
    60  		return 0, err
    61  	}
    62  
    63  	n = copy(p, data[off:])
    64  	if n < len(p) {
    65  		err = io.ErrUnexpectedEOF
    66  	}
    67  	return
    68  }
    69  
    70  type ddbTableStore struct {
    71  	ddb    ddbsvc
    72  	table  string
    73  	readRl chan struct{}
    74  	cache  *sizecache.SizeCache // TODO: merge this with tableCache as part of BUG 3601
    75  }
    76  
    77  func (dts *ddbTableStore) ReadTable(ctx context.Context, name addr, stats *Stats) (data []byte, err error) {
    78  	t1 := time.Now()
    79  	if dts.cache != nil {
    80  		if i, present := dts.cache.Get(name); present {
    81  			data = i.([]byte)
    82  			defer func() {
    83  				stats.MemBytesPerRead.Sample(uint64(len(data)))
    84  				stats.MemReadLatency.SampleTimeSince(t1)
    85  			}()
    86  			return data, nil
    87  		}
    88  	}
    89  
    90  	data, err = dts.readTable(ctx, name)
    91  	if data != nil {
    92  		defer func() {
    93  			stats.DynamoBytesPerRead.Sample(uint64(len(data)))
    94  			stats.DynamoReadLatency.SampleTimeSince(t1)
    95  		}()
    96  	}
    97  
    98  	if dts.cache != nil && err == nil {
    99  		dts.cache.Add(name, uint64(len(data)), data)
   100  	}
   101  	return data, err
   102  }
   103  
   104  func (dts *ddbTableStore) readTable(ctx context.Context, name addr) (data []byte, err error) {
   105  	try := func(input *dynamodb.GetItemInput) (data []byte, err error) {
   106  		if dts.readRl != nil {
   107  			dts.readRl <- struct{}{}
   108  			defer func() {
   109  				<-dts.readRl
   110  			}()
   111  		}
   112  		result, rerr := dts.ddb.GetItemWithContext(ctx, input)
   113  		if rerr != nil {
   114  			return nil, rerr
   115  		} else if len(result.Item) == 0 {
   116  			return nil, tableNotInDynamoErr{name.String(), dts.table}
   117  		} else if result.Item[dataAttr] == nil || result.Item[dataAttr].B == nil {
   118  			return nil, fmt.Errorf("NBS table %s in DynamoDB table %s is malformed", name, dts.table)
   119  		}
   120  		return result.Item[dataAttr].B, nil
   121  	}
   122  
   123  	input := dynamodb.GetItemInput{
   124  		TableName: aws.String(dts.table),
   125  		Key: map[string]*dynamodb.AttributeValue{
   126  			dbAttr: {S: aws.String(fmtTableName(name))},
   127  		},
   128  	}
   129  	data, err = try(&input)
   130  	if _, isNotFound := err.(tableNotInDynamoErr); isNotFound {
   131  		verbose.Logger(ctx).Sugar().Debugf("Eventually consistent read for %s failed; trying fully-consistent", name)
   132  		input.ConsistentRead = aws.Bool(true)
   133  		return try(&input)
   134  	}
   135  	return data, err
   136  }
   137  
   138  func fmtTableName(name addr) string {
   139  	return tablePrefix + name.String()
   140  }
   141  
   142  func (dts *ddbTableStore) Write(ctx context.Context, name addr, data []byte) error {
   143  	_, err := dts.ddb.PutItemWithContext(ctx, &dynamodb.PutItemInput{
   144  		TableName: aws.String(dts.table),
   145  		Item: map[string]*dynamodb.AttributeValue{
   146  			dbAttr:   {S: aws.String(fmtTableName(name))},
   147  			dataAttr: {B: data},
   148  		},
   149  	})
   150  
   151  	if dts.cache != nil && err == nil {
   152  		dts.cache.Add(name, uint64(len(data)), data)
   153  	}
   154  	return err
   155  }