Add timeout when for Rsync

Author: bougou

rsync.go

@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"strconv"
 	"strings"
+	"time"
 )
 
 // Rsync is wrapper under rsync
@@ -15,6 +16,8 @@ type Rsync struct {
 	Destination string
 
 	cmd *exec.Cmd
+
+	timeout time.Duration
 }
 
 // RsyncOptions for rsync
@@ -217,6 +220,10 @@ func (r Rsync) StderrPipe() (io.ReadCloser, error) {
 	return r.cmd.StderrPipe()
 }
 
+func (r *Rsync) WithTimeout(timeout time.Duration) {
+	r.timeout = timeout
+}
+
 // Run start rsync task
 func (r Rsync) Run() error {
 	if !isExist(r.Destination) {
@@ -229,6 +236,10 @@ func (r Rsync) Run() error {
 		return err
 	}
 
+	if r.timeout != 0 {
+		return WaitTimeout(r.cmd, r.timeout)
+	}
+
 	return r.cmd.Wait()
 }
 
@@ -245,6 +256,7 @@ func NewRsync(source, destination string, options RsyncOptions) *Rsync {
 		Source:      source,
 		Destination: destination,
 		cmd:         exec.Command(binaryPath, arguments...),
+		timeout:     0 * time.Second,
 	}
 }
 

timeout_cmd.go

@@ -0,0 +1,66 @@
+package grsync
+
+import (
+	"errors"
+	"log"
+	"os/exec"
+	"syscall"
+	"time"
+)
+
+var ErrTimeout = errors.New("command timed out")
+
+// KillGrace is the amount of time we allow a process to shutdown before
+// sending a SIGKILL.
+const KillGrace = 5 * time.Second
+
+// WaitTimeout waits for the given command to finish with a timeout.
+// It assumes the command has already been started.
+// If the command times out, it attempts to kill the process and returns
+// a ErrTimeout error.
+func WaitTimeout(c *exec.Cmd, timeout time.Duration) error {
+	var kill *time.Timer
+
+	term := time.AfterFunc(timeout, func() {
+		err := c.Process.Signal(syscall.SIGTERM)
+		if err != nil {
+			log.Printf("Error terminating process: %s", err)
+			return
+		}
+
+		kill = time.AfterFunc(KillGrace, func() {
+			err := c.Process.Kill()
+			if err != nil {
+				log.Printf("Error killing process: %s", err)
+				return
+			}
+		})
+	})
+
+	err := c.Wait()
+
+	// Shutdown all timers (the kill timer and the term timer) before checking cmd err,
+	// otherwise there is no chance to turn off these timers that have not expired.
+	if kill != nil {
+		kill.Stop()
+	}
+	termSent := !term.Stop()
+	// For a timer created with AfterFunc(d, f), if t.Stop returns false, then
+	// the timer has already expired and the function f has been started in its own goroutine.
+	// So if termSent is true, it means the cmd does not finished before the term timer expired.
+
+	// Now, we can check cmd err.
+	// If the process exited without error treat it as success.
+	// This allows a process to do a clean shutdown on signal.
+	if err == nil {
+		return nil
+	}
+
+	// If SIGTERM was sent then treat any process error as a timeout.
+	if termSent {
+		return ErrTimeout
+	}
+
+	// Otherwise there was an cmd error unrelated to termination.
+	return err
+}