diff --git a/agent/stats/engine.go b/agent/stats/engine.go index e09d3120ca8..214ebce382f 100644 --- a/agent/stats/engine.go +++ b/agent/stats/engine.go @@ -957,11 +957,10 @@ func (engine *DockerStatsEngine) doRemoveContainerUnsafe(container *StatsContain // Delete will do nothing if the specified key doesn't exist. delete(engine.tasksToDefinitions, taskArn) seelog.Debugf("Deleted task from tasks, arn: %s", taskArn) - } - if _, ok := engine.taskToServiceConnectStats[taskArn]; ok { + // Stop collecting service connect stats for this task delete(engine.taskToServiceConnectStats, taskArn) - seelog.Debugf("Deleted task from service connect stats watch list, arn: %s", taskArn) + seelog.Debugf("Deleted task from the Service Connect stats watchlist, task ARN: %s", taskArn) } // Remove the container from health container watch list diff --git a/agent/stats/engine_test.go b/agent/stats/engine_test.go index 96f1a687d95..869ad9fe58c 100644 --- a/agent/stats/engine_test.go +++ b/agent/stats/engine_test.go @@ -218,6 +218,46 @@ func TestStatsEngineAddRemoveContainers(t *testing.T) { validateIdleContainerMetrics(t, engine) } +func TestDoRemoveContainerUnsafeServiceConnectStats(t *testing.T) { + engine := NewDockerStatsEngine(&cfg, nil, eventStream("TestDoRemoveContainerUnsafeServiceConnectStats"), nil, nil, nil) + defer engine.removeAll() + + taskArn := "arn:aws:ecs:us-west-2:123456789012:task/test-task" + + // Setup task with service connect stats + engine.taskToServiceConnectStats[taskArn] = &ServiceConnectStats{} + engine.tasksToContainers[taskArn] = make(map[string]*StatsContainer) + + // Add two containers to the task with proper initialization + _, cancel1 := context.WithCancel(context.Background()) + container1 := &StatsContainer{ + containerMetadata: &ContainerMetadata{DockerID: "container1"}, + cancel: cancel1, + } + _, cancel2 := context.WithCancel(context.Background()) + container2 := &StatsContainer{ + containerMetadata: &ContainerMetadata{DockerID: "container2"}, + cancel: cancel2, + } + + engine.tasksToContainers[taskArn]["container1"] = container1 + engine.tasksToContainers[taskArn]["container2"] = container2 + + // Remove first container - service connect stats should remain + engine.doRemoveContainerUnsafe(container1, taskArn) + + // Verify service connect stats still exist (task has remaining containers) + _, exists := engine.taskToServiceConnectStats[taskArn] + assert.True(t, exists, "Service connect stats should remain when task has running containers") + + // Remove second container - service connect stats should be cleaned up + engine.doRemoveContainerUnsafe(container2, taskArn) + + // Verify service connect stats are cleaned up (no containers left in task) + _, exists = engine.taskToServiceConnectStats[taskArn] + assert.False(t, exists, "Service connect stats should be cleaned up when no containers remain in task") +} + func TestStatsEngineMetadataInStatsSets(t *testing.T) { mockCtrl := gomock.NewController(t) defer mockCtrl.Finish()