os: overhaul handling of PID vs pidfd within Process

There are several issues with pidfd handling today:

* The zero value of a Process makes the handle field appear valid, so
  methods attempt to use it as a pidfd rather than falling back to the
  PID as they should (#67634).

* If a process doesn't exist, FindProcess returns a Process with Pid ==
  -2, which is not a compatible change (#67640).

* pidfd close is racy as-is. A Release call or successful Wait will
  clear the handle field and close the pidfd. However, a concurrent call
  may have already loaded the handle field and could then proceed to use
  the closed FD (which could have been reopened as a different pidfd,
  targeting a different process) (#67641).

This CL performs multiple structural changes to the internals of
Process.

First and foremost, each method is refactored to clearly select either
pidfd or raw pid mode. Previously, raw pid mode was structured as a
fallback when pidfd mode is unavailable. This works fine, but it does
not make it clear that a given Process object either always uses pidfd
or always uses raw pid. Since each mode needs to handle different race
conditions, it helps to make it clear that we can't switch between modes
within a single Process object.

Second, pidfd close safety is handled by reference counting uses of the
FD. The last user of the FD will close the FD. For example, this means
that with concurrent Release and Signal, the Signal call may be the one
to close the FD. This is the bulk of this CL, though I find the end
result makes the overall implementation easier to reason about.

Third, the PID path handles a similar race condtion between Wait and
Kill: Wait frees the PID value in the kernel, which could be reallocated
causing Kill to target the wrong process. This is handled with a done
flag and a mutex. The done flag now shares the same state field used for
the handle.

Similarly, the Windows implementation reuses all of the handle reference
counting that Linux uses. This means the implementations more
consistent, and make Windows safe against the same handle reuse
problems. (Though I am unsure if Windows ever reuses handles).

Wait has a slight behavior change on Windows: previously Wait after
Release or an earlier Wait would hang indefinitely (WaitForSingleObject
on syscall.InvalidHandle waits indefinitely). Now it returns the same
errors as Linux (EINVAL and ErrProcessDone, respectively).

Similarly, Release on Windows no longer returns close errors, as it may
not actually be the place where the close occurs.

Fixes #67634.
Fixes #67640.
Fixes #67641.
Updates #67642.

Cq-Include-Trybots: luci.golang.try:gotip-linux-amd64-longtest,gotip-windows-amd64-longtest
Change-Id: I2ad998f7b67d32031e6f870e8533dbd55d3c3d10
Reviewed-on: https://go-review.googlesource.com/c/go/+/588675
Reviewed-by: Austin Clements <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: prattmic

src/os/exec.go

@@ -17,27 +17,234 @@ import (
 // ErrProcessDone indicates a [Process] has finished.
 var ErrProcessDone = errors.New("os: process already finished")
 
+type processMode uint8
+
+const (
+	// modePID means that Process operations such use the raw PID from the
+	// Pid field. handle is not used.
+	//
+	// This may be due to the host not supporting handles, or because
+	// Process was created as a literal, leaving handle unset.
+	//
+	// This must be the zero value so Process literals get modePID.
+	modePID processMode = iota
+
+	// modeHandle means that Process operations use handle, which is
+	// initialized with an OS process handle.
+	//
+	// Note that Release and Wait will deactivate and eventually close the
+	// handle, so acquire may fail, indicating the reason.
+	modeHandle
+)
+
+type processStatus uint64
+
+const (
+	// PID/handle OK to use.
+	statusOK processStatus = 0
+
+	// statusDone indicates that the PID/handle should not be used because
+	// the process is done (has been successfully Wait'd on).
+	statusDone processStatus = 1 << 62
+
+	// statusReleased indicates that the PID/handle should not be used
+	// because the process is released.
+	statusReleased processStatus = 1 << 63
+
+	processStatusMask = 0x3 << 62
+)
+
 // Process stores the information about a process created by [StartProcess].
 type Process struct {
-	Pid    int
-	handle atomic.Uintptr // Process handle for Windows, pidfd for Linux.
-	isdone atomic.Bool    // process has been successfully waited on
-	sigMu  sync.RWMutex   // avoid race between wait and signal
+	Pid int
+
+	mode processMode
+
+	// State contains the atomic process state.
+	//
+	// In modePID, this consists only of the processStatus fields, which
+	// indicate if the process is done/released.
+	//
+	// In modeHandle, the lower bits also contain a reference count for the
+	// handle field.
+	//
+	// The Process itself initially holds 1 persistent reference. Any
+	// operation that uses the handle with a system call temporarily holds
+	// an additional transient reference. This prevents the handle from
+	// being closed prematurely, which could result in the OS allocating a
+	// different handle with the same value, leading to Process' methods
+	// operating on the wrong process.
+	//
+	// Release and Wait both drop the Process' persistent reference, but
+	// other concurrent references may delay actually closing the handle
+	// because they hold a transient reference.
+	//
+	// Regardless, we want new method calls to immediately treat the handle
+	// as unavailable after Release or Wait to avoid extending this delay.
+	// This is achieved by setting either processStatus flag when the
+	// Process' persistent reference is dropped. The only difference in the
+	// flags is the reason the handle is unavailable, which affects the
+	// errors returned by concurrent calls.
+	state atomic.Uint64
+
+	// Used only in modePID.
+	sigMu sync.RWMutex // avoid race between wait and signal
+
+	// handle is the OS handle for process actions, used only in
+	// modeHandle.
+	//
+	// handle must be accessed only via the handleTransientAcquire method
+	// (or during closeHandle), not directly! handle is immutable.
+	//
+	// On Windows, it is a handle from OpenProcess.
+	// On Linux, it is a pidfd.
+	// It is unused on other GOOSes.
+	handle uintptr
+}
+
+func newPIDProcess(pid int) *Process {
+	p := &Process{
+		Pid:  pid,
+		mode: modePID,
+	}
+	runtime.SetFinalizer(p, (*Process).Release)
+	return p
+}
+
+func newHandleProcess(pid int, handle uintptr) *Process {
+	p := &Process{
+		Pid:    pid,
+		mode:   modeHandle,
+		handle: handle,
+	}
+	p.state.Store(1) // 1 persistent reference
+	runtime.SetFinalizer(p, (*Process).Release)
+	return p
 }
 
-func newProcess(pid int, handle uintptr) *Process {
-	p := &Process{Pid: pid}
-	p.handle.Store(handle)
+func newDoneProcess(pid int) *Process {
+	p := &Process{
+		Pid:  pid,
+		mode: modeHandle,
+		// N.B Since we set statusDone, handle will never actually be
+		// used, so its value doesn't matter.
+	}
+	p.state.Store(uint64(statusDone)) // No persistent reference, as there is no handle.
 	runtime.SetFinalizer(p, (*Process).Release)
 	return p
 }
 
-func (p *Process) setDone() {
-	p.isdone.Store(true)
+func (p *Process) handleTransientAcquire() (uintptr, processStatus) {
+	if p.mode != modeHandle {
+		panic("handleTransientAcquire called in invalid mode")
+	}
+
+	for {
+		refs := p.state.Load()
+		if refs&processStatusMask != 0 {
+			return 0, processStatus(refs & processStatusMask)
+		}
+		new := refs + 1
+		if !p.state.CompareAndSwap(refs, new) {
+			continue
+		}
+		return p.handle, statusOK
+	}
+}
+
+func (p *Process) handleTransientRelease() {
+	if p.mode != modeHandle {
+		panic("handleTransientRelease called in invalid mode")
+	}
+
+	for {
+		state := p.state.Load()
+		refs := state &^ processStatusMask
+		status := processStatus(state & processStatusMask)
+		if refs == 0 {
+			// This should never happen because
+			// handleTransientRelease is always paired with
+			// handleTransientAcquire.
+			panic("release of handle with refcount 0")
+		}
+		if refs == 1 && status == statusOK {
+			// Process holds a persistent reference and always sets
+			// a status when releasing that reference
+			// (handlePersistentRelease). Thus something has gone
+			// wrong if this is the last release but a status has
+			// not always been set.
+			panic("final release of handle without processStatus")
+		}
+		new := state - 1
+		if !p.state.CompareAndSwap(state, new) {
+			continue
+		}
+		if new&^processStatusMask == 0 {
+			p.closeHandle()
+		}
+		return
+	}
 }
 
-func (p *Process) done() bool {
-	return p.isdone.Load()
+// Drop the Process' persistent reference on the handle, deactivating future
+// Wait/Signal calls with the passed reason.
+//
+// Returns the status prior to this call. If this is not statusOK, then the
+// reference was not dropped or status changed.
+func (p *Process) handlePersistentRelease(reason processStatus) processStatus {
+	if p.mode != modeHandle {
+		panic("handlePersistentRelease called in invalid mode")
+	}
+
+	for {
+		refs := p.state.Load()
+		status := processStatus(refs & processStatusMask)
+		if status != statusOK {
+			// Both Release and successful Wait will drop the
+			// Process' persistent reference on the handle. We
+			// can't allow concurrent calls to drop the reference
+			// twice, so we use the status as a guard to ensure the
+			// reference is dropped exactly once.
+			return status
+		}
+		if refs == 0 {
+			// This should never happen because dropping the
+			// persistent reference always sets a status.
+			panic("release of handle with refcount 0")
+		}
+		new := (refs - 1) | uint64(reason)
+		if !p.state.CompareAndSwap(refs, new) {
+			continue
+		}
+		if new&^processStatusMask == 0 {
+			p.closeHandle()
+		}
+		return status
+	}
+}
+
+func (p *Process) pidStatus() processStatus {
+	if p.mode != modePID {
+		panic("pidStatus called in invalid mode")
+	}
+
+	return processStatus(p.state.Load())
+}
+
+func (p *Process) pidDeactivate(reason processStatus) {
+	if p.mode != modePID {
+		panic("pidDeactivate called in invalid mode")
+	}
+
+	// Both Release and successful Wait will deactivate the PID. Only one
+	// of those should win, so nothing left to do here if the compare
+	// fails.
+	//
+	// N.B. This means that results can be inconsistent. e.g., with a
+	// racing Release and Wait, Wait may successfully wait on the process,
+	// returning the wait status, while future calls error with "process
+	// released" rather than "process done".
+	p.state.CompareAndSwap(0, uint64(reason))
 }
 
 // ProcAttr holds the attributes that will be applied to a new process
@@ -116,6 +323,22 @@ func StartProcess(name string, argv []string, attr *ProcAttr) (*Process, error)
 // rendering it unusable in the future.
 // Release only needs to be called if [Process.Wait] is not.
 func (p *Process) Release() error {
+	// Note to future authors: the Release API is cursed.
+	//
+	// On Unix and Plan 9, Release sets p.Pid = -1. This is the only part of the
+	// Process API that is not thread-safe, but it can't be changed now.
+	//
+	// On Windows, Release does _not_ modify p.Pid.
+	//
+	// On Windows, Wait calls Release after successfully waiting to
+	// proactively clean up resources.
+	//
+	// On Unix and Plan 9, Wait also proactively cleans up resources, but
+	// can not call Release, as Wait does not set p.Pid = -1.
+	//
+	// On Unix and Plan 9, calling Release a second time has no effect.
+	//
+	// On Windows, calling Release a second time returns EINVAL.
 	return p.release()
 }
 

src/os/exec_linux.go

@@ -0,0 +1,13 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package os
+
+import (
+	"syscall"
+)
+
+func (p *Process) closeHandle() {
+	syscall.Close(int(p.handle))
+}

src/os/exec_nohandle.go

@@ -0,0 +1,9 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !linux && !windows
+
+package os
+
+func (p *Process) closeHandle() {}

src/os/exec_plan9.go

@@ -32,12 +32,12 @@ func startProcess(name string, argv []string, attr *ProcAttr) (p *Process, err e
 		sysattr.Files = append(sysattr.Files, f.Fd())
 	}
 
-	pid, h, e := syscall.StartProcess(name, argv, sysattr)
+	pid, _, e := syscall.StartProcess(name, argv, sysattr)
 	if e != nil {
 		return nil, &PathError{Op: "fork/exec", Path: name, Err: e}
 	}
 
-	return newProcess(pid, h), nil
+	return newPIDProcess(pid), nil
 }
 
 func (p *Process) writeProcFile(file string, data string) error {
@@ -51,9 +51,13 @@ func (p *Process) writeProcFile(file string, data string) error {
 }
 
 func (p *Process) signal(sig Signal) error {
-	if p.done() {
+	switch p.pidStatus() {
+	case statusDone:
 		return ErrProcessDone
+	case statusReleased:
+		return syscall.ENOENT
 	}
+
 	if e := p.writeProcFile("note", sig.String()); e != nil {
 		return NewSyscallError("signal", e)
 	}
@@ -67,15 +71,17 @@ func (p *Process) kill() error {
 func (p *Process) wait() (ps *ProcessState, err error) {
 	var waitmsg syscall.Waitmsg
 
-	if p.Pid == -1 {
+	switch p.pidStatus() {
+	case statusReleased:
 		return nil, ErrInvalid
 	}
+
 	err = syscall.WaitProcess(p.Pid, &waitmsg)
 	if err != nil {
 		return nil, NewSyscallError("wait", err)
 	}
 
-	p.setDone()
+	p.pidDeactivate(statusDone)
 	ps = &ProcessState{
 		pid:    waitmsg.Pid,
 		status: &waitmsg,
@@ -84,16 +90,19 @@ func (p *Process) wait() (ps *ProcessState, err error) {
 }
 
 func (p *Process) release() error {
-	// NOOP for Plan 9.
 	p.Pid = -1
+
+	// Just mark the PID unusable.
+	p.pidDeactivate(statusReleased)
+
 	// no need for a finalizer anymore
 	runtime.SetFinalizer(p, nil)
 	return nil
 }
 
 func findProcess(pid int) (p *Process, err error) {
 	// NOOP for Plan 9.
-	return newProcess(pid, 0), nil
+	return newPIDProcess(pid), nil
 }
 
 // ProcessState stores information about a process, as reported by Wait.

src/os/exec_posix.go

@@ -13,10 +13,6 @@ import (
 	"syscall"
 )
 
-// unsetHandle is a value for Process.handle used when the handle is not set.
-// Same as syscall.InvalidHandle for Windows.
-const unsetHandle = ^uintptr(0)
-
 // The only signal values guaranteed to be present in the os package on all
 // systems are os.Interrupt (send the process an interrupt) and os.Kill (force
 // the process to exit). On Windows, sending os.Interrupt to a process with
@@ -66,10 +62,14 @@ func startProcess(name string, argv []string, attr *ProcAttr) (p *Process, err e
 
 	// For Windows, syscall.StartProcess above already returned a process handle.
 	if runtime.GOOS != "windows" {
-		h = getPidfd(sysattr.Sys)
+		var ok bool
+		h, ok = getPidfd(sysattr.Sys)
+		if !ok {
+			return newPIDProcess(pid), nil
+		}
 	}
 
-	return newProcess(pid, h), nil
+	return newHandleProcess(pid, h), nil
 }
 
 func (p *Process) kill() error {