Write pidfile atomically

Author: stepancheg

runsc/cmd/BUILD

@@ -57,6 +57,7 @@ go_library(
         "mitigate_extras.go",
         "path.go",
         "pause.go",
+        "pidfile.go",
         "platforms.go",
         "portforward.go",
         "ps.go",
@@ -144,6 +145,7 @@ go_test(
         "install_test.go",
         "list_test.go",
         "mitigate_test.go",
+        "pidfile_test.go",
     ],
     data = [
         "//runsc",

runsc/cmd/exec.go

@@ -197,17 +197,16 @@ func (ex *Exec) exec(conf *config.Config, c *container.Container, e *control.Exe
 
 	// Write the sandbox-internal pid if required.
 	if ex.internalPidFile != "" {
-		pidStr := []byte(strconv.Itoa(int(pid)))
-		if err := os.WriteFile(ex.internalPidFile, pidStr, 0644); err != nil {
-			return util.Errorf("writing internal pid file %q: %v", ex.internalPidFile, err)
+		if err := WritePidFile(ex.internalPidFile, int(pid)); err != nil {
+			return util.Errorf("writing internal pid file: %v", err)
 		}
 	}
 
 	// Generate the pid file after the internal pid file is generated, so that
 	// users can safely assume that the internal pid file is ready after
 	// `runsc exec -d` returns.
 	if ex.pidFile != "" {
-		if err := os.WriteFile(ex.pidFile, []byte(strconv.Itoa(os.Getpid())), 0644); err != nil {
+		if err := WritePidFile(ex.pidFile, os.Getpid()); err != nil {
 			return util.Errorf("writing pid file: %v", err)
 		}
 	}

runsc/cmd/pidfile.go

@@ -0,0 +1,44 @@
+package cmd
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"os"
+	"strconv"
+)
+
+func fileExists(path string) bool {
+	_, err := os.Stat(path)
+	return !os.IsNotExist(err)
+}
+
+// WritePidFile writes pid file atomically if possible.
+func WritePidFile(path string, pid int) error {
+	pidStr := []byte(strconv.Itoa(pid))
+
+	if fileExists(path) {
+		// If path exists, write in place, because file could be pipe or something.
+		if err := os.WriteFile(path, pidStr, 0644); err != nil {
+			return fmt.Errorf("failed to write pid file %s: %w", path, err)
+		}
+	} else {
+		// Otherwise write using temp file to make write atomic.
+		b := make([]byte, 8)
+		_, err := rand.Read(b)
+		if err != nil {
+			return fmt.Errorf("failed to generate random bytes: %w", err)
+		}
+		tempPath := path + fmt.Sprintf(".%x", hex.EncodeToString(b))
+
+		if err := os.WriteFile(tempPath, pidStr, 0644); err != nil {
+			return fmt.Errorf("failed to write temp pid file %s: %w", tempPath, err)
+		}
+
+		if err := os.Rename(tempPath, path); err != nil {
+			return fmt.Errorf("failed to rename temp pid file %s -> %s: %w", tempPath, path, err)
+		}
+	}
+
+	return nil
+}

runsc/cmd/pidfile_test.go

@@ -0,0 +1,45 @@
+package cmd
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestWritePidFile(t *testing.T) {
+	tempDir, err := os.MkdirTemp("", "test-*")
+	if err != nil {
+		t.Fatalf("failed to create temp dir: %v", err)
+	}
+	defer func() {
+		_ = os.RemoveAll(tempDir)
+	}()
+
+	// Write new file.
+
+	path := filepath.Join(tempDir, "test.pid")
+	if err := WritePidFile(path, 17); err != nil {
+		t.Fatalf("failed to write pid file: %v", err)
+	}
+
+	pidStr, err := os.ReadFile(path)
+	if err != nil {
+		t.Fatalf("failed to read pid file: %v", err)
+	}
+	if string(pidStr) != "17" {
+		t.Fatalf("pid file did not contain pid '17'")
+	}
+
+	// Overwrite existing file.
+
+	if err := WritePidFile(path, 19); err != nil {
+		t.Fatalf("failed to overwrite write pid file: %v", err)
+	}
+	pidStr, err = os.ReadFile(path)
+	if err != nil {
+		t.Fatalf("failed to read pid file: %v", err)
+	}
+	if string(pidStr) != "19" {
+		t.Fatalf("pid file did not contain pid '19'")
+	}
+}