Move v4 TMDS container and task stats endpoint handlers to ecs-agent module

agent/go.sum

@@ -491,8 +491,8 @@ gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gotest.tools/v3 v3.0.3 h1:4AuOwCGf4lLR9u3YOe2awrHygurzhO/HeQ6laiA6Sx0=
 gotest.tools/v3 v3.0.3/go.mod h1:Z7Lb0S5l+klDB31fvDQX8ss/FlKDxtlFlw3Oa8Ymbl8=
+gotest.tools/v3 v3.5.0 h1:Ljk6PdHdOhAb5aDMWXjDLMMhph+BpztA4v1QdqEW2eY=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=

agent/handlers/task_server_setup.go

@@ -77,7 +77,7 @@ func taskServerSetup(
 	muxRouter.HandleFunc(tmdsv1.CredentialsPath,
 		tmdsv1.CredentialsHandler(credentialsManager, auditLogger))
 
-	tmdsAgentState := v4.NewTMDSAgentState(state, ecsClient, cluster, availabilityZone, vpcID, containerInstanceArn)
+	tmdsAgentState := v4.NewTMDSAgentState(state, statsEngine, ecsClient, cluster, availabilityZone, vpcID, containerInstanceArn)
 	metricsFactory := metrics.NewNopEntryFactory()
 
 	v2HandlersSetup(muxRouter, state, ecsClient, statsEngine, cluster, credentialsManager, auditLogger, availabilityZone, containerInstanceArn)
@@ -153,8 +153,8 @@ func v4HandlersSetup(muxRouter *mux.Router,
 	muxRouter.HandleFunc(tmdsv4.ContainerMetadataPath(), tmdsv4.ContainerMetadataHandler(tmdsAgentState, metricsFactory))
 	muxRouter.HandleFunc(tmdsv4.TaskMetadataPath(), tmdsv4.TaskMetadataHandler(tmdsAgentState, metricsFactory))
 	muxRouter.HandleFunc(tmdsv4.TaskMetadataWithTagsPath(), tmdsv4.TaskMetadataWithTagsHandler(tmdsAgentState, metricsFactory))
-	muxRouter.HandleFunc(v4.ContainerStatsPath, v4.ContainerStatsHandler(state, statsEngine))
-	muxRouter.HandleFunc(v4.TaskStatsPath, v4.TaskStatsHandler(state, statsEngine))
+	muxRouter.HandleFunc(tmdsv4.ContainerStatsPath(), tmdsv4.ContainerStatsHandler(tmdsAgentState, metricsFactory))
+	muxRouter.HandleFunc(tmdsv4.TaskStatsPath(), tmdsv4.TaskStatsHandler(tmdsAgentState, metricsFactory))
 	muxRouter.HandleFunc(v4.ContainerAssociationsPath, v4.ContainerAssociationsHandler(state))
 	muxRouter.HandleFunc(v4.ContainerAssociationPathWithSlash, v4.ContainerAssociationHandler(state))
 	muxRouter.HandleFunc(v4.ContainerAssociationPath, v4.ContainerAssociationHandler(state))

agent/handlers/v4/stats_response.go

@@ -25,7 +25,7 @@ import (
 // NewV4TaskStatsResponse returns a new v4 task stats response object
 func NewV4TaskStatsResponse(taskARN string,
 	state dockerstate.TaskEngineState,
-	statsEngine stats.Engine) (map[string]response.StatsResponse, error) {
+	statsEngine stats.Engine) (map[string]*response.StatsResponse, error) {
 
 	containerMap, ok := state.ContainerMapByArn(taskARN)
 	if !ok {
@@ -34,14 +34,14 @@ func NewV4TaskStatsResponse(taskARN string,
 			taskARN)
 	}
 
-	resp := make(map[string]response.StatsResponse)
+	resp := make(map[string]*response.StatsResponse)
 	for _, dockerContainer := range containerMap {
 		containerID := dockerContainer.DockerID
 		dockerStats, network_rate_stats, err := statsEngine.ContainerDockerStats(taskARN, containerID)
 		if err != nil {
 			seelog.Warnf("V4 task stats response: Unable to get stats for container '%s' for task '%s': %v",
 				containerID, taskARN, err)
-			resp[containerID] = response.StatsResponse{}
+			resp[containerID] = &response.StatsResponse{}
 			continue
 		}
 
@@ -50,7 +50,7 @@ func NewV4TaskStatsResponse(taskARN string,
 			Network_rate_stats: network_rate_stats,
 		}
 
-		resp[containerID] = statsResponse
+		resp[containerID] = &statsResponse
 	}
 
 	return resp, nil

agent/handlers/v4/tmdsstate.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/aws/amazon-ecs-agent/agent/api"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
+	"github.com/aws/amazon-ecs-agent/agent/stats"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/logger/field"
 	tmdsv4 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state"
@@ -25,6 +26,7 @@ import (
 // Implements AgentState interface for TMDS v4.
 type TMDSAgentState struct {
 	state                dockerstate.TaskEngineState
+	statsEngine          stats.Engine
 	ecsClient            api.ECSClient
 	cluster              string
 	availabilityZone     string
@@ -34,6 +36,7 @@ type TMDSAgentState struct {
 
 func NewTMDSAgentState(
 	state dockerstate.TaskEngineState,
+	statsEngine stats.Engine,
 	ecsClient api.ECSClient,
 	cluster string,
 	availabilityZone string,
@@ -42,6 +45,7 @@ func NewTMDSAgentState(
 ) *TMDSAgentState {
 	return &TMDSAgentState{
 		state:                state,
+		statsEngine:          statsEngine,
 		ecsClient:            ecsClient,
 		cluster:              cluster,
 		availabilityZone:     availabilityZone,
@@ -149,3 +153,49 @@ func (s *TMDSAgentState) getTaskMetadata(v3EndpointID string, includeTags bool)
 
 	return *taskResponse, nil
 }
+
+func (s *TMDSAgentState) GetContainerStats(v3EndpointID string) (tmdsv4.StatsResponse, error) {
+	taskARN, ok := s.state.TaskARNByV3EndpointID(v3EndpointID)
+	if !ok {
+		return tmdsv4.StatsResponse{}, tmdsv4.NewErrorStatsLookupFailure(fmt.Sprintf(
+			"V4 container handler: unable to get task arn from request: unable to get task Arn from v3 endpoint ID: %s",
+			v3EndpointID))
+	}
+
+	containerID, ok := s.state.DockerIDByV3EndpointID(v3EndpointID)
+	if !ok {
+		return tmdsv4.StatsResponse{}, tmdsv4.NewErrorStatsLookupFailure(fmt.Sprintf(
+			"V4 container stats handler: unable to get container ID from request: unable to get docker ID from v3 endpoint ID: %s",
+			v3EndpointID))
+	}
+
+	dockerStats, network_rate_stats, err := s.statsEngine.ContainerDockerStats(taskARN, containerID)
+	if err != nil {
+		return tmdsv4.StatsResponse{}, tmdsv4.NewErrorStatsFetchFailure(
+			fmt.Sprintf("Unable to get container stats for: %s", containerID),
+			err)
+	}
+
+	return tmdsv4.StatsResponse{
+		StatsJSON:          dockerStats,
+		Network_rate_stats: network_rate_stats,
+	}, nil
+}
+
+func (s *TMDSAgentState) GetTaskStats(v3EndpointID string) (map[string]*tmdsv4.StatsResponse, error) {
+	taskARN, ok := s.state.TaskARNByV3EndpointID(v3EndpointID)
+	if !ok {
+		return nil, tmdsv4.NewErrorStatsLookupFailure(fmt.Sprintf(
+			"V4 task stats handler: unable to get task arn from request: unable to get task Arn from v3 endpoint ID: %s",
+			v3EndpointID))
+	}
+
+	taskStatsResponse, err := NewV4TaskStatsResponse(taskARN, s.state, s.statsEngine)
+	if err != nil {
+		return nil, tmdsv4.NewErrorStatsFetchFailure(
+			fmt.Sprintf("Unable to get task stats for: %s", taskARN),
+			err)
+	}
+
+	return taskStatsResponse, nil
+}