func (a *Assoc) Get(ctx context.Context, kind assoc.Kind, k digest.Digest) (digest.Digest, digest.Digest, error) {
	reqStartTime := time.Now()
	defer func() {
		metrics.GetDydbassocOpLatencySecondsHistogram(ctx, "get").Observe(time.Now().Sub(reqStartTime).Seconds())
	}()
	var v digest.Digest
	switch kind {
	case assoc.Fileset, assoc.ExecInspect, assoc.Logs, assoc.Bundle:
	default:
		return k, v, errors.E(errors.NotSupported, errors.Errorf("mappings of kind %v are not supported", kind))
	}
	var col string
	switch kind {
	case assoc.Fileset:
		col = "Value"
	case assoc.ExecInspect:
		col = "ExecInspect"
	case assoc.Logs:
		col = "Logs"
	case assoc.Bundle:
		col = "Bundle"
	}
	if err := a.Limiter.Acquire(ctx, 1); err != nil {
		return k, v, err
	}
	defer a.Limiter.Release(1)
	ctx, cancel := context.WithTimeout(ctx, getTimeout)
	defer cancel()
	var item *dynamodb.AttributeValue
	if k.IsAbbrev() {
		resp, err := a.DB.Query(&dynamodb.QueryInput{
			TableName:              aws.String(a.TableName),
			IndexName:              aws.String("ID4-ID-index"),
			KeyConditionExpression: aws.String("ID4 = :id4"),
			ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
				":id4": {S: aws.String(k.HexN(4))},
			},
		})
		if err != nil {
			return k, v, err
		}
		expanded := make(map[digest.Digest]*dynamodb.AttributeValue)
		for _, it := range resp.Items {
			v := it["ID"]
			if v == nil || v.S == nil {
				continue
			}
			kit, err := reflow.Digester.Parse(*v.S)
			if err != nil {
				log.Debugf("invalid dynamodb entry %v", it)
				continue
			}
			if kit.Expands(k) {
				expanded[kit] = it[col]
			}
		}
		switch len(expanded) {
		case 0:
		case 1:
			for k1, v1 := range expanded {
				k = k1
				item = v1
			}
		default:
			return k, v, errors.E("lookup", k, errors.Invalid, errors.New("more than one key matched"))
		}
	} else {
		resp, err := a.DB.GetItemWithContext(ctx, &dynamodb.GetItemInput{
			Key: map[string]*dynamodb.AttributeValue{
				"ID": {
					S: aws.String(k.String()),
				},
			},
			TableName: aws.String(a.TableName),
		})
		if err != nil {
			return k, v, err
		}
		item = resp.Item[col]
	}
	if item == nil || (kind == assoc.Fileset && item.S == nil) || (kind == assoc.ExecInspect && item.L == nil) || (kind == assoc.Logs && item.L == nil) || (kind == assoc.Bundle && item.L == nil) {
		return k, v, errors.E("lookup", k, errors.NotExist)
	}
	if item.L != nil {
		item = item.L[0]
	}
	var err error
	v, err = reflow.Digester.Parse(*item.S)
	if err != nil {
		return k, v, errors.E("lookup", k, err)
	}
	_, err = a.DB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		Key: map[string]*dynamodb.AttributeValue{
			"ID": {
				S: aws.String(k.String()),
			},
		},
		TableName:        aws.String(a.TableName),
		UpdateExpression: aws.String("SET LastAccessTime = :time ADD AccessCount :one"),
		ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
			":time": {N: aws.String(fmt.Sprint(time.Now().Unix()))},
			":one":  {N: aws.String("1")},
		},
	})
	if err != nil && err != ctx.Err() {
		awserr, ok := err.(awserr.Error)
		if !ok || awserr.Code() != "RequestCanceled" {
			log.Errorf("dynamodb: update %v: %v", k, err)
		}
	}
	return k, v, nil
}
Get returns the digest associated with key digest k. Lookup
returns an error flagged errors.NotExist when no such mapping
exists. Lookup also modifies the item's last-accessed time, which
can be used for LRU object garbage collection.
Get expands abbreviated keys by making use of a DynamoDB index.
package dydbassoc
import (
	"context"
	"fmt"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/grailbio/base/digest"
	"github.com/grailbio/base/limiter"
	"github.com/grailbio/base/retry"
	"github.com/grailbio/base/traverse"
	"github.com/grailbio/infra"
	"github.com/grailbio/reflow"
	"github.com/grailbio/reflow/assoc"
	"github.com/grailbio/reflow/errors"
	"github.com/grailbio/reflow/flow"
	"github.com/grailbio/reflow/liveset"
	"github.com/grailbio/reflow/log"
	"github.com/grailbio/reflow/metrics"
	"github.com/grailbio/reflow/pool"
	"golang.org/x/sync/errgroup"
	"golang.org/x/time/rate"
)
func init() {
	infra.Register("dynamodbassoc", new(Assoc))
}
const (
	getTimeout = 30 * time.Second
	writecap = 10
	readcap  = 20
)
var (
	colmap = map[assoc.Kind]string{
		assoc.Fileset:     "Value",
		assoc.Logs:        "Logs",
		assoc.Bundle:      "Bundle",
		assoc.ExecInspect: "ExecInspect",
	}
	backOffPolicy = retry.MaxRetries(retry.Backoff(2*time.Millisecond, time.Minute, 1), 10)
)
type Assoc struct {
	assoc.AssocFlagsTrait
	DB      dynamodbiface.DynamoDBAPI `yaml:"-"`
	Limiter *limiter.Limiter          `yaml:"-"`
	Labels pool.Labels `yaml:"-"`
	labelsOnce sync.Once `yaml:"-"`
	labels     []*string `yaml:"-"`
}
func (a *Assoc) String() string {
	return fmt.Sprintf("%T,TableName=%s", a, a.TableName)
}
func (a *Assoc) Help() string {
	return "configure an assoc using the provided DynamoDB table name"
}
func (a *Assoc) Init(sess *session.Session, labels pool.Labels) error {
	if a.TableName == "" {
		return fmt.Errorf("Assoc table name cannot be empty")
	}
	lim := limiter.New()
	lim.Release(32)
	a.DB = dynamodb.New(sess)
	a.Limiter = lim
	a.Labels = labels.Copy()
	return nil
}
func (a *Assoc) Setup(sess *session.Session, logger *log.Logger) error {
	log.Printf("creating DynamoDB table %s", a.TableName)
	db := dynamodb.New(sess)
	_, err := db.CreateTable(&dynamodb.CreateTableInput{
		AttributeDefinitions: []*dynamodb.AttributeDefinition{
			{
				AttributeName: aws.String("ID"),
				AttributeType: aws.String("S"),
			},
		},
		KeySchema: []*dynamodb.KeySchemaElement{
			{
				AttributeName: aws.String("ID"),
				KeyType:       aws.String("HASH"),
			},
		},
		BillingMode: aws.String("PAY_PER_REQUEST"),
		TableName:   aws.String(a.TableName),
	})
	if err != nil {
		if aerr, ok := err.(awserr.Error); !ok || aerr.Code() != dynamodb.ErrCodeResourceInUseException {
			log.Fatal(err)
		}
		log.Printf("dynamodb table %s already exists", a.TableName)
	} else {
		log.Printf("created DynamoDB table %s", a.TableName)
	}
	const indexName = "ID4-ID-index"
	var describe *dynamodb.DescribeTableOutput
	start := time.Now()
	for {
		describe, err = db.DescribeTable(&dynamodb.DescribeTableInput{
			TableName: aws.String(a.TableName),
		})
		if err != nil {
			return err
		}
		status := *describe.Table.TableStatus
		if status == "ACTIVE" {
			break
		}
		if time.Since(start) > time.Minute {
			return errors.New("waited for table to become active for too long; try again later")
		}
		log.Printf("waiting for table to become active; current status: %v", status)
		time.Sleep(4 * time.Second)
	}
	var exists bool
	for _, index := range describe.Table.GlobalSecondaryIndexes {
		if *index.IndexName == indexName {
			exists = true
			break
		}
	}
	if exists {
		log.Printf("dynamodb index %s already exists", indexName)
	} else {
		_, err = db.UpdateTable(&dynamodb.UpdateTableInput{
			TableName: aws.String(a.TableName),
			AttributeDefinitions: []*dynamodb.AttributeDefinition{
				{
					AttributeName: aws.String("ID"),
					AttributeType: aws.String("S"),
				},
				{
					AttributeName: aws.String("ID4"),
					AttributeType: aws.String("S"),
				},
			},
			GlobalSecondaryIndexUpdates: []*dynamodb.GlobalSecondaryIndexUpdate{
				{
					Create: &dynamodb.CreateGlobalSecondaryIndexAction{
						IndexName: aws.String(indexName),
						KeySchema: []*dynamodb.KeySchemaElement{
							{
								KeyType:       aws.String("HASH"),
								AttributeName: aws.String("ID4"),
							},
							{
								KeyType:       aws.String("RANGE"),
								AttributeName: aws.String("ID"),
							},
						},
						Projection: &dynamodb.Projection{
							ProjectionType: aws.String("ALL"),
						},
					},
				},
			},
		})
		if err != nil {
			return errors.E("error creating secondary index: %v", err)
		}
		log.Printf("created secondary index %s", indexName)
	}
	return nil
}
func (a *Assoc) Version() int {
	return 1
}
func (a *Assoc) Store(ctx context.Context, kind assoc.Kind, k, v digest.Digest) error {
	reqStartTime := time.Now()
	defer func() {
		metrics.GetDydbassocOpLatencySecondsHistogram(ctx, "store").Observe(time.Now().Sub(reqStartTime).Seconds())
	}()
	switch kind {
	case assoc.Fileset, assoc.ExecInspect, assoc.Logs, assoc.Bundle:
	default:
		return errors.E(errors.NotSupported, errors.Errorf("mappings of kind %v are not supported", kind))
	}
	if err := a.Limiter.Acquire(ctx, 1); err != nil {
		return err
	}
	defer a.Limiter.Release(1)
	updateExpr, attrValues, attrNames := a.getUpdateComponents(kind, k, v)
	input := &dynamodb.UpdateItemInput{
		Key: map[string]*dynamodb.AttributeValue{
			"ID": {S: aws.String(k.String())},
		},
		UpdateExpression: &updateExpr,
		TableName:        aws.String(a.TableName),
	}
	if len(attrValues) > 0 {
		input.ExpressionAttributeValues = attrValues
	}
	if len(attrNames) > 0 {
		input.ExpressionAttributeNames = attrNames
	}
	_, err := a.DB.UpdateItemWithContext(ctx, input)
	if err == nil {
		return nil
	}
	awserr, ok := err.(awserr.Error)
	if !ok {
		return err
	}
	switch awserr.Code() {
	case "ConditionalCheckFailedException":
		return errors.E(errors.Precondition, err)
	}
	return err
}
func (a *Assoc) Delete(ctx context.Context, k digest.Digest) error {
	reqStartTime := time.Now()
	defer func() {
		metrics.GetDydbassocOpLatencySecondsHistogram(ctx, "delete").Observe(time.Now().Sub(reqStartTime).Seconds())
	}()
	input := &dynamodb.DeleteItemInput{
		Key: map[string]*dynamodb.AttributeValue{
			"ID": {
				S: aws.String(k.String()),
			},
		},
		ReturnValues: aws.String("ALL_OLD"),
		TableName:    aws.String(a.TableName),
	}
	output, err := a.DB.DeleteItemWithContext(ctx, input)
	if err != nil {
		return err
	}
	var keyFound bool
	if output != nil {
		_, keyFound = output.Attributes["ID"]
	}
	if !keyFound {
		return errors.E(errors.NotExist, fmt.Errorf("key %s not found", k))
	}
	return nil
}
func (a *Assoc) getUpdateComponents(kind assoc.Kind, k, v digest.Digest) (expr string, av map[string]*dynamodb.AttributeValue, an map[string]*string) {
	av = make(map[string]*dynamodb.AttributeValue)
	an = make(map[string]*string)
	switch kind {
	case assoc.Fileset:
		k4 := k
		k4.Truncate(4)
		switch {
		case v.IsZero():
			expr = "REMOVE #v"
		default:
			expr = "SET #v = :value, ID4 = :id4, LastAccessTime = :lastaccess"
			av[":value"] = &dynamodb.AttributeValue{S: aws.String(v.String())}
			av[":id4"] = &dynamodb.AttributeValue{S: aws.String(k4.HexN(4))}
			av[":lastaccess"] = &dynamodb.AttributeValue{N: aws.String(fmt.Sprint(time.Now().Unix()))}
		}
		an["#v"] = aws.String("Value")
	case assoc.ExecInspect:
		switch {
		case v.IsZero():
			expr = "REMOVE ExecInspect"
		default:
			expr = "SET ExecInspect = list_append(:inspect, if_not_exists(ExecInspect, :empty_list))"
			av[":inspect"] = &dynamodb.AttributeValue{L: []*dynamodb.AttributeValue{&dynamodb.AttributeValue{S: aws.String(v.String())}}}
			av[":empty_list"] = &dynamodb.AttributeValue{L: []*dynamodb.AttributeValue{}}
		}
	case assoc.Logs:
		switch {
		case v.IsZero():
			expr = "REMOVE Logs"
		default:
			expr = "SET Logs = list_append(:logs, if_not_exists(Logs, :empty_list))"
			av[":logs"] = &dynamodb.AttributeValue{L: []*dynamodb.AttributeValue{&dynamodb.AttributeValue{S: aws.String(v.String())}}}
			av[":empty_list"] = &dynamodb.AttributeValue{L: []*dynamodb.AttributeValue{}}
		}
	case assoc.Bundle:
		k4 := k
		k4.Truncate(4)
		switch {
		case v.IsZero():
			expr = "REMOVE Bundle"
		default:
			expr = "SET ID4 = :id4, Bundle= list_append(:bundle, if_not_exists(Bundle, :empty_list))"
			av[":bundle"] = &dynamodb.AttributeValue{L: []*dynamodb.AttributeValue{&dynamodb.AttributeValue{S: aws.String(v.String())}}}
			av[":id4"] = &dynamodb.AttributeValue{S: aws.String(k4.HexN(4))}
			av[":empty_list"] = &dynamodb.AttributeValue{L: []*dynamodb.AttributeValue{}}
		}
	}
	if !v.IsZero() && len(a.Labels) > 0 {
		a.labelsOnce.Do(func() {
			for k, v := range a.Labels {
				a.labels = append(a.labels, aws.String(fmt.Sprintf("%s=%s", k, v)))
			}
		})
		an["#l"] = aws.String("Labels")
		expr += " ADD #l :labels"
		av[":labels"] = &dynamodb.AttributeValue{SS: a.labels}
	}
	return
}