Support timeout for commands

utils.go

@@ -2,6 +2,7 @@ package zfs
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -10,18 +11,57 @@ 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
 	Stdout  io.Writer
 }
 
 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