From 2afe07ab2d4c0377ccd3b5d690b872f6faac1090 Mon Sep 17 00:00:00 2001 From: wlggraham Date: Mon, 13 Jul 2026 12:39:18 -0600 Subject: [PATCH] fix(ecs): share single AWS client across all ECS plugin executions Previously, each ECS job execution called config.LoadDefaultConfig and constructed new ecs.Client and cloudwatchlogs.Client instances in buildExecutionContext and PullLogs. Under high concurrency (e.g. many parallel amazon_serp_backfill_replay shards), this caused N independent HTTP connection pools and TLS sessions to accumulate in the heimdall process, contributing to OOM kills. Clients are now initialized once in New() on the commandContext and reused across all Execute() calls. The AWS SDK v2 clients are goroutine-safe so no locking is required. Co-Authored-By: Claude Sonnet 4.6 --- internal/pkg/aws/cloudwatch.go | 12 +----------- internal/pkg/object/command/ecs/ecs.go | 22 ++++++++++++++-------- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/internal/pkg/aws/cloudwatch.go b/internal/pkg/aws/cloudwatch.go index 543c9f9d..8140670e 100644 --- a/internal/pkg/aws/cloudwatch.go +++ b/internal/pkg/aws/cloudwatch.go @@ -7,20 +7,10 @@ import ( "time" "github.com/aws/aws-sdk-go-v2/aws" - "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs" ) -func PullLogs(ctx context.Context, writer *os.File, logGroup, logStream string, chunkSize int, memoryLimit int64) error { - - // initialize AWS session - cfg, err := config.LoadDefaultConfig(ctx) - if err != nil { - return err - } - - // create a new cloudwatch logs client - client := cloudwatchlogs.NewFromConfig(cfg) +func PullLogs(ctx context.Context, client *cloudwatchlogs.Client, writer *os.File, logGroup, logStream string, chunkSize int, memoryLimit int64) error { // Write log stream header to differentiate streams header := fmt.Sprintf("=== %s ===\n\n", logStream) diff --git a/internal/pkg/object/command/ecs/ecs.go b/internal/pkg/object/command/ecs/ecs.go index 685a7810..66244b5a 100644 --- a/internal/pkg/object/command/ecs/ecs.go +++ b/internal/pkg/object/command/ecs/ecs.go @@ -40,6 +40,9 @@ type commandContext struct { Timeout duration.Duration `yaml:"timeout,omitempty" json:"timeout,omitempty"` MaxFailCount int `yaml:"max_fail_count,omitempty" json:"max_fail_count,omitempty"` // max failures before giving up Tags map[string]string `yaml:"tags,omitempty" json:"tags,omitempty"` + + ecsClient *ecs.Client + logsClient *cloudwatchlogs.Client } // ECS cluster context structure @@ -155,6 +158,13 @@ func New(commandCtx *heimdallContext.Context) (plugin.Handler, error) { } } + cfg, err := config.LoadDefaultConfig(context.Background()) + if err != nil { + return nil, err + } + e.ecsClient = ecs.NewFromConfig(cfg) + e.logsClient = cloudwatchlogs.NewFromConfig(cfg) + return e, nil } @@ -450,7 +460,7 @@ func (execCtx *executionContext) pollForCompletion(ctx context.Context) error { } -func buildExecutionContext(ctx context.Context, commandCtx *commandContext, j *job.Job, c *cluster.Cluster, runtime *plugin.Runtime) (*executionContext, error) { +func buildExecutionContext(_ context.Context, commandCtx *commandContext, j *job.Job, c *cluster.Cluster, runtime *plugin.Runtime) (*executionContext, error) { execCtx := &executionContext{ tasks: make(map[string]*taskTracker), @@ -501,12 +511,8 @@ func buildExecutionContext(ctx context.Context, commandCtx *commandContext, j *j return nil, err } - // initialize AWS session - cfg, err := config.LoadDefaultConfig(ctx) - if err != nil { - return nil, err - } - execCtx.ecsClient = ecs.NewFromConfig(cfg) + execCtx.ecsClient = commandCtx.ecsClient + execCtx.logsClient = commandCtx.logsClient return execCtx, nil @@ -754,7 +760,7 @@ func (execCtx *executionContext) retrieveLogs(ctx context.Context) error { case types.LogDriverAwslogs: logGroup := logInfo.options["awslogs-group"] logStream := fmt.Sprintf("%s/%s/%s", logInfo.options["awslogs-stream-prefix"], logInfo.containerName, taskID) - if err := heimdallAws.PullLogs(ctx, writer, logGroup, logStream, maxLogChunkSize, maxLogMemoryBytes); err != nil { + if err := heimdallAws.PullLogs(ctx, execCtx.logsClient, writer, logGroup, logStream, maxLogChunkSize, maxLogMemoryBytes); err != nil { return err } default: