github.com/sequix/cortex@v1.1.6/pkg/chunk/aws/dynamodb_table_client.go (about)

     1  package aws
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  
     7  	"github.com/aws/aws-sdk-go/aws"
     8  	"github.com/aws/aws-sdk-go/aws/awserr"
     9  	"github.com/aws/aws-sdk-go/service/dynamodb"
    10  	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
    11  	"github.com/go-kit/kit/log/level"
    12  	"github.com/pkg/errors"
    13  	"golang.org/x/time/rate"
    14  
    15  	"github.com/sequix/cortex/pkg/chunk"
    16  	"github.com/sequix/cortex/pkg/util"
    17  	"github.com/weaveworks/common/instrument"
    18  )
    19  
    20  // Pluggable auto-scaler implementation
    21  type autoscale interface {
    22  	PostCreateTable(ctx context.Context, desc chunk.TableDesc) error
    23  	// This whole interface is very similar to chunk.TableClient, but
    24  	// DescribeTable needs to mutate desc
    25  	DescribeTable(ctx context.Context, desc *chunk.TableDesc) error
    26  	UpdateTable(ctx context.Context, current chunk.TableDesc, expected *chunk.TableDesc) error
    27  }
    28  
    29  type callManager struct {
    30  	limiter       *rate.Limiter
    31  	backoffConfig util.BackoffConfig
    32  }
    33  
    34  type dynamoTableClient struct {
    35  	DynamoDB    dynamodbiface.DynamoDBAPI
    36  	callManager callManager
    37  	autoscale   autoscale
    38  }
    39  
    40  // NewDynamoDBTableClient makes a new DynamoTableClient.
    41  func NewDynamoDBTableClient(cfg DynamoDBConfig) (chunk.TableClient, error) {
    42  	dynamoDB, err := dynamoClientFromURL(cfg.DynamoDB.URL)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  
    47  	callManager := callManager{
    48  		limiter:       rate.NewLimiter(rate.Limit(cfg.APILimit), 1),
    49  		backoffConfig: cfg.backoffConfig,
    50  	}
    51  
    52  	var autoscale autoscale
    53  	if cfg.ApplicationAutoScaling.URL != nil {
    54  		autoscale, err = newAWSAutoscale(cfg, callManager)
    55  		if err != nil {
    56  			return nil, err
    57  		}
    58  	}
    59  
    60  	if cfg.Metrics.URL != "" {
    61  		autoscale, err = newMetrics(cfg)
    62  		if err != nil {
    63  			return nil, err
    64  		}
    65  	}
    66  
    67  	return dynamoTableClient{
    68  		DynamoDB:    dynamoDB,
    69  		callManager: callManager,
    70  		autoscale:   autoscale,
    71  	}, nil
    72  }
    73  
    74  func (d *dynamoTableClient) Stop() {
    75  }
    76  
    77  func (d dynamoTableClient) backoffAndRetry(ctx context.Context, fn func(context.Context) error) error {
    78  	return d.callManager.backoffAndRetry(ctx, fn)
    79  }
    80  
    81  func (d callManager) backoffAndRetry(ctx context.Context, fn func(context.Context) error) error {
    82  	if d.limiter != nil { // Tests will have a nil limiter.
    83  		d.limiter.Wait(ctx)
    84  	}
    85  
    86  	backoff := util.NewBackoff(ctx, d.backoffConfig)
    87  	for backoff.Ongoing() {
    88  		if err := fn(ctx); err != nil {
    89  			if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "ThrottlingException" {
    90  				level.Warn(util.WithContext(ctx, util.Logger)).Log("msg", "got error, backing off and retrying", "err", err, "retry", backoff.NumRetries())
    91  				backoff.Wait()
    92  				continue
    93  			} else {
    94  				return err
    95  			}
    96  		}
    97  		return nil
    98  	}
    99  	return backoff.Err()
   100  }
   101  
   102  func (d dynamoTableClient) ListTables(ctx context.Context) ([]string, error) {
   103  	table := []string{}
   104  	err := d.backoffAndRetry(ctx, func(ctx context.Context) error {
   105  		return instrument.CollectedRequest(ctx, "DynamoDB.ListTablesPages", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   106  			return d.DynamoDB.ListTablesPagesWithContext(ctx, &dynamodb.ListTablesInput{}, func(resp *dynamodb.ListTablesOutput, _ bool) bool {
   107  				for _, s := range resp.TableNames {
   108  					table = append(table, *s)
   109  				}
   110  				return true
   111  			})
   112  		})
   113  	})
   114  	return table, err
   115  }
   116  
   117  func chunkTagsToDynamoDB(ts chunk.Tags) []*dynamodb.Tag {
   118  	var result []*dynamodb.Tag
   119  	for k, v := range ts {
   120  		result = append(result, &dynamodb.Tag{
   121  			Key:   aws.String(k),
   122  			Value: aws.String(v),
   123  		})
   124  	}
   125  	return result
   126  }
   127  
   128  func (d dynamoTableClient) CreateTable(ctx context.Context, desc chunk.TableDesc) error {
   129  	var tableARN *string
   130  	if err := d.backoffAndRetry(ctx, func(ctx context.Context) error {
   131  		return instrument.CollectedRequest(ctx, "DynamoDB.CreateTable", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   132  			input := &dynamodb.CreateTableInput{
   133  				TableName: aws.String(desc.Name),
   134  				AttributeDefinitions: []*dynamodb.AttributeDefinition{
   135  					{
   136  						AttributeName: aws.String(hashKey),
   137  						AttributeType: aws.String(dynamodb.ScalarAttributeTypeS),
   138  					},
   139  					{
   140  						AttributeName: aws.String(rangeKey),
   141  						AttributeType: aws.String(dynamodb.ScalarAttributeTypeB),
   142  					},
   143  				},
   144  				KeySchema: []*dynamodb.KeySchemaElement{
   145  					{
   146  						AttributeName: aws.String(hashKey),
   147  						KeyType:       aws.String(dynamodb.KeyTypeHash),
   148  					},
   149  					{
   150  						AttributeName: aws.String(rangeKey),
   151  						KeyType:       aws.String(dynamodb.KeyTypeRange),
   152  					},
   153  				},
   154  				ProvisionedThroughput: &dynamodb.ProvisionedThroughput{
   155  					ReadCapacityUnits:  aws.Int64(desc.ProvisionedRead),
   156  					WriteCapacityUnits: aws.Int64(desc.ProvisionedWrite),
   157  				},
   158  			}
   159  			output, err := d.DynamoDB.CreateTableWithContext(ctx, input)
   160  			if err != nil {
   161  				return err
   162  			}
   163  			if output.TableDescription != nil {
   164  				tableARN = output.TableDescription.TableArn
   165  			}
   166  			return nil
   167  		})
   168  	}); err != nil {
   169  		return err
   170  	}
   171  
   172  	if d.autoscale != nil {
   173  		err := d.autoscale.PostCreateTable(ctx, desc)
   174  		if err != nil {
   175  			return err
   176  		}
   177  	}
   178  
   179  	tags := chunkTagsToDynamoDB(desc.Tags)
   180  	if len(tags) > 0 {
   181  		return d.backoffAndRetry(ctx, func(ctx context.Context) error {
   182  			return instrument.CollectedRequest(ctx, "DynamoDB.TagResource", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   183  				_, err := d.DynamoDB.TagResourceWithContext(ctx, &dynamodb.TagResourceInput{
   184  					ResourceArn: tableARN,
   185  					Tags:        tags,
   186  				})
   187  				if relevantError(err) {
   188  					return err
   189  				}
   190  				return nil
   191  			})
   192  		})
   193  	}
   194  	return nil
   195  }
   196  
   197  func (d dynamoTableClient) DeleteTable(ctx context.Context, name string) error {
   198  	if err := d.backoffAndRetry(ctx, func(ctx context.Context) error {
   199  		return instrument.CollectedRequest(ctx, "DynamoDB.DeleteTable", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   200  			input := &dynamodb.DeleteTableInput{TableName: aws.String(name)}
   201  			_, err := d.DynamoDB.DeleteTableWithContext(ctx, input)
   202  			if err != nil {
   203  				return err
   204  			}
   205  
   206  			return nil
   207  		})
   208  	}); err != nil {
   209  		return err
   210  	}
   211  
   212  	return nil
   213  }
   214  
   215  func (d dynamoTableClient) DescribeTable(ctx context.Context, name string) (desc chunk.TableDesc, isActive bool, err error) {
   216  	var tableARN *string
   217  	err = d.backoffAndRetry(ctx, func(ctx context.Context) error {
   218  		return instrument.CollectedRequest(ctx, "DynamoDB.DescribeTable", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   219  			out, err := d.DynamoDB.DescribeTableWithContext(ctx, &dynamodb.DescribeTableInput{
   220  				TableName: aws.String(name),
   221  			})
   222  			if err != nil {
   223  				return err
   224  			}
   225  			desc.Name = name
   226  			if out.Table != nil {
   227  				if provision := out.Table.ProvisionedThroughput; provision != nil {
   228  					if provision.ReadCapacityUnits != nil {
   229  						desc.ProvisionedRead = *provision.ReadCapacityUnits
   230  					}
   231  					if provision.WriteCapacityUnits != nil {
   232  						desc.ProvisionedWrite = *provision.WriteCapacityUnits
   233  					}
   234  				}
   235  				if out.Table.TableStatus != nil {
   236  					isActive = (*out.Table.TableStatus == dynamodb.TableStatusActive)
   237  				}
   238  				if out.Table.BillingModeSummary != nil {
   239  					desc.UseOnDemandIOMode = *out.Table.BillingModeSummary.BillingMode == dynamodb.BillingModePayPerRequest
   240  				}
   241  				tableARN = out.Table.TableArn
   242  			}
   243  			return err
   244  		})
   245  	})
   246  	if err != nil {
   247  		return
   248  	}
   249  
   250  	err = d.backoffAndRetry(ctx, func(ctx context.Context) error {
   251  		return instrument.CollectedRequest(ctx, "DynamoDB.ListTagsOfResource", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   252  			out, err := d.DynamoDB.ListTagsOfResourceWithContext(ctx, &dynamodb.ListTagsOfResourceInput{
   253  				ResourceArn: tableARN,
   254  			})
   255  			if relevantError(err) {
   256  				return err
   257  			}
   258  			desc.Tags = make(map[string]string, len(out.Tags))
   259  			for _, tag := range out.Tags {
   260  				desc.Tags[*tag.Key] = *tag.Value
   261  			}
   262  			return nil
   263  		})
   264  	})
   265  
   266  	if d.autoscale != nil {
   267  		err = d.autoscale.DescribeTable(ctx, &desc)
   268  	}
   269  	return
   270  }
   271  
   272  // Filter out errors that we don't want to see
   273  // (currently only relevant in integration tests)
   274  func relevantError(err error) bool {
   275  	if err == nil {
   276  		return false
   277  	}
   278  	if strings.Contains(err.Error(), "Tagging is not currently supported in DynamoDB Local.") {
   279  		return false
   280  	}
   281  	return true
   282  }
   283  
   284  func (d dynamoTableClient) UpdateTable(ctx context.Context, current, expected chunk.TableDesc) error {
   285  	if d.autoscale != nil {
   286  		err := d.autoscale.UpdateTable(ctx, current, &expected)
   287  		if err != nil {
   288  			return err
   289  		}
   290  	}
   291  	level.Debug(util.Logger).Log("msg", "Updating Table",
   292  		"expectedWrite", expected.ProvisionedWrite,
   293  		"currentWrite", current.ProvisionedWrite,
   294  		"expectedRead", expected.ProvisionedRead,
   295  		"currentRead", current.ProvisionedRead,
   296  		"expectedOnDemandMode", expected.UseOnDemandIOMode,
   297  		"currentOnDemandMode", current.UseOnDemandIOMode)
   298  	if (current.ProvisionedRead != expected.ProvisionedRead ||
   299  		current.ProvisionedWrite != expected.ProvisionedWrite) &&
   300  		!expected.UseOnDemandIOMode {
   301  		level.Info(util.Logger).Log("msg", "updating provisioned throughput on table", "table", expected.Name, "old_read", current.ProvisionedRead, "old_write", current.ProvisionedWrite, "new_read", expected.ProvisionedRead, "new_write", expected.ProvisionedWrite)
   302  		if err := d.backoffAndRetry(ctx, func(ctx context.Context) error {
   303  			return instrument.CollectedRequest(ctx, "DynamoDB.UpdateTable", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   304  				var dynamoBillingMode string
   305  				updateTableInput := &dynamodb.UpdateTableInput{TableName: aws.String(expected.Name),
   306  					ProvisionedThroughput: &dynamodb.ProvisionedThroughput{
   307  						ReadCapacityUnits:  aws.Int64(expected.ProvisionedRead),
   308  						WriteCapacityUnits: aws.Int64(expected.ProvisionedWrite),
   309  					},
   310  				}
   311  				// we need this to be a separate check for the billing mode, as aws returns
   312  				// an error if we set a table to the billing mode it is currently on.
   313  				if current.UseOnDemandIOMode != expected.UseOnDemandIOMode {
   314  					dynamoBillingMode = dynamodb.BillingModeProvisioned
   315  					level.Info(util.Logger).Log("msg", "updating billing mode on table", "table", expected.Name, "old_mode", current.UseOnDemandIOMode, "new_mode", expected.UseOnDemandIOMode)
   316  					updateTableInput.BillingMode = aws.String(dynamoBillingMode)
   317  				}
   318  
   319  				_, err := d.DynamoDB.UpdateTableWithContext(ctx, updateTableInput)
   320  				return err
   321  			})
   322  		}); err != nil {
   323  			recordDynamoError(expected.Name, err, "DynamoDB.UpdateTable")
   324  			if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "LimitExceededException" {
   325  				level.Warn(util.Logger).Log("msg", "update limit exceeded", "err", err)
   326  			} else {
   327  				return err
   328  			}
   329  		}
   330  	} else if expected.UseOnDemandIOMode && current.UseOnDemandIOMode != expected.UseOnDemandIOMode {
   331  		// moved the enabling of OnDemand mode to it's own block to reduce complexities & interactions with the various
   332  		// settings used in provisioned mode. Unfortunately the boilerplate wrappers for retry and tracking needed to be copied.
   333  		if err := d.backoffAndRetry(ctx, func(ctx context.Context) error {
   334  			return instrument.CollectedRequest(ctx, "DynamoDB.UpdateTable", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   335  				level.Info(util.Logger).Log("msg", "updating billing mode on table", "table", expected.Name, "old_mode", current.UseOnDemandIOMode, "new_mode", expected.UseOnDemandIOMode)
   336  				updateTableInput := &dynamodb.UpdateTableInput{TableName: aws.String(expected.Name), BillingMode: aws.String(dynamodb.BillingModePayPerRequest)}
   337  				_, err := d.DynamoDB.UpdateTableWithContext(ctx, updateTableInput)
   338  				return err
   339  			})
   340  		}); err != nil {
   341  			recordDynamoError(expected.Name, err, "DynamoDB.UpdateTable")
   342  			if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "LimitExceededException" {
   343  				level.Warn(util.Logger).Log("msg", "update limit exceeded", "err", err)
   344  			} else {
   345  				return err
   346  			}
   347  		}
   348  	}
   349  
   350  	if !current.Tags.Equals(expected.Tags) {
   351  		var tableARN *string
   352  		if err := d.backoffAndRetry(ctx, func(ctx context.Context) error {
   353  			return instrument.CollectedRequest(ctx, "DynamoDB.DescribeTable", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   354  				out, err := d.DynamoDB.DescribeTableWithContext(ctx, &dynamodb.DescribeTableInput{
   355  					TableName: aws.String(expected.Name),
   356  				})
   357  				if err != nil {
   358  					return err
   359  				}
   360  				if out.Table != nil {
   361  					tableARN = out.Table.TableArn
   362  				}
   363  				return nil
   364  			})
   365  		}); err != nil {
   366  			return err
   367  		}
   368  
   369  		return d.backoffAndRetry(ctx, func(ctx context.Context) error {
   370  			return instrument.CollectedRequest(ctx, "DynamoDB.TagResource", dynamoRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   371  				_, err := d.DynamoDB.TagResourceWithContext(ctx, &dynamodb.TagResourceInput{
   372  					ResourceArn: tableARN,
   373  					Tags:        chunkTagsToDynamoDB(expected.Tags),
   374  				})
   375  				if relevantError(err) {
   376  					return errors.Wrap(err, "applying tags")
   377  				}
   378  				return nil
   379  			})
   380  		})
   381  	}
   382  	return nil
   383  }