Move TMDS v4 container stats types to ecs-agent module (#3785)

aws · Jul 7, 2023 · e9ce7c7 · e9ce7c7
1 parent 235ac90
commit e9ce7c7
Show file tree

Hide file tree

Showing 80 changed files with 6,537 additions and 41 deletions.
diff --git a/agent/handlers/task_server_setup_test.go b/agent/handlers/task_server_setup_test.go
@@ -37,15 +37,14 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/config"
 	mock_dockerstate "github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
-	v4stats "github.com/aws/amazon-ecs-agent/agent/handlers/v4"
-	"github.com/aws/amazon-ecs-agent/agent/stats"
 	mock_stats "github.com/aws/amazon-ecs-agent/agent/stats/mock"
 	apieni "github.com/aws/amazon-ecs-agent/ecs-agent/api/eni"
 	mock_taskprotection "github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	mock_credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/ecs_client/model/ecs"
 	mock_audit "github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit/mocks"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/stats"
 	tmdsresponse "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/response"
 	tp "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers"
 	tptypes "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/types"
@@ -1286,9 +1285,9 @@ type TMDSResponse interface {
 		v4.TaskResponse |
 		tptypes.TaskProtectionResponse |
 		types.StatsJSON |
-		v4stats.StatsResponse |
+		v4.StatsResponse |
 		map[string]*types.StatsJSON |
-		map[string]*v4stats.StatsResponse |
+		map[string]*v4.StatsResponse |
 		string
 }
 
@@ -2861,7 +2860,7 @@ func TestV4ContainerStats(t *testing.T) {
 			RxBytesPerSecond: 52,
 			TxBytesPerSecond: 84,
 		}
-		testTMDSRequest(t, TMDSTestCase[v4stats.StatsResponse]{
+		testTMDSRequest(t, TMDSTestCase[v4.StatsResponse]{
 			path: path,
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				gomock.InOrder(
@@ -2874,7 +2873,7 @@ func TestV4ContainerStats(t *testing.T) {
 					Return(&dockerStats, &networkStats, nil)
 			},
 			expectedStatusCode: http.StatusOK,
-			expectedResponseBody: v4stats.StatsResponse{
+			expectedResponseBody: v4.StatsResponse{
 				StatsJSON:          &dockerStats,
 				Network_rate_stats: &networkStats,
 			},
@@ -2913,7 +2912,7 @@ func TestV4TaskStats(t *testing.T) {
 	})
 	t.Run("containerMap empty", func(t *testing.T) {
 		containerMap := map[string]*apicontainer.DockerContainer{}
-		testTMDSRequest(t, TMDSTestCase[map[string]*v4stats.StatsResponse]{
+		testTMDSRequest(t, TMDSTestCase[map[string]*v4.StatsResponse]{
 			path: path,
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				gomock.InOrder(
@@ -2922,14 +2921,14 @@ func TestV4TaskStats(t *testing.T) {
 				)
 			},
 			expectedStatusCode:   http.StatusOK,
-			expectedResponseBody: map[string]*v4stats.StatsResponse{},
+			expectedResponseBody: map[string]*v4.StatsResponse{},
 		})
 	})
 	t.Run("stats not found for a container", func(t *testing.T) {
 		containerMap := map[string]*apicontainer.DockerContainer{
 			containerName: {DockerID: containerID},
 		}
-		testTMDSRequest(t, TMDSTestCase[map[string]*v4stats.StatsResponse]{
+		testTMDSRequest(t, TMDSTestCase[map[string]*v4.StatsResponse]{
 			path: path,
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				gomock.InOrder(
@@ -2942,7 +2941,7 @@ func TestV4TaskStats(t *testing.T) {
 					Return(nil, nil, errors.New("some error"))
 			},
 			expectedStatusCode: http.StatusOK,
-			expectedResponseBody: map[string]*v4stats.StatsResponse{
+			expectedResponseBody: map[string]*v4.StatsResponse{
 				containerID: {
 					StatsJSON: nil, Network_rate_stats: nil,
 				}},
@@ -2957,7 +2956,7 @@ func TestV4TaskStats(t *testing.T) {
 			TxBytesPerSecond: 84,
 		}
 		dockerStats := types.StatsJSON{Stats: types.Stats{NumProcs: 2}}
-		testTMDSRequest(t, TMDSTestCase[map[string]*v4stats.StatsResponse]{
+		testTMDSRequest(t, TMDSTestCase[map[string]*v4.StatsResponse]{
 			path: path,
 			setStateExpectations: func(state *mock_dockerstate.MockTaskEngineState) {
 				gomock.InOrder(
@@ -2970,7 +2969,7 @@ func TestV4TaskStats(t *testing.T) {
 					Return(&dockerStats, &networkStats, nil)
 			},
 			expectedStatusCode: http.StatusOK,
-			expectedResponseBody: map[string]*v4stats.StatsResponse{containerID: {
+			expectedResponseBody: map[string]*v4.StatsResponse{containerID: {
 				StatsJSON:          &dockerStats,
 				Network_rate_stats: &networkStats,
 			}},

diff --git a/agent/handlers/v2/stats_response_test.go b/agent/handlers/v2/stats_response_test.go
@@ -19,7 +19,7 @@ package v2
 import (
 	"testing"
 
-	"github.com/aws/amazon-ecs-agent/agent/stats"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/stats"
 
 	apicontainer "github.com/aws/amazon-ecs-agent/agent/api/container"
 	mock_dockerstate "github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"

diff --git a/agent/handlers/v4/container_stats_handler.go b/agent/handlers/v4/container_stats_handler.go
@@ -10,7 +10,6 @@
 // on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
-
 package v4
 
 import (
@@ -22,6 +21,8 @@ import (
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
 	"github.com/aws/amazon-ecs-agent/agent/stats"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/utils"
+	response "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state"
+
 	"github.com/cihub/seelog"
 )
 
@@ -72,7 +73,7 @@ func WriteV4ContainerStatsResponse(w http.ResponseWriter,
 		return
 	}
 
-	containerStatsResponse := StatsResponse{
+	containerStatsResponse := response.StatsResponse{
 		StatsJSON:          dockerStats,
 		Network_rate_stats: network_rate_stats,
 	}

diff --git a/agent/handlers/v4/stats_response.go b/agent/handlers/v4/stats_response.go
@@ -16,22 +16,16 @@ package v4
 import (
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
 	"github.com/aws/amazon-ecs-agent/agent/stats"
+	response "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state"
+
 	"github.com/cihub/seelog"
-	"github.com/docker/docker/api/types"
 	"github.com/pkg/errors"
 )
 
-// StatsResponse is the v4 Stats response. It augments the v4 Stats response
-// with the docker stats.
-type StatsResponse struct {
-	*types.StatsJSON
-	Network_rate_stats *stats.NetworkStatsPerSec `json:"network_rate_stats,omitempty"`
-}
-
 // NewV4TaskStatsResponse returns a new v4 task stats response object
 func NewV4TaskStatsResponse(taskARN string,
 	state dockerstate.TaskEngineState,
-	statsEngine stats.Engine) (map[string]StatsResponse, error) {
+	statsEngine stats.Engine) (map[string]response.StatsResponse, error) {
 
 	containerMap, ok := state.ContainerMapByArn(taskARN)
 	if !ok {
@@ -40,18 +34,18 @@ func NewV4TaskStatsResponse(taskARN string,
 			taskARN)
 	}
 
-	resp := make(map[string]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] = StatsResponse{}
+			resp[containerID] = response.StatsResponse{}
 			continue
 		}
 
-		statsResponse := StatsResponse{
+		statsResponse := response.StatsResponse{
 			StatsJSON:          dockerStats,
 			Network_rate_stats: network_rate_stats,
 		}

diff --git a/agent/stats/engine.go b/agent/stats/engine.go
@@ -38,6 +38,7 @@ import (
 	ecsengine "github.com/aws/amazon-ecs-agent/agent/engine"
 	"github.com/aws/amazon-ecs-agent/agent/stats/resolver"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/eventstream"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/stats"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/docker/docker/api/types"
@@ -78,7 +79,7 @@ type DockerContainerMetadataResolver struct {
 // defined to make testing easier.
 type Engine interface {
 	GetInstanceMetrics(includeServiceConnectStats bool) (*ecstcs.MetricsMetadata, []*ecstcs.TaskMetric, error)
-	ContainerDockerStats(taskARN string, containerID string) (*types.StatsJSON, *NetworkStatsPerSec, error)
+	ContainerDockerStats(taskARN string, containerID string) (*types.StatsJSON, *stats.NetworkStatsPerSec, error)
 	GetTaskHealthMetrics() (*ecstcs.HealthMetadata, []*ecstcs.TaskHealth, error)
 	GetPublishServiceConnectTickerInterval() int32
 	SetPublishServiceConnectTickerInterval(int32)
@@ -957,7 +958,7 @@ func (engine *DockerStatsEngine) resetStatsUnsafe() {
 }
 
 // ContainerDockerStats returns the last stored raw docker stats object for a container
-func (engine *DockerStatsEngine) ContainerDockerStats(taskARN string, containerID string) (*types.StatsJSON, *NetworkStatsPerSec, error) {
+func (engine *DockerStatsEngine) ContainerDockerStats(taskARN string, containerID string) (*types.StatsJSON, *stats.NetworkStatsPerSec, error) {
 	engine.lock.RLock()
 	defer engine.lock.RUnlock()
 

diff --git a/agent/stats/engine_unix_integ_test.go b/agent/stats/engine_unix_integ_test.go
@@ -46,7 +46,7 @@ const (
 )
 
 const (
-	stats = `# TYPE MetricFamily3 histogram
+	testStats = `# TYPE MetricFamily3 histogram
 		MetricFamily3{dimensionX="value1", dimensionY="value2", le="0.5"} 1
 		MetricFamily3{dimensionX="value1", dimensionY="value2", le="1"} 2
 		MetricFamily3{dimensionX="value1", dimensionY="value2", le="5"} 3
@@ -144,7 +144,7 @@ func TestStatsEngineWithServiceConnectMetrics(t *testing.T) {
 			// simulate appnet server providing service connect metrics
 			r := mux.NewRouter()
 			r.HandleFunc(testStatsRestPath, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				fmt.Fprintf(w, "%v", stats)
+				fmt.Fprintf(w, "%v", testStats)
 			}))
 			ts := httptest.NewUnstartedServer(r)
 			l, err := net.Listen("unix", testUDSPath)

diff --git a/agent/stats/mock/engine.go b/agent/stats/mock/engine.go
diff --git a/agent/stats/queue.go b/agent/stats/queue.go
@@ -19,6 +19,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/aws/amazon-ecs-agent/ecs-agent/stats"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
 	"github.com/cihub/seelog"
 	"github.com/docker/docker/api/types"
@@ -36,7 +37,7 @@ type Queue struct {
 	buffer                []UsageStats
 	maxSize               int
 	lastStat              *types.StatsJSON
-	lastNetworkStatPerSec *NetworkStatsPerSec
+	lastNetworkStatPerSec *stats.NetworkStatsPerSec
 	lock                  sync.RWMutex
 }
 
@@ -129,7 +130,7 @@ func (queue *Queue) add(rawStat *ContainerStats) {
 		}
 
 		if stat.NetworkStats != nil {
-			networkStatPerSec := &NetworkStatsPerSec{
+			networkStatPerSec := &stats.NetworkStatsPerSec{
 				RxBytesPerSecond: stat.NetworkStats.RxBytesPerSecond,
 				TxBytesPerSecond: stat.NetworkStats.TxBytesPerSecond,
 			}
@@ -148,7 +149,7 @@ func (queue *Queue) GetLastStat() *types.StatsJSON {
 	return queue.lastStat
 }
 
-func (queue *Queue) GetLastNetworkStatPerSec() *NetworkStatsPerSec {
+func (queue *Queue) GetLastNetworkStatPerSec() *stats.NetworkStatsPerSec {
 	queue.lock.RLock()
 	defer queue.lock.RUnlock()
 

diff --git a/agent/stats/types.go b/agent/stats/types.go
@@ -94,8 +94,3 @@ type taskDefinition struct {
 	family  string
 	version string
 }
-
-type NetworkStatsPerSec struct {
-	RxBytesPerSecond float32 `json:"rx_bytes_per_sec"`
-	TxBytesPerSecond float32 `json:"tx_bytes_per_sec"`
-}
diff --git a/agent/tcs/handler/handler_test.go b/agent/tcs/handler/handler_test.go
@@ -32,10 +32,10 @@ import (
 	mock_api "github.com/aws/amazon-ecs-agent/agent/api/mocks"
 	"github.com/aws/amazon-ecs-agent/agent/config"
 	mock_engine "github.com/aws/amazon-ecs-agent/agent/engine/mocks"
-	"github.com/aws/amazon-ecs-agent/agent/stats"
 	"github.com/aws/amazon-ecs-agent/agent/version"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/eventstream"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/stats"
 	tcsclient "github.com/aws/amazon-ecs-agent/ecs-agent/tcs/client"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"

diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/stats/types.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/stats/types.go
diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state/response.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state/response.go
diff --git a/agent/vendor/modules.txt b/agent/vendor/modules.txt
@@ -29,6 +29,7 @@ github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit/mocks
 github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit/request
 github.com/aws/amazon-ecs-agent/ecs-agent/logger/field
 github.com/aws/amazon-ecs-agent/ecs-agent/metrics
+github.com/aws/amazon-ecs-agent/ecs-agent/stats
 github.com/aws/amazon-ecs-agent/ecs-agent/tcs/client
 github.com/aws/amazon-ecs-agent/ecs-agent/tcs/handler
 github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs

diff --git a/ecs-agent/go.mod b/ecs-agent/go.mod
@@ -35,6 +35,8 @@ require (
 	github.com/google/uuid v1.3.0 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
+	github.com/opencontainers/go-digest v1.0.0 // indirect
+	github.com/opencontainers/image-spec v1.0.3-0.20211202183452-c5a74bcca799 // indirect
 	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/mod v0.10.0 // indirect
@@ -44,4 +46,5 @@ require (
 	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	gotest.tools/v3 v3.5.0 // indirect
 )
diff --git a/ecs-agent/go.sum b/ecs-agent/go.sum
@@ -95,6 +95,10 @@ github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3Rllmb
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
+github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
+github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
+github.com/opencontainers/image-spec v1.0.3-0.20211202183452-c5a74bcca799 h1:rc3tiVYb5z54aKaDfakKn0dDjIyPpTtszkjuMzyt7ec=
+github.com/opencontainers/image-spec v1.0.3-0.20211202183452-c5a74bcca799/go.mod h1:BtxoFyWECRxE4U/7sNtV5W15zMzWCbyJoFRP3s7yZA0=
 github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
 github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pborman/uuid v1.2.1 h1:+ZZIw58t/ozdjRaXh/3awHfmWRbzYxJoAdNJxe/3pvw=
@@ -220,6 +224,8 @@ gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 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.5.0 h1:Ljk6PdHdOhAb5aDMWXjDLMMhph+BpztA4v1QdqEW2eY=
+gotest.tools/v3 v3.5.0/go.mod h1:isy3WKz7GK6uNw/sbHzfKBLvlvXwUyV06n6brMxxopU=
 k8s.io/apimachinery v0.26.0-alpha.0 h1:cnXW2EigxCOrD+s52R9r5AZOcu1Nbv508gOCgSQkbo4=
 k8s.io/apimachinery v0.26.0-alpha.0/go.mod h1:qMx9eAk0sZQGsXGu86fab8tZdffHbwUfsvzqKn4mfB0=
 k8s.io/klog/v2 v2.100.1 h1:7WCHKK6K8fNhTqfBhISHQ97KrnJNFZMcQvKp7gP/tmg=

diff --git a/ecs-agent/stats/types.go b/ecs-agent/stats/types.go
@@ -0,0 +1,18 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+package stats
+
+type NetworkStatsPerSec struct {
+	RxBytesPerSecond float32 `json:"rx_bytes_per_sec"`
+	TxBytesPerSecond float32 `json:"tx_bytes_per_sec"`
+}