Write pidfile atomically

Author: stepancheg

runsc/cmd/BUILD

@@ -58,6 +58,7 @@ go_library(
         "mitigate_extras.go",
         "path.go",
         "pause.go",
+        "pidfile.go",
         "platforms.go",
         "portforward.go",
         "ps.go",
@@ -145,6 +146,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,59 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strconv"
+)
+
+// WritePidFile writes pid file atomically if possible.
+func WritePidFile(path string, pid int) error {
+	pidStr := []byte(strconv.Itoa(pid))
+
+	_, err := os.Stat(path)
+	if err != nil && !os.IsNotExist(err) {
+		// 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)
+		}
+		return nil
+	}
+	if err != nil {
+		return fmt.Errorf("stat file %s failed: %w", path, err)
+	}
+
+	// Otherwise write using temp file to make write atomic.
+	dir := filepath.Dir(path)
+	tempFile, err := os.CreateTemp(dir, "pid-tmp-*")
+	if err != nil {
+		return fmt.Errorf("failed to create temp pid file in dir %s: %w", dir, err)
+	}
+
+	tempFileRenamed := false
+	defer func(tempFile *os.File) {
+		_ = tempFile.Close()
+		if !tempFileRenamed {
+			_ = os.Remove(tempFile.Name())
+		}
+	}(tempFile)
+
+	if err := os.Chmod(tempFile.Name(), 0644); err != nil {
+		return fmt.Errorf("failed to chmod pid file %s: %w", tempFile.Name(), err)
+	}
+
+	if _, err := tempFile.Write(pidStr); err != nil {
+		return fmt.Errorf("failed to write pid file %s: %w", tempFile.Name(), err)
+	}
+
+	if err := tempFile.Close(); err != nil {
+		return fmt.Errorf("failed to close temp pid file %s: %w", tempFile.Name(), err)
+	}
+
+	if err := os.Rename(tempFile.Name(), path); err != nil {
+		return fmt.Errorf("failed to rename temp pid file %s -> %s: %w", tempFile.Name(), path, err)
+	}
+	tempFileRenamed = true
+
+	return nil
+}

runsc/cmd/pidfile_test.go

@@ -0,0 +1,50 @@
+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)
+	}()
+
+	t.Run("Write new file", func(t *testing.T) {
+		path := filepath.Join(tempDir, "test-new.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'")
+		}
+	})
+
+	t.Run("Overwrite existing file", func(t *testing.T) {
+		path := filepath.Join(tempDir, "test-overwrite.pid")
+		if err := os.WriteFile(path, []byte("11"), 0600); err != nil {
+			t.Fatalf("failed to write pid file: %v", err)
+		}
+
+		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'")
+		}
+	})
+}