diff --git a/utils.go b/utils.go index b69942b..8e49be3 100644 --- a/utils.go +++ b/utils.go @@ -2,6 +2,7 @@ package zfs import ( "bytes" + "context" "errors" "fmt" "io" @@ -10,10 +11,37 @@ import ( "runtime" "strconv" "strings" + "sync/atomic" + "syscall" + "time" "github.com/google/uuid" ) +// Runner specifies the parameters used when executing ZFS commands. +type Runner struct { + // Timeout specifies how long to wait before sending a SIGTERM signal to the running process. + Timeout time.Duration + + // Grace specifies the time waited after signaling the running process with SIGTERM before it is forcefully + // killed with SIGKILL. + Grace time.Duration +} + +var defaultRunner atomic.Value + +func init() { + defaultRunner.Store(&Runner{}) +} + +func Default() *Runner { + return defaultRunner.Load().(*Runner) //nolint: forcetypeassert // Impossible for it to be anything else. +} + +func SetRunner(runner *Runner) { + defaultRunner.Store(runner) +} + type command struct { Command string Stdin io.Reader @@ -21,7 +49,19 @@ type command struct { } func (c *command) Run(arg ...string) ([][]string, error) { - cmd := exec.Command(c.Command, arg...) + var cmd *exec.Cmd + if Default().Timeout == 0 { + cmd = exec.Command(c.Command, arg...) + } else { + ctx, cancel := context.WithTimeout(context.Background(), Default().Timeout) + defer cancel() + + cmd = exec.CommandContext(ctx, c.Command, arg...) + cmd.Cancel = func() error { + return cmd.Process.Signal(syscall.SIGTERM) + } + cmd.WaitDelay = Default().Grace + } var stdout, stderr bytes.Buffer