mirror of https://github.com/coder/coder.git
fix: cleanup reaper implementation (#2563)
- Clean up the agent/reaper API to be a more isolated and reusable package.
This commit is contained in:
parent
0585372170
commit
ee5918217b
|
@ -0,0 +1,28 @@
|
|||
package reaper
|
||||
|
||||
import "github.com/hashicorp/go-reap"
|
||||
|
||||
type Option func(o *options)
|
||||
|
||||
// WithExecArgs specifies the exec arguments for the fork exec call.
|
||||
// By default the same arguments as the parent are used as dictated by
|
||||
// os.Args. Since ForkReap calls a fork-exec it is the responsibility of
|
||||
// the caller to avoid fork-bombing oneself.
|
||||
func WithExecArgs(args ...string) Option {
|
||||
return func(o *options) {
|
||||
o.ExecArgs = args
|
||||
}
|
||||
}
|
||||
|
||||
// WithPIDCallback sets the channel that reaped child process PIDs are pushed
|
||||
// onto.
|
||||
func WithPIDCallback(ch reap.PidCh) Option {
|
||||
return func(o *options) {
|
||||
o.PIDs = ch
|
||||
}
|
||||
}
|
||||
|
||||
type options struct {
|
||||
ExecArgs []string
|
||||
PIDs reap.PidCh
|
||||
}
|
|
@ -2,18 +2,11 @@
|
|||
|
||||
package reaper
|
||||
|
||||
import "github.com/hashicorp/go-reap"
|
||||
|
||||
// IsChild returns true if we're the forked process.
|
||||
func IsChild() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// IsInitProcess returns true if the current process's PID is 1.
|
||||
func IsInitProcess() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func ForkReap(_ reap.PidCh) error {
|
||||
func ForkReap(opt ...Option) error {
|
||||
return nil
|
||||
}
|
||||
|
|
|
@ -24,17 +24,16 @@ func TestReap(t *testing.T) {
|
|||
t.Skip("Detected CI, skipping reaper tests")
|
||||
}
|
||||
|
||||
// Because we're forkexecing these tests will try to run twice...
|
||||
if reaper.IsChild() {
|
||||
t.Skip("I'm a child!")
|
||||
}
|
||||
|
||||
// OK checks that's the reaper is successfully reaping
|
||||
// exited processes and passing the PIDs through the shared
|
||||
// channel.
|
||||
t.Run("OK", func(t *testing.T) {
|
||||
pids := make(reap.PidCh, 1)
|
||||
err := reaper.ForkReap(pids)
|
||||
err := reaper.ForkReap(
|
||||
reaper.WithPIDCallback(pids),
|
||||
// Provide some argument that immediately exits.
|
||||
reaper.WithExecArgs("/bin/sh", "-c", "exit 0"),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := exec.Command("tail", "-f", "/dev/null")
|
||||
|
|
|
@ -3,7 +3,6 @@
|
|||
package reaper
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"syscall"
|
||||
|
||||
|
@ -11,17 +10,6 @@ import (
|
|||
"golang.org/x/xerrors"
|
||||
)
|
||||
|
||||
// agentEnvMark is a simple environment variable that we use as a marker
|
||||
// to indicated that the process is a child as opposed to the reaper.
|
||||
// Since we are forkexec'ing we need to be able to differentiate between
|
||||
// the two to avoid fork bombing ourselves.
|
||||
const agentEnvMark = "CODER_DO_NOT_REAP"
|
||||
|
||||
// IsChild returns true if we're the forked process.
|
||||
func IsChild() bool {
|
||||
return os.Getenv(agentEnvMark) != ""
|
||||
}
|
||||
|
||||
// IsInitProcess returns true if the current process's PID is 1.
|
||||
func IsInitProcess() bool {
|
||||
return os.Getpid() == 1
|
||||
|
@ -33,19 +21,16 @@ func IsInitProcess() bool {
|
|||
// the reaper and an exec.Command waiting for its process to complete.
|
||||
// The provided 'pids' channel may be nil if the caller does not care about the
|
||||
// reaped children PIDs.
|
||||
func ForkReap(pids reap.PidCh) error {
|
||||
// Check if the process is the parent or the child.
|
||||
// If it's the child we want to skip attempting to reap.
|
||||
if IsChild() {
|
||||
return nil
|
||||
func ForkReap(opt ...Option) error {
|
||||
opts := &options{
|
||||
ExecArgs: os.Args,
|
||||
}
|
||||
|
||||
go reap.ReapChildren(pids, nil, nil, nil)
|
||||
for _, o := range opt {
|
||||
o(opts)
|
||||
}
|
||||
|
||||
args := os.Args
|
||||
// This is simply done to help identify the real agent process
|
||||
// when viewing in something like 'ps'.
|
||||
args = append(args, "#Agent")
|
||||
go reap.ReapChildren(opts.PIDs, nil, nil, nil)
|
||||
|
||||
pwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
|
@ -54,8 +39,7 @@ func ForkReap(pids reap.PidCh) error {
|
|||
|
||||
pattrs := &syscall.ProcAttr{
|
||||
Dir: pwd,
|
||||
// Add our marker for identifying the child process.
|
||||
Env: append(os.Environ(), fmt.Sprintf("%s=true", agentEnvMark)),
|
||||
Env: os.Environ(),
|
||||
Sys: &syscall.SysProcAttr{
|
||||
Setsid: true,
|
||||
},
|
||||
|
@ -67,7 +51,7 @@ func ForkReap(pids reap.PidCh) error {
|
|||
}
|
||||
|
||||
//#nosec G204
|
||||
pid, _ := syscall.ForkExec(args[0], args, pattrs)
|
||||
pid, _ := syscall.ForkExec(opts.ExecArgs[0], opts.ExecArgs, pattrs)
|
||||
|
||||
var wstatus syscall.WaitStatus
|
||||
_, err = syscall.Wait4(pid, &wstatus, 0, nil)
|
||||
|
|
|
@ -32,6 +32,7 @@ func workspaceAgent() *cobra.Command {
|
|||
auth string
|
||||
pprofEnabled bool
|
||||
pprofAddress string
|
||||
noReap bool
|
||||
)
|
||||
cmd := &cobra.Command{
|
||||
Use: "agent",
|
||||
|
@ -58,9 +59,12 @@ func workspaceAgent() *cobra.Command {
|
|||
|
||||
// Spawn a reaper so that we don't accumulate a ton
|
||||
// of zombie processes.
|
||||
if reaper.IsInitProcess() && !reaper.IsChild() && isLinux {
|
||||
if reaper.IsInitProcess() && !noReap && isLinux {
|
||||
logger.Info(cmd.Context(), "spawning reaper process")
|
||||
err := reaper.ForkReap(nil)
|
||||
// Do not start a reaper on the child process. It's important
|
||||
// to do this else we fork bomb ourselves.
|
||||
args := append(os.Args, "--no-reap")
|
||||
err := reaper.ForkReap(reaper.WithExecArgs(args...))
|
||||
if err != nil {
|
||||
logger.Error(cmd.Context(), "failed to reap", slog.Error(err))
|
||||
return xerrors.Errorf("fork reap: %w", err)
|
||||
|
@ -182,6 +186,7 @@ func workspaceAgent() *cobra.Command {
|
|||
|
||||
cliflag.StringVarP(cmd.Flags(), &auth, "auth", "", "CODER_AGENT_AUTH", "token", "Specify the authentication type to use for the agent")
|
||||
cliflag.BoolVarP(cmd.Flags(), &pprofEnabled, "pprof-enable", "", "CODER_AGENT_PPROF_ENABLE", false, "Enable serving pprof metrics on the address defined by --pprof-address.")
|
||||
cliflag.BoolVarP(cmd.Flags(), &noReap, "no-reap", "", "", false, "Do not start a process reaper.")
|
||||
cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
|
||||
return cmd
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue