adjust OOM score of main agent process

agent/agent_unix.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/aws/amazon-ssm-agent/agent/appconfig"
 	logger "github.com/aws/amazon-ssm-agent/agent/log/ssmlog"
+	"github.com/aws/amazon-ssm-agent/agent/platform"
 )
 
 func main() {
@@ -20,6 +21,9 @@ func main() {
 	defer log.Close()
 	defer log.Flush()
 
+	// Adjust OOM score to protect the SSM agent from being killed by the OOM killer
+	platform.SetOOMScoreAdjust(log)
+
 	// parse input parameters
 	parseFlags(log)
 

agent/platform/platform_darwin.go

@@ -132,3 +132,7 @@ func fullyQualifiedDomainName(_ log.T) string {
 func isPlatformNanoServer(_ log.T) (bool, error) {
 	return false, nil
 }
+
+func SetOOMScoreAdjust(_ log.T) {
+	// No-op on macOS
+}

agent/platform/platform_unix.go

@@ -19,6 +19,7 @@ package platform
 
 import (
 	"errors"
+	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -49,6 +50,7 @@ const (
 var (
 	readAllText = fileutil.ReadAllText
 	fileExists  = fileutil.Exists
+	writeFile   = os.WriteFile
 
 	ErrFileNotFound   = errors.New("file not found")
 	ErrFilePermission = errors.New("no sufficient permissions")
@@ -309,3 +311,19 @@ func fullyQualifiedDomainName(log log.T) string {
 func isPlatformNanoServer(_ log.T) (bool, error) {
 	return false, nil
 }
+
+// Set the OOM score adjustment for the current process on supported platforms.
+// -1000 makes the process immune to the OOM killer under normal circumstances.
+func SetOOMScoreAdjust(log log.T) {
+	const oomScoreAdj = -1000
+	oomScoreAdjPath := "/proc/self/oom_score_adj"
+
+	if fileExists(oomScoreAdjPath) {
+		err := writeFile(oomScoreAdjPath, []byte(fmt.Sprintf("%d", oomScoreAdj)), 0644)
+		if err != nil {
+			log.Warnf("Failed to set OOM score adjustment: %v. Agent will be vulnerable to OOM killer.", err)
+			return
+	    }
+        log.Debugf("Successfully set OOM score adjustment to %d for SSM agent process", oomScoreAdj)
+    }
+}

agent/platform/platform_unix_test.go

@@ -19,6 +19,7 @@ package platform
 
 import (
 	"fmt"
+	"os"
 	"testing"
 
 	logger "github.com/aws/amazon-ssm-agent/agent/mocks/log"
@@ -223,3 +224,80 @@ func TestGetSystemInfoWithNonExistingParam(t *testing.T) {
 	GetSystemInfo(logMock, XenUuidSystemInfoParamKey)
 	assert.Equal(t, 2, cacheInitCount)
 }
+
+func TestSetOOMScoreAdjust(t *testing.T) {
+	// Save original functions
+	tmpFileExists := fileExists
+	tmpWriteFile := writeFile
+	defer func() {
+		fileExists = tmpFileExists
+		writeFile = tmpWriteFile
+	}()
+
+	// Mock file exists
+	fileExists = func(filePath string) bool {
+		return filePath == "/proc/self/oom_score_adj"
+	}
+
+	// Mock successful write
+	writeCalled := false
+	var writtenData []byte
+	writeFile = func(name string, data []byte, perm os.FileMode) error {
+		writeCalled = true
+		writtenData = data
+		return nil
+	}
+
+	logMock := logger.NewMockLog()
+	SetOOMScoreAdjust(logMock)
+
+	assert.True(t, writeCalled, "WriteFile should have been called")
+	assert.Equal(t, "-1000", string(writtenData), "Should write -1000")
+}
+
+func TestSetOOMScoreAdjustError(t *testing.T) {
+	// Save original functions
+	tmpFileExists := fileExists
+	tmpWriteFile := writeFile
+	defer func() {
+		fileExists = tmpFileExists
+		writeFile = tmpWriteFile
+	}()
+
+	fileExists = func(filePath string) bool {
+		return filePath == "/proc/self/oom_score_adj"
+	}
+
+	writeFile = func(name string, data []byte, perm os.FileMode) error {
+		return fmt.Errorf("permission denied")
+	}
+
+	logMock := logger.NewMockLog()
+	SetOOMScoreAdjust(logMock)
+
+	assert.NotNil(t, logMock)
+}
+
+func TestSetOOMScoreAdjustNotSupported(t *testing.T) {
+	tmpFileExists := fileExists
+	tmpWriteFile := writeFile
+	defer func() {
+		fileExists = tmpFileExists
+		writeFile = tmpWriteFile
+	}()
+
+	fileExists = func(filePath string) bool {
+		return false
+	}
+
+	writeCalled := false
+	writeFile = func(name string, data []byte, perm os.FileMode) error {
+		writeCalled = true
+		return nil
+	}
+
+	logMock := logger.NewMockLog()
+	SetOOMScoreAdjust(logMock)
+
+	assert.False(t, writeCalled, "WriteFile should not have been called when file doesn't exist")
+}