Add RestartCount to introspection endpoint container response

aws · May 23, 2024 · b4decae · b4decae
1 parent ee0eae9
commit b4decae
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 10 deletions.
diff --git a/agent/handlers/v1/response.go b/agent/handlers/v1/response.go
@@ -48,16 +48,17 @@ type TasksResponse struct {
 
 // ContainerResponse is the schema for the container response JSON object
 type ContainerResponse struct {
-	DockerID   string                        `json:"DockerId"`
-	DockerName string                        `json:"DockerName"`
-	Name       string                        `json:"Name"`
-	Image      string                        `json:"Image"`
-	ImageID    string                        `json:"ImageID"`
-	CreatedAt  *time.Time                    `json:"CreatedAt,omitempty"`
-	StartedAt  *time.Time                    `json:"StartedAt,omitempty"`
-	Ports      []tmdsresponse.PortResponse   `json:"Ports,omitempty"`
-	Networks   []tmdsresponse.Network        `json:"Networks,omitempty"`
-	Volumes    []tmdsresponse.VolumeResponse `json:"Volumes,omitempty"`
+	DockerID     string                        `json:"DockerId"`
+	DockerName   string                        `json:"DockerName"`
+	Name         string                        `json:"Name"`
+	Image        string                        `json:"Image"`
+	ImageID      string                        `json:"ImageID"`
+	CreatedAt    *time.Time                    `json:"CreatedAt,omitempty"`
+	StartedAt    *time.Time                    `json:"StartedAt,omitempty"`
+	Ports        []tmdsresponse.PortResponse   `json:"Ports,omitempty"`
+	Networks     []tmdsresponse.Network        `json:"Networks,omitempty"`
+	Volumes      []tmdsresponse.VolumeResponse `json:"Volumes,omitempty"`
+	RestartCount *int                          `json:"RestartCount,omitempty"`
 }
 
 // NewTaskResponse creates a TaskResponse for a task.
@@ -121,6 +122,10 @@ func NewContainerResponse(dockerContainer *apicontainer.DockerContainer, eni *ni
 		startedAt = startedAt.UTC()
 		resp.StartedAt = &startedAt
 	}
+	if container.RestartPolicyEnabled() {
+		restartCount := container.RestartTracker.GetRestartCount()
+		resp.RestartCount = &restartCount
+	}
 	return resp
 }
 

diff --git a/agent/handlers/v1/response_test.go b/agent/handlers/v1/response_test.go
@@ -23,6 +23,7 @@ import (
 
 	apicontainer "github.com/aws/amazon-ecs-agent/agent/api/container"
 	apitask "github.com/aws/amazon-ecs-agent/agent/api/task"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/container/restart"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
 	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/response"
@@ -197,6 +198,60 @@ func TestContainerResponse(t *testing.T) {
 	assert.Equal(t, expectedContainerResponse, containerResponse)
 }
 
+func TestContainerResponseWithRestartPolicy(t *testing.T) {
+	restartPolicy := &restart.RestartPolicy{
+		Enabled: true,
+	}
+	container := &apicontainer.Container{
+		Name:    containerName,
+		Image:   imageName,
+		ImageID: imageID,
+		Ports: []apicontainer.PortBinding{
+			{
+				ContainerPort: 80,
+				Protocol:      apicontainer.TransportProtocolTCP,
+			},
+		},
+		VolumesUnsafe: []types.MountPoint{
+			{
+				Name:        volName,
+				Source:      volSource,
+				Destination: volDestination,
+			},
+		},
+		RestartPolicy:  restartPolicy,
+		RestartTracker: restart.NewRestartTracker(*restartPolicy),
+	}
+
+	container.SetCreatedAt(containerTime)
+	container.SetStartedAt(containerTime)
+
+	dockerContainer := &apicontainer.DockerContainer{
+		DockerID:   containerID,
+		DockerName: containerName,
+		Container:  container,
+	}
+
+	eni := &ni.NetworkInterface{
+		IPV4Addresses: []*ni.IPV4Address{
+			{
+				Address: eniIPv4Address,
+			},
+		},
+	}
+
+	containerResponse := NewContainerResponse(dockerContainer, eni)
+
+	_, err := json.Marshal(containerResponse)
+	assert.NoError(t, err)
+
+	expectedContainerResponseWithRestartPolicy := expectedContainerResponse
+	rc := int(1)
+	expectedContainerResponseWithRestartPolicy.RestartCount = &rc
+
+	assert.Equal(t, expectedContainerResponse, containerResponse)
+}
+
 func TestPortBindingsResponse(t *testing.T) {
 	container := &apicontainer.Container{
 		Name: containerName,