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: