Fixing bug to successfully kill AppHealth Gracefully & VMWatch.

.devcontainer/Dockerfile

@@ -1,4 +1,4 @@
-FROM mcr.microsoft.com/devcontainers/go:1.18-bullseye
+FROM mcr.microsoft.com/devcontainers/go:1.21-bullseye
 
 RUN apt-get -qqy update && \
 	apt-get -qqy install jq openssl ca-certificates && \
@@ -16,3 +16,32 @@ RUN mkdir -p /var/lib/waagent && \
 COPY extension-settings.json /var/lib/waagent/Extension/config/0.settings
 # install go tools we need for build
 RUN go install github.com/ahmetb/govvv@latest
+
+# Install npm
+RUN apt-get update && \
+    apt-get install -y npm
+
+# Updating npm to the latest version
+RUN npm cache clean -f && npm install -g n && n stable
+
+# Install dev enviroment dependencies
+RUN npm install bats -g
+
+#Install Bats-Assert and Bats-Support
+RUN npm install -g https://github.com/bats-core/bats-assert && \
+        npm install -g https://github.com/bats-core/bats-support
+
+# Install Parallel 
+RUN apt-get install -y parallel
+
+# Install Docker
+RUN apt-get install runc -y && \
+        apt-get install containerd -y && \
+        apt-get install docker.io -y
+
+# Install Docker Engine
+RUN curl -fsSL https://test.docker.com -o test-docker.sh && \
+        sh test-docker.sh
+
+# Creating ENV variables
+ENV CUSTOM_BATS_LIB_PATH /usr/local/lib/node_modules

.devcontainer/devcontainer.json

@@ -23,14 +23,22 @@
             "settings": {},
             // Add the IDs of extensions you want installed when the container is created.
             "extensions": [
-                "golang.go"
-            ]
+                "golang.go",
+                "ms-azuretools.vscode-docker"
+            ],
+            "recommendations": [
+                "GitHub.copilot",
+                "GitHub.copilot-chat",
+                "GitHub.vscode-pull-request-github"
+            ] 
         }
     },
     "remoteUser": "root",
     // Use 'postCreateCommand' to run commands after the container is created.
-    "postCreateCommand": "echo hello",
-    "mounts": [ ]
+    "postCreateCommand": "go mod download && echo hello",
+    "mounts": [ 
+        "source=/var/run/docker.sock,target=/var/run/docker.sock,type=bind" 
+    ]
 
     // // Use 'forwardPorts' to make a list of ports inside the container available locally.
     // "forwardPorts": [5000, 5001],B

.devcontainer/extension-settings.json

@@ -8,7 +8,7 @@
                     "port": 8080,
                     "numberOfProbes": 1,
                     "intervalInSeconds": 5,
-                    "gracePeriod": 600,
+                    "gracePeriod": 10,
                     "vmWatchSettings": {
                         "enabled": true,
                         "signalFilters": {

.vscode/launch.json

@@ -27,5 +27,24 @@
             ],
             "preLaunchTask": "make devcontainer"
         },
+        {
+            "name": "Run integration tests",
+            "type": "node-terminal",
+            "request": "launch",
+            "command": "./integration-test/run.sh",
+            "cwd": "${workspaceFolder}",
+            "preLaunchTask": "make binary"
+        },
+        {
+            "name": "devcontainer run - enable NOBUILD",
+            "type": "go",
+            "request": "launch",
+            "mode": "exec",
+            "program": "/var/lib/waagent/Extension/bin/applicationhealth-extension",
+            "cwd": "${workspaceFolder}",
+            "args" : [
+                "enable"
+            ]
+        }
     ]
 }

integration-test/env/extension-test-helpers.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+set -uex
+logFilePath="/var/log/azure/Extension/force-kill-extension.txt"
+
+force_kill_apphealth() {
+    app_health_pid=$(ps -ef | grep "applicationhealth-extension" | grep -v grep | grep -v tee | awk '{print $2}')
+
+    if [ -z "$app_health_pid" ]; then
+        echo "Applicationhealth extension is not running" > $logFilePath
+        return 0
+    fi
+    echo "Killing the applicationhealth extension forcefully" >> $logFilePath
+    kill -9 $app_health_pid
+
+    output=$(check_running_processes)
+    if [ "$output" == "Applicationhealth and VMWatch are not running" ]; then
+        echo "$output" >> /var/log/azure/Extension/force-kill-extension.txt
+        echo "Successfully killed the apphealth extension" >> $logFilePath
+        echo "Successfully killed the VMWatch extension" >> $logFilePath
+    else
+        echo "$output" >> /var/log/azure/Extension/force-kill-extension.txt
+        echo "Failed to kill the apphealth extension" >> $logFilePath
+    fi
+
+}
+
+check_running_processes() {
+    local output=$(ps -ef | grep -e "applicationhealth-extension" -e "vmwatch_linux_amd64" | grep -v grep | grep -v tee)
+    if [ -z "$output" ]; then
+        echo "Applicationhealth and VMWatch are not running"
+    else
+        if [ -n "$(echo $output | grep "applicationhealth-extension")" ]; then
+            echo "Applicationhealth is running"
+        fi
+        if [ -n "$(echo $output | grep "vmwatch_linux_amd64")" ]; then
+            echo "VMWatch is running"
+        fi
+        echo "$output"
+    fi
+}

integration-test/run.sh

@@ -21,10 +21,10 @@ done
 source integration-test/test/test_helper.bash
 create_certificate
 # Run Sequential Integration Tests
-sudo bats integration-test/test/sequential -T --trace 
+bats integration-test/test/sequential -T --trace 
 err1=$?
 # Run Parallel Integration Tests
-sudo bats integration-test/test/parallel  --jobs 10 -T --trace $FILTER
+bats integration-test/test/parallel  --jobs 10 -T --trace $FILTER
 err2=$?
 delete_certificate
 exit $((err1 + err2))

integration-test/test/parallel/7_vmwatch.bats

@@ -382,7 +382,6 @@ teardown(){
 
     [[ "$output" == *'Invoking: /var/lib/waagent/Extension/bin/applicationhealth-shim disable'* ]]
     [[ "$output" == *'applicationhealth-extension process terminated'* ]]
-    [[ "$output" == *'vmwatch_linux_amd64 process terminated'* ]]
 
     status_file="$(container_read_extension_status)"
     verify_status_item "$status_file" Disable success "Disable succeeded"
@@ -411,10 +410,86 @@ teardown(){
 
     [[ "$output" == *'Invoking: /var/lib/waagent/Extension/bin/applicationhealth-shim uninstall'* ]]
     [[ "$output" == *'applicationhealth-extension process terminated'* ]]
-    [[ "$output" == *'vmwatch_linux_amd64 process terminated'* ]]
     [[ "$output" == *'operation=uninstall seq=0 path=/var/lib/waagent/apphealth event=uninstalled'* ]]
 }
 
+@test "handler command: enable - Graceful Shutdown - vm watch killed when Apphealth is killed gracefully with SIGTERM" {
+    mk_container $container_name bash -c "nc -l localhost 22 -k & fake-waagent install && export RUNNING_IN_DEV_CONTAINER=1 && export ALLOW_VMWATCH_CGROUP_ASSIGNMENT_FAILURE=1 && fake-waagent enable && wait-for-enable webserverexit && sleep 2 && source /var/lib/waagent/test_helper.bash;kill_apphealth_extension_gracefully SIGTERM & sleep 2"
+    push_settings '
+    {
+        "protocol": "http",
+        "requestPath": "/",
+        "port": 8080,
+        "numberOfProbes": 1,
+        "intervalInSeconds": 5,
+        "gracePeriod": 600,
+        "vmWatchSettings": {
+            "enabled": true
+        }
+    }' ''
+    run start_container
+    echo "$output"
+    [[ "$output" == *'Setup VMWatch command: /var/lib/waagent/Extension/bin/VMWatch/vmwatch_linux_amd64'* ]]
+    [[ "$output" == *'VMWatch process started'* ]]
+    [[ "$output" == *'VMWatch is running'* ]]
+
+    [[ "$output" == *'event="Received shutdown request"'* ]]
+    [[ "$output" == *'Successfully killed VMWatch process with PID'* ]]
+    [[ "$output" == *'Application health process terminated'* ]]
+}
+
+@test "handler command: enable - Graceful Shutdown - vm watch killed when Apphealth is killed gracefully with SIGINT" {
+    mk_container $container_name bash -c "nc -l localhost 22 -k & fake-waagent install && export RUNNING_IN_DEV_CONTAINER=1 && export ALLOW_VMWATCH_CGROUP_ASSIGNMENT_FAILURE=1 && fake-waagent enable && wait-for-enable webserverexit && sleep 2 && source /var/lib/waagent/test_helper.bash;kill_apphealth_extension_gracefully SIGINT & sleep 2"
+    push_settings '
+    {
+        "protocol": "http",
+        "requestPath": "/",
+        "port": 8080,
+        "numberOfProbes": 1,
+        "intervalInSeconds": 5,
+        "gracePeriod": 600,
+        "vmWatchSettings": {
+            "enabled": true
+        }
+    }' ''
+    run start_container
+    echo "$output"
+    [[ "$output" == *'Setup VMWatch command: /var/lib/waagent/Extension/bin/VMWatch/vmwatch_linux_amd64'* ]]
+    [[ "$output" == *'VMWatch process started'* ]]
+    [[ "$output" == *'VMWatch is running'* ]]
+
+    [[ "$output" == *'event="Received shutdown request"'* ]]
+    [[ "$output" == *'Successfully killed VMWatch process with PID'* ]]
+    [[ "$output" == *'Application health process terminated'* ]]
+}
+
+@test "handler command: enable - Forced Shutdown - vm watch killed when Apphealth is killed gracefully with SIGKILL" {
+    mk_container $container_name bash -c "nc -l localhost 22 -k & export RUNNING_IN_DEV_CONTAINER=1 && export ALLOW_VMWATCH_CGROUP_ASSIGNMENT_FAILURE=1 && fake-waagent enable && wait-for-enable webserverexit && sleep 10 && source /var/lib/waagent/extension-test-helpers.sh;force_kill_apphealth"
+    push_settings '
+    {
+        "protocol": "tcp",
+        "requestPath": "",
+        "port": 22,
+        "numberOfProbes": 1,
+        "intervalInSeconds": 5,
+        "gracePeriod": 600,
+        "vmWatchSettings": {
+            "enabled": true
+        }
+    }' ''
+    run start_container
+
+    echo "$output"
+    shutdown_log="$(container_read_file /var/log/azure/Extension/force-kill-extension.txt)"
+    echo "$shutdown_log"
+    [[ "$output" == *'Setup VMWatch command: /var/lib/waagent/Extension/bin/VMWatch/vmwatch_linux_amd64'* ]]
+    [[ "$output" == *'VMWatch process started'* ]]
+    [[ "$output" == *'VMWatch is running'* ]]
+
+    [[ "$shutdown_log" == *'Successfully killed the apphealth extension'* ]]
+    [[ "$shutdown_log" == *'Successfully killed the VMWatch extension'* ]]
+}
+
 @test "handler command: enable/uninstall - vm passes memory to commandline" {
     mk_container $container_name sh -c "nc -l localhost 22 -k & fake-waagent install && export RUNNING_IN_DEV_CONTAINER=1 && export ALLOW_VMWATCH_CGROUP_ASSIGNMENT_FAILURE=1 && fake-waagent enable && wait-for-enable webserverexit && sleep 5 && fake-waagent uninstall"
     push_settings '

integration-test/test/test_helper.bash

@@ -271,3 +271,21 @@ get_extension_version() {
     version=$(awk -F'[<>]' '/<Version>/ {print $3}' misc/manifest.xml)
     echo $version
 }
+# Accepted Kill Signals SIGINT SIGTERM
+kill_apphealth_extension_gracefully() {
+    # kill the applicationhealth extension gracefully
+    # echo "Printing the process list Before killing the applicationhealth extension"
+    ps -ef | grep -e "applicationhealth-extension" -e "vmwatch_linux_amd64" | grep -v grep
+    kill_signal=$1
+    [[ $kill_signal == "SIGINT" || $kill_signal == "SIGTERM" ]] || { echo "Invalid signal: $kill_signal"; return 1; }
+    app_health_pid=$(ps -ef | grep "applicationhealth-extension" | grep -v grep | grep -v tee | awk '{print $2}')
+    if [ -z "$app_health_pid" ]; then
+        echo "Applicationhealth extension is not running"
+        return 0
+    fi
+    # echo "Killing applicationhealth extension with signal: $kill_signal"
+    # echo "PID: $app_health_pid"
+    kill -s $kill_signal $app_health_pid
+    # echo "Printing the process list after killing the applicationhealth extension"
+    ps -ef | grep -e "applicationhealth-extension" -e "vmwatch_linux_amd64" | grep -v grep
+}

main/main.go

@@ -35,7 +35,7 @@ func main() {
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
 	go func() {
 		<-sigs
-		ctx.Log("event", fmt.Sprintf("Received shutdown request"))
+		ctx.Log("event", "Received shutdown request")
 		shutdown = true
 		err := killVMWatch(ctx, vmWatchCommand)
 		if err != nil {

main/vmWatch.go

@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/containerd/cgroups/v3"
@@ -101,7 +102,6 @@ func executeVMWatch(ctx *log.Context, s *vmWatchSettings, hEnv HandlerEnvironmen
 
 func executeVMWatchHelper(ctx *log.Context, attempt int, vmWatchSettings *vmWatchSettings, hEnv HandlerEnvironment) (err error) {
 	pid := -1
-	var cmd *exec.Cmd
 	defer func() {
 		if r := recover(); r != nil {
 			err = fmt.Errorf("Error: %w\n Additonal Details: %+v", err, r)
@@ -110,28 +110,29 @@ func executeVMWatchHelper(ctx *log.Context, attempt int, vmWatchSettings *vmWatc
 	}()
 
 	// Setup command
-	cmd, err = setupVMWatchCommand(vmWatchSettings, hEnv)
+	vmWatchCommand, err = setupVMWatchCommand(vmWatchSettings, hEnv)
 	if err != nil {
 		err = fmt.Errorf("[%v][PID -1] Attempt %d: VMWatch setup failed. Error: %w", time.Now().UTC().Format(time.RFC3339), attempt, err)
 		ctx.Log("error", err.Error())
 		return err
 	}
 
-	ctx.Log("event", fmt.Sprintf("Attempt %d: Setup VMWatch command: %s\nArgs: %v\nDir: %s\nEnv: %v\n", attempt, cmd.Path, cmd.Args, cmd.Dir, cmd.Env))
+	ctx.Log("event", fmt.Sprintf("Attempt %d: Setup VMWatch command: %s\nArgs: %v\nDir: %s\nEnv: %v\n", attempt, vmWatchCommand.Path, vmWatchCommand.Args, vmWatchCommand.Dir, vmWatchCommand.Env))
 
 	// TODO: Combined output may get excessively long, especially since VMWatch is a long running process
 	// We should trim the output or only get from Stderr
 	combinedOutput := &bytes.Buffer{}
-	cmd.Stdout = combinedOutput
-	cmd.Stderr = combinedOutput
+	vmWatchCommand.Stdout = combinedOutput
+	vmWatchCommand.Stderr = combinedOutput
+	vmWatchCommand.SysProcAttr = &syscall.SysProcAttr{ Pdeathsig: syscall.SIGTERM }
 
 	// Start command
-	if err := cmd.Start(); err != nil {
+	if err := vmWatchCommand.Start(); err != nil {
 		err = fmt.Errorf("[%v][PID -1] Attempt %d: VMWatch failed to start. Error: %w\nOutput: %s", time.Now().UTC().Format(time.RFC3339), attempt, err, combinedOutput.String())
 		ctx.Log("error", err.Error())
 		return err
 	}
-	pid = cmd.Process.Pid // cmd.Process should be populated on success
+	pid = vmWatchCommand.Process.Pid // cmd.Process should be populated on success
 	ctx.Log("event", fmt.Sprintf("Attempt %d: VMWatch process started with pid %d", attempt, pid))
 
 	err = createAndAssignCgroups(ctx, vmWatchSettings, pid)
@@ -147,7 +148,7 @@ func executeVMWatchHelper(ctx *log.Context, attempt int, vmWatchSettings *vmWatc
 		// this allows us to test both cases
 		if os.Getenv(AllowVMWatchCgroupAssignmentFailureVariableName) == "" || os.Getenv(RunningInDevContainerVariableName) == "" {
 			ctx.Log("event", fmt.Sprintf("Killing VMWatch process as cgroup assigment failed"))
-			_ = killVMWatch(ctx, cmd)
+			_ = killVMWatch(ctx, vmWatchCommand)
 			return err
 		}
 	}
@@ -160,15 +161,15 @@ func executeVMWatchHelper(ctx *log.Context, attempt int, vmWatchSettings *vmWatc
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		err = cmd.Wait()
+		err = vmWatchCommand.Wait()
 		processDone <- true
 		close(processDone)
 	}()
 	// add a task to monitor heartbeat
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		monitorHeartBeat(ctx, GetVMWatchHeartbeatFilePath(hEnv), processDone, cmd)
+		monitorHeartBeat(ctx, GetVMWatchHeartbeatFilePath(hEnv), processDone, vmWatchCommand)
 	}()
 	wg.Wait()
 	err = fmt.Errorf("[%v][PID %d] Attempt %d: VMWatch process exited. Error: %w\nOutput: %s", time.Now().UTC().Format(time.RFC3339), pid, attempt, err, combinedOutput.String())
@@ -205,7 +206,7 @@ func monitorHeartBeat(ctx *log.Context, heartBeatFile string, processDone chan b
 
 func killVMWatch(ctx *log.Context, cmd *exec.Cmd) error {
 	if cmd == nil || cmd.Process == nil || cmd.ProcessState != nil {
-		ctx.Log("event", fmt.Sprintf("VMWatch is not running, killing process is not necessary."))
+		ctx.Log("event", "VMWatch is not running, killing process is not necessary.")
 		return nil
 	}
 

test.Dockerfile

@@ -27,3 +27,6 @@ RUN ln -s /var/lib/waagent/fake-waagent /sbin/fake-waagent && \
 COPY misc/HandlerManifest.json ./Extension/
 COPY misc/applicationhealth-shim ./Extension/bin/
 COPY bin/applicationhealth-extension ./Extension/bin/
+
+# Copy Helper functions and scripts
+COPY integration-test/test/test_helper.bash /var/lib/waagent