diff --git a/cmd/pgoload/clients.go b/cmd/pgoload/clients.go new file mode 100644 index 0000000000..73e6b2e575 --- /dev/null +++ b/cmd/pgoload/clients.go @@ -0,0 +1,163 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + awscfg "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/cloudwatch" + "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs" + "github.com/aws/aws-sdk-go-v2/service/dynamodb" + "github.com/aws/aws-sdk-go-v2/service/ec2" + "github.com/aws/aws-sdk-go-v2/service/eventbridge" + "github.com/aws/aws-sdk-go-v2/service/iam" + "github.com/aws/aws-sdk-go-v2/service/kinesis" + "github.com/aws/aws-sdk-go-v2/service/kms" + "github.com/aws/aws-sdk-go-v2/service/lambda" + "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/aws/aws-sdk-go-v2/service/secretsmanager" + "github.com/aws/aws-sdk-go-v2/service/sfn" + "github.com/aws/aws-sdk-go-v2/service/sns" + "github.com/aws/aws-sdk-go-v2/service/sqs" + "github.com/aws/aws-sdk-go-v2/service/ssm" + "github.com/aws/aws-sdk-go-v2/service/sts" +) + +// clients bundles one SDK client per service exercised by pgoload, all +// pointed at the same Gopherstack endpoint. +type clients struct { + ddb *dynamodb.Client + s3 *s3.Client + sqs *sqs.Client + sns *sns.Client + kinesis *kinesis.Client + iam *iam.Client + sts *sts.Client + ssm *ssm.Client + secretsmanager *secretsmanager.Client + cloudwatch *cloudwatch.Client + logs *cloudwatchlogs.Client + ec2 *ec2.Client + lambda *lambda.Client + kms *kms.Client + eventbridge *eventbridge.Client + stepfunctions *sfn.Client +} + +// buildClients constructs every service client used by pgoload, using +// static test credentials against cfg.endpoint. +func buildClients(ctx context.Context, cfg config) (*clients, error) { + awsCfg, err := awscfg.LoadDefaultConfig(ctx, + awscfg.WithRegion(awsRegion), + awscfg.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("test", "test", "")), + ) + if err != nil { + return nil, fmt.Errorf("load aws config: %w", err) + } + + ep := aws.String(cfg.endpoint) + + return &clients{ + ddb: dynamodb.NewFromConfig(awsCfg, func(o *dynamodb.Options) { o.BaseEndpoint = ep }), + s3: s3.NewFromConfig(awsCfg, func(o *s3.Options) { + o.BaseEndpoint = ep + o.UsePathStyle = true + }), + sqs: sqs.NewFromConfig(awsCfg, func(o *sqs.Options) { o.BaseEndpoint = ep }), + sns: sns.NewFromConfig(awsCfg, func(o *sns.Options) { o.BaseEndpoint = ep }), + kinesis: kinesis.NewFromConfig(awsCfg, func(o *kinesis.Options) { o.BaseEndpoint = ep }), + iam: iam.NewFromConfig(awsCfg, func(o *iam.Options) { o.BaseEndpoint = ep }), + sts: sts.NewFromConfig(awsCfg, func(o *sts.Options) { o.BaseEndpoint = ep }), + ssm: ssm.NewFromConfig(awsCfg, func(o *ssm.Options) { o.BaseEndpoint = ep }), + secretsmanager: secretsmanager.NewFromConfig(awsCfg, func(o *secretsmanager.Options) { o.BaseEndpoint = ep }), + cloudwatch: cloudwatch.NewFromConfig(awsCfg, func(o *cloudwatch.Options) { o.BaseEndpoint = ep }), + logs: cloudwatchlogs.NewFromConfig(awsCfg, func(o *cloudwatchlogs.Options) { o.BaseEndpoint = ep }), + ec2: ec2.NewFromConfig(awsCfg, func(o *ec2.Options) { o.BaseEndpoint = ep }), + lambda: lambda.NewFromConfig(awsCfg, func(o *lambda.Options) { o.BaseEndpoint = ep }), + kms: kms.NewFromConfig(awsCfg, func(o *kms.Options) { o.BaseEndpoint = ep }), + eventbridge: eventbridge.NewFromConfig(awsCfg, func(o *eventbridge.Options) { o.BaseEndpoint = ep }), + stepfunctions: sfn.NewFromConfig(awsCfg, func(o *sfn.Options) { o.BaseEndpoint = ep }), + }, nil +} + +// resources holds the identifiers of everything provisioned once during +// setup and then reused by every breadth-scenario worker. +type resources struct { + roleArn string + queueURL string + queueArn string + topicArn string + kinesisShardID string + kmsKeyID string + functionName string + stateMachineArn string + s3Buckets []string +} + +// setupResources provisions the DynamoDB table, S3 buckets, and every +// breadth-scenario's resources, bounded by setupTimeout so a short +// -duration still leaves time for the load phase itself. Each service's +// setup is independent; a failure in one is fatal (unlike op-level errors +// during the load phase, which are only counted) because a missing +// resource would make that whole scenario error-spin instead of exercising +// its real hot paths. +func setupResources(ctx context.Context, cls *clients, log *slog.Logger) (*resources, error) { + setupCtx, cancel := context.WithTimeout(ctx, setupTimeout) + defer cancel() + + if err := ensureDDBTable(setupCtx, cls.ddb, log); err != nil { + return nil, fmt.Errorf("ddb table setup: %w", err) + } + + buckets, err := ensureS3Buckets(setupCtx, cls.s3, log) + if err != nil { + return nil, fmt.Errorf("s3 bucket setup: %w", err) + } + + res := &resources{s3Buckets: buckets} + + if res.roleArn, err = ensureIAMRole(setupCtx, cls.iam, log); err != nil { + return nil, fmt.Errorf("iam role setup: %w", err) + } + + if res.queueURL, res.queueArn, err = ensureSQSQueue(setupCtx, cls.sqs, log); err != nil { + return nil, fmt.Errorf("sqs queue setup: %w", err) + } + + if res.topicArn, err = ensureSNSTopic(setupCtx, cls.sns, res.queueArn, log); err != nil { + return nil, fmt.Errorf("sns topic setup: %w", err) + } + + if err = ensureSecrets(setupCtx, cls.secretsmanager, log); err != nil { + return nil, fmt.Errorf("secretsmanager setup: %w", err) + } + + if res.kinesisShardID, err = ensureKinesisStream(setupCtx, cls.kinesis, log); err != nil { + return nil, fmt.Errorf("kinesis stream setup: %w", err) + } + + if err = ensureLogGroup(setupCtx, cls.logs, log); err != nil { + return nil, fmt.Errorf("logs group setup: %w", err) + } + + if res.kmsKeyID, err = ensureKMSKey(setupCtx, cls.kms, log); err != nil { + return nil, fmt.Errorf("kms key setup: %w", err) + } + + if res.functionName, err = ensureLambdaFunction(setupCtx, cls.lambda, res.roleArn, log); err != nil { + return nil, fmt.Errorf("lambda function setup: %w", err) + } + + if err = ensureEventBridgeRule(setupCtx, cls.eventbridge, res.queueArn, log); err != nil { + return nil, fmt.Errorf("eventbridge rule setup: %w", err) + } + + if res.stateMachineArn, err = ensureStateMachine(setupCtx, cls.stepfunctions, res.roleArn, log); err != nil { + return nil, fmt.Errorf("step functions state machine setup: %w", err) + } + + return res, nil +} diff --git a/cmd/pgoload/cloudwatch.go b/cmd/pgoload/cloudwatch.go new file mode 100644 index 0000000000..6abe76694e --- /dev/null +++ b/cmd/pgoload/cloudwatch.go @@ -0,0 +1,83 @@ +package main + +import ( + "context" + "log/slog" + "strconv" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/cloudwatch" + cwtypes "github.com/aws/aws-sdk-go-v2/service/cloudwatch/types" +) + +const ( + cwNamespace = "pgoload" + cwMetricName = "Ops" + cwStatWindow = 10 * time.Minute + cwStatPeriodSeconds = 60 + cwDatapointsPerBatch = 3 +) + +// cloudwatchWorker repeatedly runs a mix of CloudWatch operations, staggered +// by workerID, until ctx is done. +// +//nolint:dupl // structurally mirrors eventBridgeWorker's op-table shape but wires an unrelated service +func cloudwatchWorker(ctx context.Context, cl *cloudwatch.Client, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { return cwPutMetricDataOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return cwPutMetricDataOp(ctx, cl, workerID, i) }, + func(ctx context.Context, _, _ int) error { return cwListMetricsOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return cwGetMetricStatisticsOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return cwDescribeAlarmsOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "cloudwatch", log) +} + +func cwPutMetricDataOp(ctx context.Context, cl *cloudwatch.Client, workerID, i int) error { + data := make([]cwtypes.MetricDatum, 0, cwDatapointsPerBatch) + for j := range cwDatapointsPerBatch { + data = append(data, cwtypes.MetricDatum{ + MetricName: aws.String(cwMetricName), + Value: aws.Float64(float64(i + j)), + Dimensions: []cwtypes.Dimension{ + {Name: aws.String("Worker"), Value: aws.String(strconv.Itoa(workerID))}, + }, + }) + } + + _, err := cl.PutMetricData(ctx, &cloudwatch.PutMetricDataInput{ + Namespace: aws.String(cwNamespace), + MetricData: data, + }) + + return err +} + +func cwListMetricsOp(ctx context.Context, cl *cloudwatch.Client) error { + _, err := cl.ListMetrics(ctx, &cloudwatch.ListMetricsInput{Namespace: aws.String(cwNamespace)}) + + return err +} + +func cwGetMetricStatisticsOp(ctx context.Context, cl *cloudwatch.Client) error { + now := time.Now() + + _, err := cl.GetMetricStatistics(ctx, &cloudwatch.GetMetricStatisticsInput{ + Namespace: aws.String(cwNamespace), + MetricName: aws.String(cwMetricName), + StartTime: aws.Time(now.Add(-cwStatWindow)), + EndTime: aws.Time(now), + Period: aws.Int32(cwStatPeriodSeconds), + Statistics: []cwtypes.Statistic{cwtypes.StatisticSum, cwtypes.StatisticAverage}, + }) + + return err +} + +func cwDescribeAlarmsOp(ctx context.Context, cl *cloudwatch.Client) error { + _, err := cl.DescribeAlarms(ctx, &cloudwatch.DescribeAlarmsInput{}) + + return err +} diff --git a/cmd/pgoload/common.go b/cmd/pgoload/common.go new file mode 100644 index 0000000000..aebd047b7c --- /dev/null +++ b/cmd/pgoload/common.go @@ -0,0 +1,49 @@ +package main + +import ( + "context" + "log/slog" + "sync/atomic" +) + +// opCounter tracks operation and error totals for one breadth-load scenario. +type opCounter struct { + ops atomic.Int64 + errors atomic.Int64 +} + +func (o *opCounter) record(err error) { + if err != nil { + o.errors.Add(1) + + return + } + + o.ops.Add(1) +} + +// opFunc is one operation exercised by a breadth-scenario worker. workerID +// and i together select the resource/key space touched by the call. +type opFunc func(ctx context.Context, workerID, i int) error + +// runOpLoop repeatedly runs a mix of operations, staggered by workerID so +// concurrent workers hit different operations on each tick, until ctx is +// done. It is the shared loop shape behind every breadth scenario added +// beyond the original DDB/S3 load (ddbWorker and s3Worker keep their own +// copy of this loop untouched). +func runOpLoop(ctx context.Context, workerID int, ops []opFunc, c *opCounter, scenario string, log *slog.Logger) { + for i := 0; ; i++ { + select { + case <-ctx.Done(): + return + default: + } + + err := ops[(workerID+i)%len(ops)](ctx, workerID, i) + c.record(err) + + if err != nil { + log.WarnContext(ctx, scenario+" operation failed", "worker", workerID, "iter", i, "error", err) + } + } +} diff --git a/cmd/pgoload/ec2.go b/cmd/pgoload/ec2.go new file mode 100644 index 0000000000..b714be9d91 --- /dev/null +++ b/cmd/pgoload/ec2.go @@ -0,0 +1,100 @@ +package main + +import ( + "context" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/ec2" + ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types" +) + +const ( + ec2ImageID = "ami-12345678" + ec2TagKey = "Owner" + ec2TagValue = "pgoload" + ec2TerminateCap = 5 +) + +// ec2Worker repeatedly runs a mix of EC2 operations, staggered by workerID, +// until ctx is done. Instance churn is self-balancing: run and terminate are +// both in the op mix, and terminate always clears out whatever the pgoload +// tag currently has, so the instance count stays bounded without any +// client-side bookkeeping across calls. +func ec2Worker(ctx context.Context, cl *ec2.Client, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, _, _ int) error { return ec2RunInstancesOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return ec2DescribeInstancesOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return ec2DescribeInstancesOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return ec2DescribeSecurityGroupsOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return ec2TerminateInstancesOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "ec2", log) +} + +func ec2OwnerFilter() []ec2types.Filter { + return []ec2types.Filter{ + {Name: aws.String("tag:" + ec2TagKey), Values: []string{ec2TagValue}}, + } +} + +func ec2RunInstancesOp(ctx context.Context, cl *ec2.Client) error { + const minMaxCount = 1 + + _, err := cl.RunInstances(ctx, &ec2.RunInstancesInput{ + ImageId: aws.String(ec2ImageID), + InstanceType: ec2types.InstanceTypeT2Micro, + MinCount: aws.Int32(minMaxCount), + MaxCount: aws.Int32(minMaxCount), + TagSpecifications: []ec2types.TagSpecification{ + { + ResourceType: ec2types.ResourceTypeInstance, + Tags: []ec2types.Tag{{Key: aws.String(ec2TagKey), Value: aws.String(ec2TagValue)}}, + }, + }, + }) + + return err +} + +func ec2DescribeInstancesOp(ctx context.Context, cl *ec2.Client) error { + _, err := cl.DescribeInstances(ctx, &ec2.DescribeInstancesInput{Filters: ec2OwnerFilter()}) + + return err +} + +func ec2DescribeSecurityGroupsOp(ctx context.Context, cl *ec2.Client) error { + _, err := cl.DescribeSecurityGroups(ctx, &ec2.DescribeSecurityGroupsInput{}) + + return err +} + +// ec2TerminateInstancesOp describes up to ec2TerminateCap pgoload-tagged +// instances and terminates them, keeping the emulator's instance count from +// growing without bound over a long run. +func ec2TerminateInstancesOp(ctx context.Context, cl *ec2.Client) error { + out, err := cl.DescribeInstances(ctx, &ec2.DescribeInstancesInput{ + Filters: ec2OwnerFilter(), + MaxResults: aws.Int32(ec2TerminateCap), + }) + if err != nil { + return err + } + + ids := make([]string, 0, ec2TerminateCap) + + for _, res := range out.Reservations { + for _, inst := range res.Instances { + ids = append(ids, aws.ToString(inst.InstanceId)) + } + } + + if len(ids) == 0 { + return nil + } + + _, err = cl.TerminateInstances(ctx, &ec2.TerminateInstancesInput{InstanceIds: ids}) + + return err +} diff --git a/cmd/pgoload/eventbridge.go b/cmd/pgoload/eventbridge.go new file mode 100644 index 0000000000..aede7bddce --- /dev/null +++ b/cmd/pgoload/eventbridge.go @@ -0,0 +1,89 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/eventbridge" + ebtypes "github.com/aws/aws-sdk-go-v2/service/eventbridge/types" +) + +const ( + ebRuleName = "pgoload-rule" + ebEventBus = "default" + ebSource = "pgoload" + ebDetailType = "pgoload.smoke" +) + +// ensureEventBridgeRule creates ebRuleName (PutRule is an upsert, so this is +// safe to re-run) and targets it at the SQS queue behind queueArn. +func ensureEventBridgeRule(ctx context.Context, cl *eventbridge.Client, queueArn string, log *slog.Logger) error { + if _, err := cl.PutRule(ctx, &eventbridge.PutRuleInput{ + Name: aws.String(ebRuleName), + ScheduleExpression: aws.String("rate(5 minutes)"), + }); err != nil { + return fmt.Errorf("put rule %s: %w", ebRuleName, err) + } + + if _, err := cl.PutTargets(ctx, &eventbridge.PutTargetsInput{ + Rule: aws.String(ebRuleName), + Targets: []ebtypes.Target{ + {Id: aws.String("pgoload-target"), Arn: aws.String(queueArn)}, + }, + }); err != nil { + log.WarnContext(ctx, "eventbridge put targets failed (continuing)", "error", err) + } + + return nil +} + +// eventBridgeWorker repeatedly runs a mix of EventBridge operations, +// staggered by workerID, until ctx is done. +// +//nolint:dupl // structurally mirrors cloudwatchWorker's op-table shape but wires an unrelated service +func eventBridgeWorker(ctx context.Context, cl *eventbridge.Client, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { return ebPutEventsOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return ebPutEventsOp(ctx, cl, workerID, i) }, + func(ctx context.Context, _, _ int) error { return ebListRulesOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return ebDescribeRuleOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return ebListTargetsByRuleOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "eventbridge", log) +} + +func ebPutEventsOp(ctx context.Context, cl *eventbridge.Client, workerID, i int) error { + _, err := cl.PutEvents(ctx, &eventbridge.PutEventsInput{ + Entries: []ebtypes.PutEventsRequestEntry{ + { + EventBusName: aws.String(ebEventBus), + Source: aws.String(ebSource), + DetailType: aws.String(ebDetailType), + Detail: aws.String(fmt.Sprintf(`{"worker":%d,"iter":%d}`, workerID, i)), + }, + }, + }) + + return err +} + +func ebListRulesOp(ctx context.Context, cl *eventbridge.Client) error { + _, err := cl.ListRules(ctx, &eventbridge.ListRulesInput{}) + + return err +} + +func ebDescribeRuleOp(ctx context.Context, cl *eventbridge.Client) error { + _, err := cl.DescribeRule(ctx, &eventbridge.DescribeRuleInput{Name: aws.String(ebRuleName)}) + + return err +} + +func ebListTargetsByRuleOp(ctx context.Context, cl *eventbridge.Client) error { + _, err := cl.ListTargetsByRule(ctx, &eventbridge.ListTargetsByRuleInput{Rule: aws.String(ebRuleName)}) + + return err +} diff --git a/cmd/pgoload/iam.go b/cmd/pgoload/iam.go new file mode 100644 index 0000000000..1bf3918ec3 --- /dev/null +++ b/cmd/pgoload/iam.go @@ -0,0 +1,127 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/iam" + iamtypes "github.com/aws/aws-sdk-go-v2/service/iam/types" +) + +const ( + iamRoleName = "pgoload-role" + + // iamUserRotation bounds the number of distinct IAM users pgoload + // cycles through so CreateUser/DeleteUser churn doesn't grow the + // backend's user set without limit over a long run. + iamUserRotation = 8 + + iamAssumeRolePolicy = `{"Version":"2012-10-17","Statement":[` + + `{"Effect":"Allow","Principal":{"Service":"lambda.amazonaws.com"},"Action":"sts:AssumeRole"}]}` + iamInlinePolicy = `{"Version":"2012-10-17","Statement":[` + + `{"Effect":"Allow","Action":"s3:GetObject","Resource":"*"}]}` +) + +// ensureIAMRole creates iamRoleName, tolerating one that already exists, and +// returns its ARN. Lambda and Step Functions setup reuse this role. +func ensureIAMRole(ctx context.Context, cl *iam.Client, log *slog.Logger) (string, error) { + out, err := cl.CreateRole(ctx, &iam.CreateRoleInput{ + RoleName: aws.String(iamRoleName), + AssumeRolePolicyDocument: aws.String(iamAssumeRolePolicy), + }) + if err == nil { + return aws.ToString(out.Role.Arn), nil + } + + var exists *iamtypes.EntityAlreadyExistsException + if !errors.As(err, &exists) { + return "", fmt.Errorf("create role %s: %w", iamRoleName, err) + } + + log.InfoContext(ctx, "iam role already exists", "role", iamRoleName) + + got, err := cl.GetRole(ctx, &iam.GetRoleInput{RoleName: aws.String(iamRoleName)}) + if err != nil { + return "", fmt.Errorf("get existing role %s: %w", iamRoleName, err) + } + + return aws.ToString(got.Role.Arn), nil +} + +func iamUserName(workerID, i int) string { + return fmt.Sprintf("pgoload-user-%d-%d", workerID, i%iamUserRotation) +} + +// iamWorker repeatedly runs a mix of IAM operations, staggered by workerID, +// until ctx is done. +func iamWorker(ctx context.Context, cl *iam.Client, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { return iamCreateUserOp(ctx, cl, workerID, i) }, + func(ctx context.Context, _, _ int) error { return iamGetRoleOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return iamListUsersOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return iamListRolesOp(ctx, cl) }, + func(ctx context.Context, workerID, i int) error { return iamPutRolePolicyOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return iamDeleteUserOp(ctx, cl, workerID, i) }, + } + + runOpLoop(ctx, workerID, ops, c, "iam", log) +} + +func iamCreateUserOp(ctx context.Context, cl *iam.Client, workerID, i int) error { + _, err := cl.CreateUser(ctx, &iam.CreateUserInput{UserName: aws.String(iamUserName(workerID, i))}) + if err == nil { + return nil + } + + var exists *iamtypes.EntityAlreadyExistsException + if errors.As(err, &exists) { + return nil + } + + return err +} + +func iamGetRoleOp(ctx context.Context, cl *iam.Client) error { + _, err := cl.GetRole(ctx, &iam.GetRoleInput{RoleName: aws.String(iamRoleName)}) + + return err +} + +func iamListUsersOp(ctx context.Context, cl *iam.Client) error { + _, err := cl.ListUsers(ctx, &iam.ListUsersInput{}) + + return err +} + +func iamListRolesOp(ctx context.Context, cl *iam.Client) error { + _, err := cl.ListRoles(ctx, &iam.ListRolesInput{}) + + return err +} + +func iamPutRolePolicyOp(ctx context.Context, cl *iam.Client, workerID, _ int) error { + _, err := cl.PutRolePolicy(ctx, &iam.PutRolePolicyInput{ + RoleName: aws.String(iamRoleName), + PolicyName: aws.String(fmt.Sprintf("pgoload-policy-%d", workerID)), + PolicyDocument: aws.String(iamInlinePolicy), + }) + + return err +} + +func iamDeleteUserOp(ctx context.Context, cl *iam.Client, workerID, i int) error { + _, err := cl.DeleteUser(ctx, &iam.DeleteUserInput{UserName: aws.String(iamUserName(workerID, i))}) + if err == nil { + return nil + } + + var notFound *iamtypes.NoSuchEntityException + if errors.As(err, ¬Found) { + return nil + } + + return err +} diff --git a/cmd/pgoload/kinesis.go b/cmd/pgoload/kinesis.go new file mode 100644 index 0000000000..0a9bbc78f7 --- /dev/null +++ b/cmd/pgoload/kinesis.go @@ -0,0 +1,161 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/kinesis" + kinesistypes "github.com/aws/aws-sdk-go-v2/service/kinesis/types" +) + +const ( + kinesisStreamName = "pgoload-stream" + kinesisShardCount = 1 + + kinesisPutBatchSize = 10 + kinesisPayloadSize = 256 +) + +// errStreamNotActive is returned when the Kinesis stream fails to reach the +// ACTIVE status within setupTimeout. +var errStreamNotActive = errors.New("kinesis stream did not become active in time") + +// ensureKinesisStream creates kinesisStreamName, tolerating one that already +// exists, waits for it to become ACTIVE, and returns its first shard's ID +// for the GetRecords op. +func ensureKinesisStream(ctx context.Context, cl *kinesis.Client, log *slog.Logger) (string, error) { + _, err := cl.CreateStream(ctx, &kinesis.CreateStreamInput{ + StreamName: aws.String(kinesisStreamName), + ShardCount: aws.Int32(kinesisShardCount), + }) + if err != nil { + var inUse *kinesistypes.ResourceInUseException + if !errors.As(err, &inUse) { + return "", fmt.Errorf("create stream %s: %w", kinesisStreamName, err) + } + + log.InfoContext(ctx, "kinesis stream already exists", "stream", kinesisStreamName) + } + + shardID, err := waitStreamActive(ctx, cl, log) + if err != nil { + return "", err + } + + return shardID, nil +} + +func waitStreamActive(ctx context.Context, cl *kinesis.Client, log *slog.Logger) (string, error) { + deadline := time.Now().Add(setupTimeout) + + for time.Now().Before(deadline) { + out, err := cl.DescribeStreamSummary(ctx, &kinesis.DescribeStreamSummaryInput{ + StreamName: aws.String(kinesisStreamName), + }) + if err == nil && out.StreamDescriptionSummary != nil && + out.StreamDescriptionSummary.StreamStatus == kinesistypes.StreamStatusActive { + shards, shardsErr := cl.ListShards(ctx, &kinesis.ListShardsInput{StreamName: aws.String(kinesisStreamName)}) + if shardsErr != nil || len(shards.Shards) == 0 { + return "", fmt.Errorf("list shards for %s: %w", kinesisStreamName, shardsErr) + } + + return aws.ToString(shards.Shards[0].ShardId), nil + } + + select { + case <-ctx.Done(): + return "", fmt.Errorf("waiting for stream %s: %w", kinesisStreamName, ctx.Err()) + case <-time.After(pollInterval): + } + } + + log.WarnContext(ctx, "gave up waiting for stream to become active", "stream", kinesisStreamName) + + return "", fmt.Errorf("%w: stream=%s", errStreamNotActive, kinesisStreamName) +} + +// kinesisWorker repeatedly runs a mix of Kinesis operations, staggered by +// workerID, until ctx is done. +func kinesisWorker( + ctx context.Context, + cl *kinesis.Client, + res *resources, + workerID int, + c *opCounter, + log *slog.Logger, +) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { return kinesisPutRecordOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return kinesisPutRecordOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return kinesisPutRecordsOp(ctx, cl, workerID, i) }, + func(ctx context.Context, _, _ int) error { + return kinesisGetRecordsOp(ctx, cl, res.kinesisShardID) + }, + func(ctx context.Context, _, _ int) error { return kinesisDescribeStreamOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "kinesis", log) +} + +func kinesisPartitionKey(workerID, i int) string { + return fmt.Sprintf("worker-%d-%d", workerID, i) +} + +func kinesisPutRecordOp(ctx context.Context, cl *kinesis.Client, workerID, i int) error { + _, err := cl.PutRecord(ctx, &kinesis.PutRecordInput{ + StreamName: aws.String(kinesisStreamName), + PartitionKey: aws.String(kinesisPartitionKey(workerID, i)), + Data: []byte(strings.Repeat("k", kinesisPayloadSize)), + }) + + return err +} + +func kinesisPutRecordsOp(ctx context.Context, cl *kinesis.Client, workerID, base int) error { + records := make([]kinesistypes.PutRecordsRequestEntry, 0, kinesisPutBatchSize) + for j := range kinesisPutBatchSize { + records = append(records, kinesistypes.PutRecordsRequestEntry{ + PartitionKey: aws.String(kinesisPartitionKey(workerID, base+j)), + Data: []byte(strings.Repeat("k", kinesisPayloadSize)), + }) + } + + _, err := cl.PutRecords(ctx, &kinesis.PutRecordsInput{ + StreamName: aws.String(kinesisStreamName), + Records: records, + }) + + return err +} + +// kinesisGetRecordsOp fetches a fresh shard iterator and reads from it, +// mirroring how a real consumer polls: two wire calls per iteration. +func kinesisGetRecordsOp(ctx context.Context, cl *kinesis.Client, shardID string) error { + iter, err := cl.GetShardIterator(ctx, &kinesis.GetShardIteratorInput{ + StreamName: aws.String(kinesisStreamName), + ShardId: aws.String(shardID), + ShardIteratorType: kinesistypes.ShardIteratorTypeLatest, + }) + if err != nil { + return fmt.Errorf("get shard iterator: %w", err) + } + + if _, getErr := cl.GetRecords(ctx, &kinesis.GetRecordsInput{ShardIterator: iter.ShardIterator}); getErr != nil { + return fmt.Errorf("get records: %w", getErr) + } + + return nil +} + +func kinesisDescribeStreamOp(ctx context.Context, cl *kinesis.Client) error { + _, err := cl.DescribeStreamSummary(ctx, &kinesis.DescribeStreamSummaryInput{ + StreamName: aws.String(kinesisStreamName), + }) + + return err +} diff --git a/cmd/pgoload/kms.go b/cmd/pgoload/kms.go new file mode 100644 index 0000000000..a3f42e591c --- /dev/null +++ b/cmd/pgoload/kms.go @@ -0,0 +1,78 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/kms" + kmstypes "github.com/aws/aws-sdk-go-v2/service/kms/types" +) + +// ensureKMSKey creates one KMS key for the whole run and returns its ID. +func ensureKMSKey(ctx context.Context, cl *kms.Client, log *slog.Logger) (string, error) { + out, err := cl.CreateKey(ctx, &kms.CreateKeyInput{Description: aws.String("pgoload key")}) + if err != nil { + return "", fmt.Errorf("create key: %w", err) + } + + log.InfoContext(ctx, "kms key ready", "keyId", aws.ToString(out.KeyMetadata.KeyId)) + + return aws.ToString(out.KeyMetadata.KeyId), nil +} + +// kmsWorker repeatedly runs a mix of KMS operations, staggered by workerID, +// until ctx is done. +func kmsWorker(ctx context.Context, cl *kms.Client, res *resources, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, _, i int) error { return kmsEncryptDecryptOp(ctx, cl, res.kmsKeyID, i) }, + func(ctx context.Context, _, _ int) error { return kmsGenerateDataKeyOp(ctx, cl, res.kmsKeyID) }, + func(ctx context.Context, _, _ int) error { return kmsDescribeKeyOp(ctx, cl, res.kmsKeyID) }, + func(ctx context.Context, _, _ int) error { return kmsListKeysOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "kms", log) +} + +// kmsEncryptDecryptOp encrypts a payload and immediately decrypts it back, +// exercising both wire paths without needing shared ciphertext state. +func kmsEncryptDecryptOp(ctx context.Context, cl *kms.Client, keyID string, i int) error { + enc, err := cl.Encrypt(ctx, &kms.EncryptInput{ + KeyId: aws.String(keyID), + Plaintext: fmt.Appendf(nil, "pgoload-plaintext-%d", i), + }) + if err != nil { + return fmt.Errorf("encrypt: %w", err) + } + + if _, decErr := cl.Decrypt(ctx, &kms.DecryptInput{ + KeyId: aws.String(keyID), + CiphertextBlob: enc.CiphertextBlob, + }); decErr != nil { + return fmt.Errorf("decrypt: %w", decErr) + } + + return nil +} + +func kmsGenerateDataKeyOp(ctx context.Context, cl *kms.Client, keyID string) error { + _, err := cl.GenerateDataKey(ctx, &kms.GenerateDataKeyInput{ + KeyId: aws.String(keyID), + KeySpec: kmstypes.DataKeySpecAes256, + }) + + return err +} + +func kmsDescribeKeyOp(ctx context.Context, cl *kms.Client, keyID string) error { + _, err := cl.DescribeKey(ctx, &kms.DescribeKeyInput{KeyId: aws.String(keyID)}) + + return err +} + +func kmsListKeysOp(ctx context.Context, cl *kms.Client) error { + _, err := cl.ListKeys(ctx, &kms.ListKeysInput{}) + + return err +} diff --git a/cmd/pgoload/lambda.go b/cmd/pgoload/lambda.go new file mode 100644 index 0000000000..f48746fc9c --- /dev/null +++ b/cmd/pgoload/lambda.go @@ -0,0 +1,116 @@ +package main + +import ( + "archive/zip" + "bytes" + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/lambda" + lambdatypes "github.com/aws/aws-sdk-go-v2/service/lambda/types" +) + +const ( + lambdaFunctionName = "pgoload-fn" + lambdaWaitTimeout = 30 * time.Second +) + +// ensureLambdaFunction creates lambdaFunctionName from an in-memory zip, +// tolerating one that already exists, and waits (best-effort, bounded) for +// it to leave the Pending state. Returns the function name. +func ensureLambdaFunction(ctx context.Context, cl *lambda.Client, roleArn string, log *slog.Logger) (string, error) { + _, err := cl.CreateFunction(ctx, &lambda.CreateFunctionInput{ + FunctionName: aws.String(lambdaFunctionName), + Runtime: lambdatypes.RuntimeNodejs20x, + Handler: aws.String("index.handler"), + Role: aws.String(roleArn), + PackageType: lambdatypes.PackageTypeZip, + Code: &lambdatypes.FunctionCode{ZipFile: lambdaDummyZip()}, + }) + if err != nil { + var exists *lambdatypes.ResourceConflictException + if !errors.As(err, &exists) { + return "", fmt.Errorf("create function %s: %w", lambdaFunctionName, err) + } + + log.InfoContext(ctx, "lambda function already exists", "function", lambdaFunctionName) + } + + waitCtx, cancel := context.WithTimeout(ctx, lambdaWaitTimeout) + defer cancel() + + for { + out, getErr := cl.GetFunction(waitCtx, &lambda.GetFunctionInput{FunctionName: aws.String(lambdaFunctionName)}) + if getErr == nil && out.Configuration != nil && out.Configuration.State != lambdatypes.StatePending { + return lambdaFunctionName, nil + } + + select { + case <-waitCtx.Done(): + log.WarnContext( + ctx, "gave up waiting for lambda function to leave Pending state", "function", lambdaFunctionName, + ) + + return lambdaFunctionName, nil + case <-time.After(pollInterval): + } + } +} + +func lambdaDummyZip() []byte { + buf := new(bytes.Buffer) + w := zip.NewWriter(buf) + + f, _ := w.Create("index.js") + _, _ = f.Write([]byte("exports.handler = async () => ({statusCode: 200, body: 'pgoload'});")) + _ = w.Close() + + return buf.Bytes() +} + +// lambdaWorker repeatedly runs a mix of Lambda operations, staggered by +// workerID, until ctx is done. +func lambdaWorker( + ctx context.Context, + cl *lambda.Client, + res *resources, + workerID int, + c *opCounter, + log *slog.Logger, +) { + ops := []opFunc{ + func(ctx context.Context, _, _ int) error { return lambdaInvokeOp(ctx, cl, res.functionName) }, + func(ctx context.Context, _, _ int) error { return lambdaInvokeOp(ctx, cl, res.functionName) }, + func(ctx context.Context, _, _ int) error { + return lambdaGetFunctionOp(ctx, cl, res.functionName) + }, + func(ctx context.Context, _, _ int) error { return lambdaListFunctionsOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "lambda", log) +} + +func lambdaInvokeOp(ctx context.Context, cl *lambda.Client, functionName string) error { + _, err := cl.Invoke(ctx, &lambda.InvokeInput{ + FunctionName: aws.String(functionName), + Payload: []byte(`{"pgoload":true}`), + }) + + return err +} + +func lambdaGetFunctionOp(ctx context.Context, cl *lambda.Client, functionName string) error { + _, err := cl.GetFunction(ctx, &lambda.GetFunctionInput{FunctionName: aws.String(functionName)}) + + return err +} + +func lambdaListFunctionsOp(ctx context.Context, cl *lambda.Client) error { + _, err := cl.ListFunctions(ctx, &lambda.ListFunctionsInput{}) + + return err +} diff --git a/cmd/pgoload/logs.go b/cmd/pgoload/logs.go new file mode 100644 index 0000000000..488e8bb3a1 --- /dev/null +++ b/cmd/pgoload/logs.go @@ -0,0 +1,113 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs" + cwltypes "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs/types" +) + +const cwlGroupName = "/pgoload/group" + +func cwlStreamName(workerID int) string { + return fmt.Sprintf("pgoload-stream-%d", workerID) +} + +// ensureLogGroup creates cwlGroupName, tolerating one that already exists. +// Log streams are created per-worker in logsWorker, once, before it enters +// its op loop. +func ensureLogGroup(ctx context.Context, cl *cloudwatchlogs.Client, log *slog.Logger) error { + _, err := cl.CreateLogGroup(ctx, &cloudwatchlogs.CreateLogGroupInput{LogGroupName: aws.String(cwlGroupName)}) + if err == nil { + return nil + } + + var exists *cwltypes.ResourceAlreadyExistsException + if errors.As(err, &exists) { + log.InfoContext(ctx, "cloudwatch logs group already exists", "group", cwlGroupName) + + return nil + } + + return fmt.Errorf("create log group %s: %w", cwlGroupName, err) +} + +// logsWorker creates its own log stream once, then repeatedly runs a mix of +// CloudWatch Logs operations, staggered by workerID, until ctx is done. +func logsWorker(ctx context.Context, cl *cloudwatchlogs.Client, workerID int, c *opCounter, log *slog.Logger) { + if err := ensureLogStream(ctx, cl, workerID); err != nil { + log.WarnContext(ctx, "cloudwatch logs stream setup failed (continuing)", "worker", workerID, "error", err) + } + + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { return logsPutLogEventsOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return logsPutLogEventsOp(ctx, cl, workerID, i) }, + func(ctx context.Context, _, _ int) error { return logsDescribeLogStreamsOp(ctx, cl) }, + func(ctx context.Context, workerID, _ int) error { return logsGetLogEventsOp(ctx, cl, workerID) }, + func(ctx context.Context, _, _ int) error { return logsFilterLogEventsOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "logs", log) +} + +func ensureLogStream(ctx context.Context, cl *cloudwatchlogs.Client, workerID int) error { + _, err := cl.CreateLogStream(ctx, &cloudwatchlogs.CreateLogStreamInput{ + LogGroupName: aws.String(cwlGroupName), + LogStreamName: aws.String(cwlStreamName(workerID)), + }) + if err == nil { + return nil + } + + var exists *cwltypes.ResourceAlreadyExistsException + if errors.As(err, &exists) { + return nil + } + + return err +} + +func logsPutLogEventsOp(ctx context.Context, cl *cloudwatchlogs.Client, workerID, i int) error { + _, err := cl.PutLogEvents(ctx, &cloudwatchlogs.PutLogEventsInput{ + LogGroupName: aws.String(cwlGroupName), + LogStreamName: aws.String(cwlStreamName(workerID)), + LogEvents: []cwltypes.InputLogEvent{ + { + Message: aws.String(fmt.Sprintf("pgoload log line worker=%d iter=%d", workerID, i)), + Timestamp: aws.Int64(time.Now().UnixMilli()), + }, + }, + }) + + return err +} + +func logsDescribeLogStreamsOp(ctx context.Context, cl *cloudwatchlogs.Client) error { + _, err := cl.DescribeLogStreams(ctx, &cloudwatchlogs.DescribeLogStreamsInput{ + LogGroupName: aws.String(cwlGroupName), + }) + + return err +} + +func logsGetLogEventsOp(ctx context.Context, cl *cloudwatchlogs.Client, workerID int) error { + _, err := cl.GetLogEvents(ctx, &cloudwatchlogs.GetLogEventsInput{ + LogGroupName: aws.String(cwlGroupName), + LogStreamName: aws.String(cwlStreamName(workerID)), + }) + + return err +} + +func logsFilterLogEventsOp(ctx context.Context, cl *cloudwatchlogs.Client) error { + _, err := cl.FilterLogEvents(ctx, &cloudwatchlogs.FilterLogEventsInput{ + LogGroupName: aws.String(cwlGroupName), + }) + + return err +} diff --git a/cmd/pgoload/main.go b/cmd/pgoload/main.go index acea7719e6..ddad36a86b 100644 --- a/cmd/pgoload/main.go +++ b/cmd/pgoload/main.go @@ -30,8 +30,6 @@ import ( "time" "github.com/aws/aws-sdk-go-v2/aws" - awscfg "github.com/aws/aws-sdk-go-v2/config" - "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/service/dynamodb" ddbtypes "github.com/aws/aws-sdk-go-v2/service/dynamodb/types" "github.com/aws/aws-sdk-go-v2/service/s3" @@ -125,11 +123,53 @@ type config struct { } // counters tracks operation and error totals across every worker goroutine. +// ddb/s3 keep their original flat fields untouched; every service added for +// breadth uses the shared opCounter type instead of four more fields each. type counters struct { ddbOps atomic.Int64 ddbErrors atomic.Int64 s3Ops atomic.Int64 s3Errors atomic.Int64 + + sqs opCounter + sns opCounter + kinesis opCounter + iam opCounter + sts opCounter + ssm opCounter + secretsmanager opCounter + cloudwatch opCounter + logs opCounter + ec2 opCounter + lambda opCounter + kms opCounter + eventbridge opCounter + stepfunctions opCounter +} + +// namedCounter pairs a scenario name with its counter, for summary printing. +type namedCounter struct { + c *opCounter + name string +} + +func (c *counters) breadthCounters() []namedCounter { + return []namedCounter{ + {name: "sqs", c: &c.sqs}, + {name: "sns", c: &c.sns}, + {name: "kinesis", c: &c.kinesis}, + {name: "iam", c: &c.iam}, + {name: "sts", c: &c.sts}, + {name: "ssm", c: &c.ssm}, + {name: "secretsmanager", c: &c.secretsmanager}, + {name: "cloudwatch", c: &c.cloudwatch}, + {name: "logs", c: &c.logs}, + {name: "ec2", c: &c.ec2}, + {name: "lambda", c: &c.lambda}, + {name: "kms", c: &c.kms}, + {name: "eventbridge", c: &c.eventbridge}, + {name: "stepfunctions", c: &c.stepfunctions}, + } } func main() { @@ -143,21 +183,21 @@ func run() int { rootCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - ddbClient, s3Client, err := buildClients(rootCtx, cfg) + cls, err := buildClients(rootCtx, cfg) if err != nil { log.ErrorContext(rootCtx, "failed to build aws clients", "error", err) return 1 } - buckets, err := setupResources(rootCtx, ddbClient, s3Client, log) + res, err := setupResources(rootCtx, cls, log) if err != nil { log.ErrorContext(rootCtx, "failed to prepare load resources", "error", err) return 1 } - return runLoad(rootCtx, cfg, ddbClient, s3Client, buckets, log) + return runLoad(rootCtx, cfg, cls, res, log) } // parseFlags resolves CLI flags, falling back to environment variables and @@ -212,63 +252,9 @@ func envIntOrDefault(key string, def int) int { return def } -// buildClients constructs the DynamoDB and S3 clients pointed at the -// Gopherstack endpoint, using static test credentials. -func buildClients(ctx context.Context, cfg config) (*dynamodb.Client, *s3.Client, error) { - awsCfg, err := awscfg.LoadDefaultConfig(ctx, - awscfg.WithRegion(awsRegion), - awscfg.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("test", "test", "")), - ) - if err != nil { - return nil, nil, fmt.Errorf("load aws config: %w", err) - } - - ddbClient := dynamodb.NewFromConfig(awsCfg, func(o *dynamodb.Options) { - o.BaseEndpoint = aws.String(cfg.endpoint) - }) - - s3Client := s3.NewFromConfig(awsCfg, func(o *s3.Options) { - o.BaseEndpoint = aws.String(cfg.endpoint) - o.UsePathStyle = true - }) - - return ddbClient, s3Client, nil -} - -// setupResources provisions the DynamoDB table and S3 buckets used by the -// load scenarios, bounded by setupTimeout so a short -duration still leaves -// time for the load phase itself. -func setupResources( - ctx context.Context, - ddbClient *dynamodb.Client, - s3Client *s3.Client, - log *slog.Logger, -) ([]string, error) { - setupCtx, cancel := context.WithTimeout(ctx, setupTimeout) - defer cancel() - - if err := ensureDDBTable(setupCtx, ddbClient, log); err != nil { - return nil, fmt.Errorf("ddb table setup: %w", err) - } - - buckets, err := ensureS3Buckets(setupCtx, s3Client, log) - if err != nil { - return nil, fmt.Errorf("s3 bucket setup: %w", err) - } - - return buckets, nil -} - // runLoad fans out cfg.concurrency workers for each scenario and lets them // run until cfg.duration elapses (or ctx is canceled), then prints a summary. -func runLoad( - ctx context.Context, - cfg config, - ddbClient *dynamodb.Client, - s3Client *s3.Client, - buckets []string, - log *slog.Logger, -) int { +func runLoad(ctx context.Context, cfg config, cls *clients, res *resources, log *slog.Logger) int { loadCtx, cancel := context.WithTimeout(ctx, cfg.duration) defer cancel() @@ -276,23 +262,37 @@ func runLoad( var wg sync.WaitGroup - for w := range cfg.concurrency { - wg.Add(1) + spawn := func(n int, worker func(id int)) { + for w := range n { + wg.Add(1) - go func(id int) { - defer wg.Done() - ddbWorker(loadCtx, ddbClient, id, c, log) - }(w) + go func(id int) { + defer wg.Done() + worker(id) + }(w) + } } - for w := range cfg.concurrency { - wg.Add(1) - - go func(id int) { - defer wg.Done() - s3Worker(loadCtx, s3Client, buckets, id, c, log) - }(w) - } + spawn(cfg.concurrency, func(id int) { ddbWorker(loadCtx, cls.ddb, id, c, log) }) + spawn(cfg.concurrency, func(id int) { s3Worker(loadCtx, cls.s3, res.s3Buckets, id, c, log) }) + spawn(cfg.concurrency, func(id int) { sqsWorker(loadCtx, cls.sqs, res, id, &c.sqs, log) }) + spawn(cfg.concurrency, func(id int) { snsWorker(loadCtx, cls.sns, res, id, &c.sns, log) }) + spawn(cfg.concurrency, func(id int) { kinesisWorker(loadCtx, cls.kinesis, res, id, &c.kinesis, log) }) + spawn(cfg.concurrency, func(id int) { iamWorker(loadCtx, cls.iam, id, &c.iam, log) }) + spawn(cfg.concurrency, func(id int) { stsWorker(loadCtx, cls.sts, res, id, &c.sts, log) }) + spawn(cfg.concurrency, func(id int) { ssmWorker(loadCtx, cls.ssm, id, &c.ssm, log) }) + spawn(cfg.concurrency, func(id int) { + secretsManagerWorker(loadCtx, cls.secretsmanager, id, &c.secretsmanager, log) + }) + spawn(cfg.concurrency, func(id int) { cloudwatchWorker(loadCtx, cls.cloudwatch, id, &c.cloudwatch, log) }) + spawn(cfg.concurrency, func(id int) { logsWorker(loadCtx, cls.logs, id, &c.logs, log) }) + spawn(cfg.concurrency, func(id int) { ec2Worker(loadCtx, cls.ec2, id, &c.ec2, log) }) + spawn(cfg.concurrency, func(id int) { lambdaWorker(loadCtx, cls.lambda, res, id, &c.lambda, log) }) + spawn(cfg.concurrency, func(id int) { kmsWorker(loadCtx, cls.kms, res, id, &c.kms, log) }) + spawn(cfg.concurrency, func(id int) { eventBridgeWorker(loadCtx, cls.eventbridge, id, &c.eventbridge, log) }) + spawn(cfg.concurrency, func(id int) { + stepFunctionsWorker(loadCtx, cls.stepfunctions, res, id, &c.stepfunctions, log) + }) wg.Wait() printSummary(c) @@ -308,6 +308,10 @@ func printSummary(c *counters) { c.s3Ops.Load(), c.ddbErrors.Load()+c.s3Errors.Load(), ) + + for _, nc := range c.breadthCounters() { + fmt.Fprintf(os.Stdout, "%s ops=%d, errors=%d\n", nc.name, nc.c.ops.Load(), nc.c.errors.Load()) + } } // --------------------------------------------------------------------------- diff --git a/cmd/pgoload/secretsmanager.go b/cmd/pgoload/secretsmanager.go new file mode 100644 index 0000000000..b304e72075 --- /dev/null +++ b/cmd/pgoload/secretsmanager.go @@ -0,0 +1,126 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/secretsmanager" + smtypes "github.com/aws/aws-sdk-go-v2/service/secretsmanager/types" +) + +// secretsCount is how many secrets pgoload provisions once during setup and +// then reads/updates/recreates in rotation, bounding the backend's secret +// set for the whole run. Deliberately not a multiple of len(ops) in +// secretsManagerWorker (5) for the same reason as ssmParamRotation: it +// keeps every secret index reachable by every operation instead of a fixed +// subset always landing on the same op (see ssm.go). +const secretsCount = 8 + +func secretName(idx int) string { + return fmt.Sprintf("pgoload-secret-%d", idx) +} + +// secretIndex spreads workers across secretsCount using their workerID as a +// phase offset. Indexing by i alone would put every worker's Nth iteration +// on the same secret at roughly the same wall-clock time (all i counters +// advance in lockstep), maximizing collisions between concurrent +// Get/Put/Recreate calls; the offset spreads that load instead. +func secretIndex(workerID, i int) int { + return (workerID + i) % secretsCount +} + +// ensureSecrets creates secretsCount secrets once, tolerating any that +// already exist from a prior run. +func ensureSecrets(ctx context.Context, cl *secretsmanager.Client, log *slog.Logger) error { + for idx := range secretsCount { + _, err := cl.CreateSecret(ctx, &secretsmanager.CreateSecretInput{ + Name: aws.String(secretName(idx)), + SecretString: aws.String("s3cr3t-0"), + }) + if err == nil { + continue + } + + var exists *smtypes.ResourceExistsException + if errors.As(err, &exists) { + continue + } + + return fmt.Errorf("create secret %s: %w", secretName(idx), err) + } + + log.InfoContext(ctx, "secretsmanager secrets ready", "count", secretsCount) + + return nil +} + +// secretsManagerWorker repeatedly runs a mix of Secrets Manager operations, +// staggered by workerID, until ctx is done. +func secretsManagerWorker( + ctx context.Context, + cl *secretsmanager.Client, + workerID int, + c *opCounter, + log *slog.Logger, +) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { return secretsGetSecretValueOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return secretsGetSecretValueOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return secretsGetSecretValueOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return secretsPutSecretValueOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return secretsPutSecretValueOp(ctx, cl, workerID, i) }, + func(ctx context.Context, _, _ int) error { return secretsListSecretsOp(ctx, cl) }, + func(ctx context.Context, workerID, i int) error { return secretsRecreateOp(ctx, cl, workerID, i) }, + } + + runOpLoop(ctx, workerID, ops, c, "secretsmanager", log) +} + +func secretsGetSecretValueOp(ctx context.Context, cl *secretsmanager.Client, workerID, i int) error { + _, err := cl.GetSecretValue(ctx, &secretsmanager.GetSecretValueInput{ + SecretId: aws.String(secretName(secretIndex(workerID, i))), + }) + + return err +} + +func secretsPutSecretValueOp(ctx context.Context, cl *secretsmanager.Client, workerID, i int) error { + _, err := cl.PutSecretValue(ctx, &secretsmanager.PutSecretValueInput{ + SecretId: aws.String(secretName(secretIndex(workerID, i))), + SecretString: aws.String(fmt.Sprintf("s3cr3t-%d", i)), + }) + + return err +} + +func secretsListSecretsOp(ctx context.Context, cl *secretsmanager.Client) error { + _, err := cl.ListSecrets(ctx, &secretsmanager.ListSecretsInput{}) + + return err +} + +// secretsRecreateOp deletes and immediately recreates one secret, exercising +// the full delete/create wire paths without ever colliding with an existing +// name (delete always runs first). +func secretsRecreateOp(ctx context.Context, cl *secretsmanager.Client, workerID, i int) error { + name := secretName(secretIndex(workerID, i)) + + if _, err := cl.DeleteSecret(ctx, &secretsmanager.DeleteSecretInput{ + SecretId: aws.String(name), + ForceDeleteWithoutRecovery: aws.Bool(true), + }); err != nil { + return fmt.Errorf("delete secret %s: %w", name, err) + } + + if _, err := cl.CreateSecret(ctx, &secretsmanager.CreateSecretInput{ + Name: aws.String(name), + SecretString: aws.String("s3cr3t-recreated"), + }); err != nil { + return fmt.Errorf("recreate secret %s: %w", name, err) + } + + return nil +} diff --git a/cmd/pgoload/sns.go b/cmd/pgoload/sns.go new file mode 100644 index 0000000000..2d3b1811d6 --- /dev/null +++ b/cmd/pgoload/sns.go @@ -0,0 +1,77 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/sns" +) + +const snsTopicName = "pgoload-topic" + +// ensureSNSTopic creates snsTopicName (CreateTopic is idempotent by name) +// and subscribes it to the SQS queue at queueArn for realistic fan-out. +// Subscribing is best-effort: a failure there doesn't fail setup, since +// Publish/ListSubscriptions/GetTopicAttributes still exercise SNS on their +// own. +func ensureSNSTopic(ctx context.Context, cl *sns.Client, queueArn string, log *slog.Logger) (string, error) { + out, err := cl.CreateTopic(ctx, &sns.CreateTopicInput{Name: aws.String(snsTopicName)}) + if err != nil { + return "", fmt.Errorf("create topic %s: %w", snsTopicName, err) + } + + topicArn := aws.ToString(out.TopicArn) + + if _, subErr := cl.Subscribe(ctx, &sns.SubscribeInput{ + TopicArn: aws.String(topicArn), + Protocol: aws.String("sqs"), + Endpoint: aws.String(queueArn), + }); subErr != nil { + log.WarnContext(ctx, "sns subscribe to sqs queue failed (continuing)", "error", subErr) + } + + return topicArn, nil +} + +// snsWorker repeatedly runs a mix of SNS operations, staggered by workerID, +// until ctx is done. +func snsWorker(ctx context.Context, cl *sns.Client, res *resources, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { + return snsPublishOp(ctx, cl, res.topicArn, workerID, i) + }, + func(ctx context.Context, workerID, i int) error { + return snsPublishOp(ctx, cl, res.topicArn, workerID, i) + }, + func(ctx context.Context, _, _ int) error { return snsListSubscriptionsOp(ctx, cl, res.topicArn) }, + func(ctx context.Context, _, _ int) error { + return snsGetTopicAttributesOp(ctx, cl, res.topicArn) + }, + } + + runOpLoop(ctx, workerID, ops, c, "sns", log) +} + +func snsPublishOp(ctx context.Context, cl *sns.Client, topicArn string, workerID, i int) error { + _, err := cl.Publish(ctx, &sns.PublishInput{ + TopicArn: aws.String(topicArn), + Message: aws.String(fmt.Sprintf("pgoload notification worker=%d iter=%d", workerID, i)), + Subject: aws.String("pgoload"), + }) + + return err +} + +func snsListSubscriptionsOp(ctx context.Context, cl *sns.Client, topicArn string) error { + _, err := cl.ListSubscriptionsByTopic(ctx, &sns.ListSubscriptionsByTopicInput{TopicArn: aws.String(topicArn)}) + + return err +} + +func snsGetTopicAttributesOp(ctx context.Context, cl *sns.Client, topicArn string) error { + _, err := cl.GetTopicAttributes(ctx, &sns.GetTopicAttributesInput{TopicArn: aws.String(topicArn)}) + + return err +} diff --git a/cmd/pgoload/sqs.go b/cmd/pgoload/sqs.go new file mode 100644 index 0000000000..1a8726c19a --- /dev/null +++ b/cmd/pgoload/sqs.go @@ -0,0 +1,135 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + "strings" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/sqs" + sqstypes "github.com/aws/aws-sdk-go-v2/service/sqs/types" +) + +const ( + sqsQueueName = "pgoload-queue" + + sqsBodySmallSize = 128 + sqsBodyMediumSize = 2048 + // sqsBodySizeMod: every sqsBodySizeMod'th message gets the larger body. + sqsBodySizeMod = 5 + + sqsBatchSize = 5 + sqsReceiveMaxMessages = 10 + sqsReceiveWaitSeconds = 1 +) + +// ensureSQSQueue creates sqsQueueName, tolerating one that already exists, +// and returns its URL and ARN. +func ensureSQSQueue(ctx context.Context, cl *sqs.Client, log *slog.Logger) (string, string, error) { + out, err := cl.CreateQueue(ctx, &sqs.CreateQueueInput{QueueName: aws.String(sqsQueueName)}) + if err != nil { + return "", "", fmt.Errorf("create queue %s: %w", sqsQueueName, err) + } + + log.InfoContext(ctx, "sqs queue ready", "queue", sqsQueueName) + + queueURL := aws.ToString(out.QueueUrl) + + attrs, err := cl.GetQueueAttributes(ctx, &sqs.GetQueueAttributesInput{ + QueueUrl: aws.String(queueURL), + AttributeNames: []sqstypes.QueueAttributeName{sqstypes.QueueAttributeNameQueueArn}, + }) + if err != nil { + return "", "", fmt.Errorf("get queue arn %s: %w", sqsQueueName, err) + } + + return queueURL, attrs.Attributes[string(sqstypes.QueueAttributeNameQueueArn)], nil +} + +func sqsBody(i int) string { + size := sqsBodySmallSize + if i%sqsBodySizeMod == 0 { + size = sqsBodyMediumSize + } + + return strings.Repeat("m", size) +} + +// sqsWorker repeatedly runs a mix of SQS operations, staggered by workerID, +// until ctx is done. +func sqsWorker(ctx context.Context, cl *sqs.Client, res *resources, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, _, i int) error { return sqsSendMessageOp(ctx, cl, res.queueURL, i) }, + func(ctx context.Context, _, i int) error { return sqsSendMessageOp(ctx, cl, res.queueURL, i) }, + func(ctx context.Context, _, i int) error { + return sqsSendMessageBatchOp(ctx, cl, res.queueURL, i) + }, + func(ctx context.Context, _, _ int) error { return sqsReceiveDeleteOp(ctx, cl, res.queueURL) }, + func(ctx context.Context, _, _ int) error { + return sqsGetQueueAttributesOp(ctx, cl, res.queueURL) + }, + } + + runOpLoop(ctx, workerID, ops, c, "sqs", log) +} + +func sqsSendMessageOp(ctx context.Context, cl *sqs.Client, queueURL string, i int) error { + _, err := cl.SendMessage(ctx, &sqs.SendMessageInput{ + QueueUrl: aws.String(queueURL), + MessageBody: aws.String(sqsBody(i)), + }) + + return err +} + +func sqsSendMessageBatchOp(ctx context.Context, cl *sqs.Client, queueURL string, base int) error { + entries := make([]sqstypes.SendMessageBatchRequestEntry, 0, sqsBatchSize) + for j := range sqsBatchSize { + entries = append(entries, sqstypes.SendMessageBatchRequestEntry{ + Id: aws.String(fmt.Sprintf("m%d", j)), + MessageBody: aws.String(sqsBody(base + j)), + }) + } + + _, err := cl.SendMessageBatch(ctx, &sqs.SendMessageBatchInput{ + QueueUrl: aws.String(queueURL), + Entries: entries, + }) + + return err +} + +// sqsReceiveDeleteOp receives a batch of messages and deletes each one, so +// the receive and delete wire paths are exercised together like a real +// consumer. +func sqsReceiveDeleteOp(ctx context.Context, cl *sqs.Client, queueURL string) error { + out, err := cl.ReceiveMessage(ctx, &sqs.ReceiveMessageInput{ + QueueUrl: aws.String(queueURL), + MaxNumberOfMessages: sqsReceiveMaxMessages, + WaitTimeSeconds: sqsReceiveWaitSeconds, + }) + if err != nil { + return err + } + + for _, msg := range out.Messages { + if _, delErr := cl.DeleteMessage(ctx, &sqs.DeleteMessageInput{ + QueueUrl: aws.String(queueURL), + ReceiptHandle: msg.ReceiptHandle, + }); delErr != nil { + return fmt.Errorf("delete message: %w", delErr) + } + } + + return nil +} + +func sqsGetQueueAttributesOp(ctx context.Context, cl *sqs.Client, queueURL string) error { + _, err := cl.GetQueueAttributes(ctx, &sqs.GetQueueAttributesInput{ + QueueUrl: aws.String(queueURL), + AttributeNames: []sqstypes.QueueAttributeName{sqstypes.QueueAttributeNameAll}, + }) + + return err +} diff --git a/cmd/pgoload/ssm.go b/cmd/pgoload/ssm.go new file mode 100644 index 0000000000..adc042d213 --- /dev/null +++ b/cmd/pgoload/ssm.go @@ -0,0 +1,71 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/ssm" + ssmtypes "github.com/aws/aws-sdk-go-v2/service/ssm/types" +) + +// ssmParamRotation bounds the number of distinct parameter names pgoload +// cycles through per worker. Deliberately not a multiple of len(ops) in +// ssmWorker (5): if it were, i%ssmParamRotation would determine i%len(ops) +// exactly, so every parameter name would always be reached by the same +// operation — e.g. always Get, never Put — producing permanent (not +// transient) ParameterNotFound errors instead of a realistic op mix. +const ssmParamRotation = 9 + +func ssmParamName(workerID, i int) string { + return fmt.Sprintf("/pgoload/%d/%d", workerID, i%ssmParamRotation) +} + +// ssmWorker repeatedly runs a mix of SSM Parameter Store operations, +// staggered by workerID, until ctx is done. +func ssmWorker(ctx context.Context, cl *ssm.Client, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { return ssmPutParameterOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return ssmPutParameterOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return ssmGetParameterOp(ctx, cl, workerID, i) }, + func(ctx context.Context, workerID, i int) error { return ssmGetParametersOp(ctx, cl, workerID, i) }, + func(ctx context.Context, _, _ int) error { return ssmDescribeParametersOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "ssm", log) +} + +func ssmPutParameterOp(ctx context.Context, cl *ssm.Client, workerID, i int) error { + _, err := cl.PutParameter(ctx, &ssm.PutParameterInput{ + Name: aws.String(ssmParamName(workerID, i)), + Value: aws.String(fmt.Sprintf("value-%d", i)), + Type: ssmtypes.ParameterTypeString, + Overwrite: aws.Bool(true), + }) + + return err +} + +func ssmGetParameterOp(ctx context.Context, cl *ssm.Client, workerID, i int) error { + _, err := cl.GetParameter(ctx, &ssm.GetParameterInput{Name: aws.String(ssmParamName(workerID, i))}) + + return err +} + +func ssmGetParametersOp(ctx context.Context, cl *ssm.Client, workerID, _ int) error { + names := make([]string, 0, ssmParamRotation) + for j := range ssmParamRotation { + names = append(names, ssmParamName(workerID, j)) + } + + _, err := cl.GetParameters(ctx, &ssm.GetParametersInput{Names: names}) + + return err +} + +func ssmDescribeParametersOp(ctx context.Context, cl *ssm.Client) error { + _, err := cl.DescribeParameters(ctx, &ssm.DescribeParametersInput{}) + + return err +} diff --git a/cmd/pgoload/stepfunctions.go b/cmd/pgoload/stepfunctions.go new file mode 100644 index 0000000000..028694eb57 --- /dev/null +++ b/cmd/pgoload/stepfunctions.go @@ -0,0 +1,124 @@ +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/sfn" + sfntypes "github.com/aws/aws-sdk-go-v2/service/sfn/types" +) + +const sfnStateMachineName = "pgoload-sm" + +const sfnPassState = "Pass" + +// errStateMachineNotFound is returned when a state machine that +// CreateStateMachine reported as already existing can't be located by a +// follow-up ListStateMachines call. +var errStateMachineNotFound = errors.New("state machine exists but was not found by ListStateMachines") + +// sfnDefinition is a minimal one-state Pass machine — enough to exercise +// CreateStateMachine/StartExecution/DescribeExecution without needing real +// task integrations. +func sfnDefinition() string { + def, _ := json.Marshal(map[string]any{ + "StartAt": sfnPassState, + "States": map[string]any{ + sfnPassState: map[string]any{"Type": "Pass", "End": true}, + }, + }) + + return string(def) +} + +// ensureStateMachine creates sfnStateMachineName, tolerating one that +// already exists, and returns its ARN. +func ensureStateMachine(ctx context.Context, cl *sfn.Client, roleArn string, log *slog.Logger) (string, error) { + out, err := cl.CreateStateMachine(ctx, &sfn.CreateStateMachineInput{ + Name: aws.String(sfnStateMachineName), + Definition: aws.String(sfnDefinition()), + RoleArn: aws.String(roleArn), + }) + if err == nil { + return aws.ToString(out.StateMachineArn), nil + } + + var exists *sfntypes.StateMachineAlreadyExists + if !errors.As(err, &exists) { + return "", fmt.Errorf("create state machine %s: %w", sfnStateMachineName, err) + } + + log.InfoContext(ctx, "step functions state machine already exists", "stateMachine", sfnStateMachineName) + + list, err := cl.ListStateMachines(ctx, &sfn.ListStateMachinesInput{}) + if err != nil { + return "", fmt.Errorf("list state machines: %w", err) + } + + for _, sm := range list.StateMachines { + if aws.ToString(sm.Name) == sfnStateMachineName { + return aws.ToString(sm.StateMachineArn), nil + } + } + + return "", fmt.Errorf("%w: stateMachine=%s", errStateMachineNotFound, sfnStateMachineName) +} + +// stepFunctionsWorker repeatedly runs a mix of Step Functions operations, +// staggered by workerID, until ctx is done. +func stepFunctionsWorker( + ctx context.Context, + cl *sfn.Client, + res *resources, + workerID int, + c *opCounter, + log *slog.Logger, +) { + ops := []opFunc{ + func(ctx context.Context, workerID, i int) error { + return sfnStartDescribeExecutionOp(ctx, cl, res.stateMachineArn, workerID, i) + }, + func(ctx context.Context, _, _ int) error { + return sfnListExecutionsOp(ctx, cl, res.stateMachineArn) + }, + func(ctx context.Context, _, _ int) error { return sfnListStateMachinesOp(ctx, cl) }, + } + + runOpLoop(ctx, workerID, ops, c, "stepfunctions", log) +} + +// sfnStartDescribeExecutionOp starts an execution and immediately describes +// it, exercising both wire paths without needing shared execution state. +func sfnStartDescribeExecutionOp(ctx context.Context, cl *sfn.Client, stateMachineArn string, workerID, i int) error { + started, err := cl.StartExecution(ctx, &sfn.StartExecutionInput{ + StateMachineArn: aws.String(stateMachineArn), + Name: aws.String(fmt.Sprintf("pgoload-exec-%d-%d", workerID, i)), + }) + if err != nil { + return fmt.Errorf("start execution: %w", err) + } + + if _, descErr := cl.DescribeExecution(ctx, &sfn.DescribeExecutionInput{ + ExecutionArn: started.ExecutionArn, + }); descErr != nil { + return fmt.Errorf("describe execution: %w", descErr) + } + + return nil +} + +func sfnListExecutionsOp(ctx context.Context, cl *sfn.Client, stateMachineArn string) error { + _, err := cl.ListExecutions(ctx, &sfn.ListExecutionsInput{StateMachineArn: aws.String(stateMachineArn)}) + + return err +} + +func sfnListStateMachinesOp(ctx context.Context, cl *sfn.Client) error { + _, err := cl.ListStateMachines(ctx, &sfn.ListStateMachinesInput{}) + + return err +} diff --git a/cmd/pgoload/sts.go b/cmd/pgoload/sts.go new file mode 100644 index 0000000000..632f37c3b9 --- /dev/null +++ b/cmd/pgoload/sts.go @@ -0,0 +1,45 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/sts" +) + +// stsWorker repeatedly runs a mix of STS operations, staggered by workerID, +// until ctx is done. +func stsWorker(ctx context.Context, cl *sts.Client, res *resources, workerID int, c *opCounter, log *slog.Logger) { + ops := []opFunc{ + func(ctx context.Context, _, _ int) error { return stsGetCallerIdentityOp(ctx, cl) }, + func(ctx context.Context, _, _ int) error { return stsGetSessionTokenOp(ctx, cl) }, + func(ctx context.Context, workerID, i int) error { + return stsAssumeRoleOp(ctx, cl, res.roleArn, workerID, i) + }, + } + + runOpLoop(ctx, workerID, ops, c, "sts", log) +} + +func stsGetCallerIdentityOp(ctx context.Context, cl *sts.Client) error { + _, err := cl.GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{}) + + return err +} + +func stsGetSessionTokenOp(ctx context.Context, cl *sts.Client) error { + _, err := cl.GetSessionToken(ctx, &sts.GetSessionTokenInput{}) + + return err +} + +func stsAssumeRoleOp(ctx context.Context, cl *sts.Client, roleArn string, workerID, i int) error { + _, err := cl.AssumeRole(ctx, &sts.AssumeRoleInput{ + RoleArn: aws.String(roleArn), + RoleSessionName: aws.String(fmt.Sprintf("pgoload-%d-%d", workerID, i)), + }) + + return err +} diff --git a/default.pgo b/default.pgo index b80feb3f4b..620c231907 100644 Binary files a/default.pgo and b/default.pgo differ diff --git a/scripts/pgo.sh b/scripts/pgo.sh index 8bf152b732..f78db00c08 100755 --- a/scripts/pgo.sh +++ b/scripts/pgo.sh @@ -8,13 +8,20 @@ # net/http/pprof on a side address, isolated from the main API port. # 3. Wait for the server's main port to accept connections. # 4. Start a pprof CPU capture in the background for PGO_SECONDS. -# 5. During the capture window, hammer the server with bin/pgoload (heavy -# DDB GSI/LSI + S3 traffic) and, best-effort, the integration test -# suite pointed at the running server (GOPHERSTACK_ENDPOINT) for -# broader coverage. The integration run is wrapped so it can never -# fail the pipeline. -# 6. Wait for the capture to finish, move it to default.pgo at repo root. -# 7. Validate: `go tool pprof -top` produces output and `go build +# 5. During the capture window, hammer the server with bin/pgoload (DDB, +# S3, and a broad set of other services) and, best-effort, the +# integration test suite pointed at the running server +# (GOPHERSTACK_ENDPOINT) for broader coverage. The integration run is +# wrapped so it can never fail the pipeline. +# 6. Best-effort: run every package with `go test`-style benchmarks +# (`func Benchmark...`), each with its own -cpuprofile. Most exercise +# in-memory backends directly with no HTTP involved, so this phase +# does not need the pprof capture window — it just adds more profiles +# to the same merge. Go PGO matches samples by function symbol, not by +# binary, so profiles from these separate test binaries still guide +# the same services/... functions the main binary calls. +# 7. Wait for the capture to finish, move it to default.pgo at repo root. +# 8. Validate: `go tool pprof -top` produces output and `go build # -pgo=auto` succeeds against the new profile. # # All knobs are overridable via environment variables. Safe to re-run. @@ -29,6 +36,11 @@ # PGO_SERVER_PORT main API port the server listens on (default 8000) # PGO_PPROF_ADDR pprof side address (default localhost:6060) # PGO_WAIT_TRIES bounded retries waiting for the server (default 30) +# PGO_BENCH_TIME per-benchmark-function -benchtime (default 300ms) +# PGO_BENCH_TIMEOUT timeout(1) budget per benchmark package, seconds +# (default 60) +# PGO_BENCH_PKGS space-separated package dirs to benchmark (default: +# auto-discovered via `grep -rl 'func Benchmark'`) set -euo pipefail @@ -41,6 +53,9 @@ PGO_INTEG_TIMEOUT="${PGO_INTEG_TIMEOUT:-120}" PGO_SERVER_PORT="${PGO_SERVER_PORT:-8000}" PGO_PPROF_ADDR="${PGO_PPROF_ADDR:-localhost:6060}" PGO_WAIT_TRIES="${PGO_WAIT_TRIES:-30}" +PGO_BENCH_TIME="${PGO_BENCH_TIME:-300ms}" +PGO_BENCH_TIMEOUT="${PGO_BENCH_TIMEOUT:-60}" +PGO_BENCH_PKGS="${PGO_BENCH_PKGS:-}" BIN_SERVER="bin/gopherstack" BIN_LOAD="bin/pgoload" @@ -58,7 +73,7 @@ cleanup() { kill "${SERVER_PID}" 2>/dev/null || true wait "${SERVER_PID}" 2>/dev/null || true fi - rm -f "${CPU_PROFILE}" cpu1.pprof cpu2.pprof + rm -f "${CPU_PROFILE}" cpu1.pprof cpu2.pprof bench_*.pprof bin/bench_*.test exit "${status}" } trap cleanup EXIT @@ -82,7 +97,7 @@ else fi log "starting server (GOPHERSTACK_PPROF_ADDR=${PGO_PPROF_ADDR}, port ${PGO_SERVER_PORT})..." -GOPHERSTACK_PPROF_ADDR="${PGO_PPROF_ADDR}" "./${BIN_SERVER}" & +GOPHERSTACK_PPROF_ADDR="${PGO_PPROF_ADDR}" PORT="${PGO_SERVER_PORT}" "./${BIN_SERVER}" & SERVER_PID=$! log "waiting for http://localhost:${PGO_SERVER_PORT} to accept connections..." @@ -141,6 +156,17 @@ if [[ "${INTEG_OK}" == "1" ]]; then CAP2_PID=$! GOPHERSTACK_ENDPOINT="http://localhost:${PGO_SERVER_PORT}" \ timeout "${PGO_INTEG_TIMEOUT}" "./${INTEG_BIN}" -test.count=1 >/dev/null 2>&1 || true + + # Some integration tests create Pipes/Scheduler resources targeting a + # queue or function that gets deleted before the pipe/schedule itself + # does, and their background pollers then retry the missing target at + # ~1Hz for as long as the server keeps running (observed: for the rest of + # this capture window and beyond). That is pure error-path retry noise, + # not representative traffic, and it lands squarely inside this capture + # window. Reset now, before any capture time is left, using the same + # /_gopherstack/reset endpoint the integration suite's own TestMain calls + # at startup — this only clears state, no functionality is disabled. + curl -s -o /dev/null -X POST "http://localhost:${PGO_SERVER_PORT}/_gopherstack/reset" || true wait "${CAP2_PID}" if [[ -s "${CPU2}" ]]; then PROFILES+=("${CPU2}") @@ -149,6 +175,51 @@ if [[ "${INTEG_OK}" == "1" ]]; then fi fi +# --- Phase 3: lots of bench calls. Every package with `func Benchmark...` +# tests gets its own `go test -bench` run with its own -cpuprofile file +# (running -cpuprofile across multiple packages in one `go test` invocation +# silently overwrites all but the last, so this must be per-package). Most +# of these benchmarks drive in-memory backends directly and never touch the +# live server, so they don't need to run inside a pprof capture window; they +# just add more profiles to the same merge below. The one package that does +# talk to the live server (test/integration, if it has benchmarks) gets +# GOPHERSTACK_ENDPOINT so its benchmark clients point at it instead of +# spinning up Docker. Best-effort throughout: a failing or hanging +# benchmark package is logged and skipped, never fails the pipeline. +if [[ -n "${PGO_BENCH_PKGS}" ]]; then + read -ra BENCH_PKGS <<<"${PGO_BENCH_PKGS}" +else + mapfile -t BENCH_PKGS < <(grep -rl 'func Benchmark' --include='*_test.go' . | xargs -n1 dirname | sort -u) +fi + +log "phase 3: benchmarking ${#BENCH_PKGS[@]} package(s) (benchtime ${PGO_BENCH_TIME}, timeout ${PGO_BENCH_TIMEOUT}s each)..." +for pkg in "${BENCH_PKGS[@]}"; do + safe_name="$(echo "${pkg}" | tr -c 'A-Za-z0-9' '_')" + bench_profile="bench_${safe_name}.pprof" + # -cpuprofile implies -c (the compiled test binary is kept so `go tool + # pprof` can symbolize it later), and go test drops that binary in the + # current directory unless told otherwise. -o sends it into bin/, which + # is already gitignored, and it's deleted right after regardless so + # repeated runs don't pile up large binaries. + bench_bin="bin/bench_${safe_name}.test" + rm -f "${bench_profile}" "${bench_bin}" + + log " benchmarking ${pkg}..." + if ! GOPHERSTACK_ENDPOINT="http://localhost:${PGO_SERVER_PORT}" \ + timeout "${PGO_BENCH_TIMEOUT}" go test "./${pkg}" \ + -run '^$' -bench=. -benchtime="${PGO_BENCH_TIME}" \ + -cpuprofile="${bench_profile}" -o "${bench_bin}" >/dev/null 2>&1; then + log " ${pkg} benchmarks failed or timed out (continuing)" + fi + rm -f "${bench_bin}" + + if [[ -s "${bench_profile}" ]]; then + PROFILES+=("${bench_profile}") + else + rm -f "${bench_profile}" + fi +done + if [[ "${#PROFILES[@]}" -eq 0 ]]; then log "no non-empty captures were produced" exit 1 @@ -167,7 +238,7 @@ if [[ "${#PROFILES[@]}" -gt 1 ]]; then else cp "${PROFILES[0]}" "${PROFILE_OUT}" fi -rm -f "${CPU1}" "${CPU2}" +rm -f "${CPU1}" "${CPU2}" bench_*.pprof log "wrote ${PROFILE_OUT}" log "validating profile with go tool pprof..."