From 9a21256b6f4cbef7f5bc5cab1cfd2e249e145836 Mon Sep 17 00:00:00 2001
From: antonag32 <antonio@vauxoo.com>
Date: Mon, 29 Jul 2024 15:22:33 -0600
Subject: [PATCH] Support timeout for commands

Commands can now be stopped after the specified timeout. The default
behaviour of no timeouts is kept the same. Users wanting to opt into
this feature must set a new DefaultRunner with properly set values.

Related to #70.
---
 utils.go | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/utils.go b/utils.go
index b69942b..37eb0dc 100644
--- a/utils.go
+++ b/utils.go
@@ -2,6 +2,7 @@ package zfs
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -10,10 +11,32 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync/atomic"
+	"syscall"
+	"time"
 
 	"github.com/google/uuid"
 )
 
+type Runner struct {
+	Timeout time.Duration
+	Grace   time.Duration
+}
+
+var defaultRunner atomic.Value
+
+func init() {
+	defaultRunner.Store(&Runner{})
+}
+
+func Default() *Runner {
+	return defaultRunner.Load().(*Runner)
+}
+
+func SetDefault(runner *Runner) {
+	defaultRunner.Store(runner)
+}
+
 type command struct {
 	Command string
 	Stdin   io.Reader
@@ -21,7 +44,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