Auto-scale DynamoDB provision based on Prometheus metrics

cmd/table-manager/main.go

@@ -37,8 +37,8 @@ func main() {
 		schemaConfig.IndexTables.WriteScale.Enabled ||
 		schemaConfig.ChunkTables.InactiveWriteScale.Enabled ||
 		schemaConfig.IndexTables.InactiveWriteScale.Enabled) &&
-		storageConfig.AWSStorageConfig.ApplicationAutoScaling.URL == nil {
-		level.Error(util.Logger).Log("msg", "WriteScale is enabled but no ApplicationAutoScaling URL has been provided")
+		(storageConfig.AWSStorageConfig.ApplicationAutoScaling.URL == nil && storageConfig.AWSStorageConfig.Metrics.URL == "") {
+		level.Error(util.Logger).Log("msg", "WriteScale is enabled but no ApplicationAutoScaling or Metrics URL has been provided")
 		os.Exit(1)
 	}
 

pkg/chunk/aws/aws_autoscaling.go

@@ -0,0 +1,226 @@
+package aws
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/applicationautoscaling"
+	"github.com/aws/aws-sdk-go/service/applicationautoscaling/applicationautoscalingiface"
+	"github.com/go-kit/kit/log/level"
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/weaveworks/common/instrument"
+	"github.com/weaveworks/cortex/pkg/chunk"
+	"github.com/weaveworks/cortex/pkg/util"
+)
+
+const (
+	autoScalingPolicyNamePrefix = "DynamoScalingPolicy_cortex_"
+)
+
+var applicationAutoScalingRequestDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{
+	Namespace: "cortex",
+	Name:      "application_autoscaling_request_duration_seconds",
+	Help:      "Time spent doing ApplicationAutoScaling requests.",
+
+	// AWS latency seems to range from a few ms to a few sec. So use 8 buckets
+	// from 128us to 2s. TODO: Confirm that this is the case for ApplicationAutoScaling.
+	Buckets: prometheus.ExponentialBuckets(0.000128, 4, 8),
+}, []string{"operation", "status_code"})
+
+func init() {
+	prometheus.MustRegister(applicationAutoScalingRequestDuration)
+}
+
+type awsAutoscale struct {
+	call                   callManager
+	ApplicationAutoScaling applicationautoscalingiface.ApplicationAutoScalingAPI
+}
+
+func newAWSAutoscale(cfg DynamoDBConfig, callManager callManager) (*awsAutoscale, error) {
+	session, err := awsSessionFromURL(cfg.ApplicationAutoScaling.URL)
+	if err != nil {
+		return nil, err
+	}
+	return &awsAutoscale{
+		call: callManager,
+		ApplicationAutoScaling: applicationautoscaling.New(session),
+	}, nil
+}
+
+func (a *awsAutoscale) PostCreateTable(ctx context.Context, desc chunk.TableDesc) error {
+	if desc.WriteScale.Enabled {
+		return a.enableAutoScaling(ctx, desc)
+	}
+	return nil
+}
+
+func (a *awsAutoscale) DescribeTable(ctx context.Context, desc *chunk.TableDesc) error {
+	err := a.call.backoffAndRetry(ctx, func(ctx context.Context) error {
+		return instrument.TimeRequestHistogram(ctx, "ApplicationAutoScaling.DescribeScalableTargetsWithContext", applicationAutoScalingRequestDuration, func(ctx context.Context) error {
+			out, err := a.ApplicationAutoScaling.DescribeScalableTargetsWithContext(ctx, &applicationautoscaling.DescribeScalableTargetsInput{
+				ResourceIds:       []*string{aws.String("table/" + desc.Name)},
+				ScalableDimension: aws.String("dynamodb:table:WriteCapacityUnits"),
+				ServiceNamespace:  aws.String("dynamodb"),
+			})
+			if err != nil {
+				return err
+			}
+			switch l := len(out.ScalableTargets); l {
+			case 0:
+				return err
+			case 1:
+				desc.WriteScale.Enabled = true
+				if target := out.ScalableTargets[0]; target != nil {
+					if target.RoleARN != nil {
+						desc.WriteScale.RoleARN = *target.RoleARN
+					}
+					if target.MinCapacity != nil {
+						desc.WriteScale.MinCapacity = *target.MinCapacity
+					}
+					if target.MaxCapacity != nil {
+						desc.WriteScale.MaxCapacity = *target.MaxCapacity
+					}
+				}
+				return err
+			default:
+				return fmt.Errorf("more than one scalable target found for DynamoDB table")
+			}
+		})
+	})
+	if err != nil {
+		return err
+	}
+
+	err = a.call.backoffAndRetry(ctx, func(ctx context.Context) error {
+		return instrument.TimeRequestHistogram(ctx, "ApplicationAutoScaling.DescribeScalingPoliciesWithContext", applicationAutoScalingRequestDuration, func(ctx context.Context) error {
+			out, err := a.ApplicationAutoScaling.DescribeScalingPoliciesWithContext(ctx, &applicationautoscaling.DescribeScalingPoliciesInput{
+				PolicyNames:       []*string{aws.String(autoScalingPolicyNamePrefix + desc.Name)},
+				ResourceId:        aws.String("table/" + desc.Name),
+				ScalableDimension: aws.String("dynamodb:table:WriteCapacityUnits"),
+				ServiceNamespace:  aws.String("dynamodb"),
+			})
+			if err != nil {
+				return err
+			}
+			switch l := len(out.ScalingPolicies); l {
+			case 0:
+				return err
+			case 1:
+				config := out.ScalingPolicies[0].TargetTrackingScalingPolicyConfiguration
+				if config != nil {
+					if config.ScaleInCooldown != nil {
+						desc.WriteScale.InCooldown = *config.ScaleInCooldown
+					}
+					if config.ScaleOutCooldown != nil {
+						desc.WriteScale.OutCooldown = *config.ScaleOutCooldown
+					}
+					if config.TargetValue != nil {
+						desc.WriteScale.TargetValue = *config.TargetValue
+					}
+				}
+				return err
+			default:
+				return fmt.Errorf("more than one scaling policy found for DynamoDB table")
+			}
+		})
+	})
+	return err
+}
+
+func (a *awsAutoscale) UpdateTable(ctx context.Context, current chunk.TableDesc, expected *chunk.TableDesc) error {
+	var err error
+	if !current.WriteScale.Enabled {
+		if expected.WriteScale.Enabled {
+			level.Info(util.Logger).Log("msg", "enabling autoscaling on table", "table")
+			err = a.enableAutoScaling(ctx, *expected)
+		}
+	} else {
+		if !expected.WriteScale.Enabled {
+			level.Info(util.Logger).Log("msg", "disabling autoscaling on table", "table")
+			err = a.disableAutoScaling(ctx, *expected)
+		} else if current.WriteScale != expected.WriteScale {
+			level.Info(util.Logger).Log("msg", "enabling autoscaling on table", "table")
+			err = a.enableAutoScaling(ctx, *expected)
+		}
+	}
+	return err
+}
+
+func (a *awsAutoscale) enableAutoScaling(ctx context.Context, desc chunk.TableDesc) error {
+	// Registers or updates a scalable target
+	if err := a.call.backoffAndRetry(ctx, func(ctx context.Context) error {
+		return instrument.TimeRequestHistogram(ctx, "ApplicationAutoScaling.RegisterScalableTarget", applicationAutoScalingRequestDuration, func(ctx context.Context) error {
+			input := &applicationautoscaling.RegisterScalableTargetInput{
+				MinCapacity:       aws.Int64(desc.WriteScale.MinCapacity),
+				MaxCapacity:       aws.Int64(desc.WriteScale.MaxCapacity),
+				ResourceId:        aws.String("table/" + desc.Name),
+				RoleARN:           aws.String(desc.WriteScale.RoleARN),
+				ScalableDimension: aws.String("dynamodb:table:WriteCapacityUnits"),
+				ServiceNamespace:  aws.String("dynamodb"),
+			}
+			_, err := a.ApplicationAutoScaling.RegisterScalableTarget(input)
+			if err != nil {
+				return err
+			}
+			return nil
+		})
+	}); err != nil {
+		return err
+	}
+
+	// Puts or updates a scaling policy
+	return a.call.backoffAndRetry(ctx, func(ctx context.Context) error {
+		return instrument.TimeRequestHistogram(ctx, "ApplicationAutoScaling.PutScalingPolicy", applicationAutoScalingRequestDuration, func(ctx context.Context) error {
+			input := &applicationautoscaling.PutScalingPolicyInput{
+				PolicyName:        aws.String(autoScalingPolicyNamePrefix + desc.Name),
+				PolicyType:        aws.String("TargetTrackingScaling"),
+				ResourceId:        aws.String("table/" + desc.Name),
+				ScalableDimension: aws.String("dynamodb:table:WriteCapacityUnits"),
+				ServiceNamespace:  aws.String("dynamodb"),
+				TargetTrackingScalingPolicyConfiguration: &applicationautoscaling.TargetTrackingScalingPolicyConfiguration{
+					PredefinedMetricSpecification: &applicationautoscaling.PredefinedMetricSpecification{
+						PredefinedMetricType: aws.String("DynamoDBWriteCapacityUtilization"),
+					},
+					ScaleInCooldown:  aws.Int64(desc.WriteScale.InCooldown),
+					ScaleOutCooldown: aws.Int64(desc.WriteScale.OutCooldown),
+					TargetValue:      aws.Float64(desc.WriteScale.TargetValue),
+				},
+			}
+			_, err := a.ApplicationAutoScaling.PutScalingPolicy(input)
+			return err
+		})
+	})
+}
+
+func (a *awsAutoscale) disableAutoScaling(ctx context.Context, desc chunk.TableDesc) error {
+	// Deregister scalable target
+	if err := a.call.backoffAndRetry(ctx, func(ctx context.Context) error {
+		return instrument.TimeRequestHistogram(ctx, "ApplicationAutoScaling.DeregisterScalableTarget", applicationAutoScalingRequestDuration, func(ctx context.Context) error {
+			input := &applicationautoscaling.DeregisterScalableTargetInput{
+				ResourceId:        aws.String("table/" + desc.Name),
+				ScalableDimension: aws.String("dynamodb:table:WriteCapacityUnits"),
+				ServiceNamespace:  aws.String("dynamodb"),
+			}
+			_, err := a.ApplicationAutoScaling.DeregisterScalableTarget(input)
+			return err
+		})
+	}); err != nil {
+		return err
+	}
+
+	// Delete scaling policy
+	return a.call.backoffAndRetry(ctx, func(ctx context.Context) error {
+		return instrument.TimeRequestHistogram(ctx, "ApplicationAutoScaling.DeleteScalingPolicy", applicationAutoScalingRequestDuration, func(ctx context.Context) error {
+			input := &applicationautoscaling.DeleteScalingPolicyInput{
+				PolicyName:        aws.String(autoScalingPolicyNamePrefix + desc.Name),
+				ResourceId:        aws.String("table/" + desc.Name),
+				ScalableDimension: aws.String("dynamodb:table:WriteCapacityUnits"),
+				ServiceNamespace:  aws.String("dynamodb"),
+			}
+			_, err := a.ApplicationAutoScaling.DeleteScalingPolicy(input)
+			return err
+		})
+	})
+}