From 69381adcb7fc396bc765efb16031512099deb54e Mon Sep 17 00:00:00 2001
From: Yajie Chu <cya@amazon.com>
Date: Wed, 6 Nov 2019 14:39:43 -0800
Subject: [PATCH] support healthcheck of ecs agent

---
 agent/app/args/flag.go                 |  5 +++++
 agent/app/healthcheck.go               | 31 ++++++++++++++++++++++++++
 agent/app/run.go                       |  2 ++
 scripts/dockerfiles/Dockerfile.release |  3 +++
 4 files changed, 41 insertions(+)
 create mode 100644 agent/app/healthcheck.go

diff --git a/agent/app/args/flag.go b/agent/app/args/flag.go
index a1a078e78af..d542dcbfe2e 100644
--- a/agent/app/args/flag.go
+++ b/agent/app/args/flag.go
@@ -23,6 +23,7 @@ const (
 	licenseUsage             = "Print the LICENSE and NOTICE files and exit"
 	blacholeEC2MetadataUsage = "Blackhole the EC2 Metadata requests. Setting this option can cause the ECS Agent to fail to work properly.  We do not recommend setting this option"
 	windowsServiceUsage      = "Run the ECS agent as a Windows Service"
+	healthcheckServiceUsage  = "Run the agent healthcheck"
 
 	versionFlagName              = "version"
 	logLevelFlagName             = "loglevel"
@@ -31,6 +32,7 @@ const (
 	licenseFlagName              = "license"
 	blackholeEC2MetadataFlagName = "blackhole-ec2-metadata"
 	windowsServiceFlagName       = "windows-service"
+	healthCheckFlagName          = "healthcheck"
 )
 
 // Args wraps various ECS Agent arguments
@@ -51,6 +53,8 @@ type Args struct {
 	ECSAttributes *bool
 	// WindowsService indicates that the agent should run as a Windows service
 	WindowsService *bool
+	// Healthcheck indicates that agent should run healthcheck
+	Healthcheck *bool
 }
 
 // New creates a new Args object from the argument list
@@ -65,6 +69,7 @@ func New(arguments []string) (*Args, error) {
 		BlackholeEC2Metadata: flagset.Bool(blackholeEC2MetadataFlagName, false, blacholeEC2MetadataUsage),
 		ECSAttributes:        flagset.Bool(ecsAttributesFlagName, false, ecsAttributesUsage),
 		WindowsService:       flagset.Bool(windowsServiceFlagName, false, windowsServiceUsage),
+		Healthcheck:          flagset.Bool(healthCheckFlagName, false, healthcheckServiceUsage),
 	}
 
 	err := flagset.Parse(arguments)
diff --git a/agent/app/healthcheck.go b/agent/app/healthcheck.go
new file mode 100644
index 00000000000..e73c06a38e2
--- /dev/null
+++ b/agent/app/healthcheck.go
@@ -0,0 +1,31 @@
+// Copyright 2019 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 app
+
+import (
+	"net/http"
+
+	"github.com/aws/amazon-ecs-agent/agent/sighandlers/exitcodes"
+	"github.com/cihub/seelog"
+)
+
+// runHealthcheck runs the Agent's healthcheck
+func runHealthcheck() int {
+	_, err := http.Get("http://localhost:51678/v1/metadata")
+	if err != nil {
+		seelog.Warnf("Health check failed with error: %v", err)
+		return exitcodes.ExitError
+	}
+	return exitcodes.ExitSuccess
+}
diff --git a/agent/app/run.go b/agent/app/run.go
index 0fda102d7b9..a73895f8203 100644
--- a/agent/app/run.go
+++ b/agent/app/run.go
@@ -38,6 +38,8 @@ func Run(arguments []string) int {
 		return printLicense()
 	} else if *parsedArgs.Version {
 		return version.PrintVersion()
+	} else if *parsedArgs.Healthcheck {
+		return runHealthcheck()
 	}
 
 	logger.SetLevel(*parsedArgs.LogLevel)
diff --git a/scripts/dockerfiles/Dockerfile.release b/scripts/dockerfiles/Dockerfile.release
index ac22a2ed1ad..50b8b2b2c4e 100644
--- a/scripts/dockerfiles/Dockerfile.release
+++ b/scripts/dockerfiles/Dockerfile.release
@@ -28,4 +28,7 @@ COPY out/cni-plugins /amazon-ecs-cni-plugins
 COPY misc/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
 
 EXPOSE 51678 51679
+
+HEALTHCHECK CMD ["/agent", "--healthcheck"]
+
 ENTRYPOINT ["/agent"]