From b4e1d5d7ae7e4b710b7ec85afff0630082825645 Mon Sep 17 00:00:00 2001
From: Jakub Nyckowski <jakub.nyckowski@goteleport.com>
Date: Mon, 19 Dec 2022 21:50:27 -0500
Subject: [PATCH] Set OOM score to 0 for child processes (#19383)

* Set OOM score to 0 for child processes

Co-authored-by: Zac Bergquist <zac.bergquist@goteleport.com>
---
 lib/srv/reexec.go       |  6 +++++-
 lib/srv/reexec_linux.go | 28 +++++++++++++++++++++++++---
 lib/srv/reexec_other.go |  2 ++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/lib/srv/reexec.go b/lib/srv/reexec.go
index 484200756381d..40fa02122d160 100644
--- a/lib/srv/reexec.go
+++ b/lib/srv/reexec.go
@@ -1,5 +1,5 @@
 /*
-Copyright 2020 Gravitational, Inc.
+Copyright 2020-2022 Gravitational, Inc.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -389,6 +389,10 @@ func RunCommand() (errw io.Writer, code int, err error) {
 		}
 	}
 
+	if err := setNeutralOOMScore(); err != nil {
+		log.WithError(err).Warnf("failed to adjust OOM score")
+	}
+
 	// Start the command.
 	err = cmd.Start()
 	if err != nil {
diff --git a/lib/srv/reexec_linux.go b/lib/srv/reexec_linux.go
index f59150f63383e..757c265bc2708 100644
--- a/lib/srv/reexec_linux.go
+++ b/lib/srv/reexec_linux.go
@@ -1,4 +1,4 @@
-// Copyright 2021 Gravitational, Inc
+// Copyright 2021-2022 Gravitational, Inc
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,15 +18,17 @@
 package srv
 
 import (
+	"os"
 	"os/exec"
 	"syscall"
 
+	"github.com/gravitational/trace"
 	"golang.org/x/sys/unix"
 )
 
 func init() {
 	// errors in open/openat are signaled by returning -1, we don't really care
-	// about the specifics anyway so we can just ignore the error value
+	// about the specifics anyway, so we can just ignore the error value
 	//
 	// we're opening with O_PATH rather than O_RDONLY because the binary might
 	// not actually be readable (but only executable)
@@ -82,7 +84,27 @@ func userCommandOSTweaks(cmd *exec.Cmd) {
 	}
 	// Linux only: when parent process (this process) dies unexpectedly, kill
 	// the child process instead of orphaning it.
-	// SIGKILL because we don't control the child process and it could choose
+	// SIGKILL because we don't control the child process, and it could choose
 	// to ignore other signals.
 	cmd.SysProcAttr.Pdeathsig = syscall.SIGKILL
 }
+
+// setNeutralOOMScore sets the OOM score for the current process to 0 (the
+// middle between -1000 and 1000). This value is inherited by all child processes.
+func setNeutralOOMScore() error {
+	// Use os.OpenFile() instead of os.WriteFile() to avoid creating the file
+	// if for some extremely weird reason doesn't exist. Permission in this case
+	// won't be used as os.O_WRONLY won't create the file.
+	f, err := os.OpenFile("/proc/self/oom_score_adj", os.O_WRONLY, 0)
+	if err != nil {
+		return trace.ConvertSystemError(err)
+	}
+
+	if _, err := f.WriteString("0"); err != nil {
+		return trace.NewAggregate(err, f.Close())
+	}
+
+	// Make sure to return errors from Close(),
+	// as sync error may be returned here.
+	return trace.Wrap(f.Close())
+}
diff --git a/lib/srv/reexec_other.go b/lib/srv/reexec_other.go
index 4756af53faf4d..3e47e990914c5 100644
--- a/lib/srv/reexec_other.go
+++ b/lib/srv/reexec_other.go
@@ -24,3 +24,5 @@ import (
 func reexecCommandOSTweaks(cmd *exec.Cmd) {}
 
 func userCommandOSTweaks(cmd *exec.Cmd) {}
+
+func setNeutralOOMScore() error { return nil }