bosun.org@v0.0.0-20210513094433-e25bc3e69a1f/util/command.go (about)

     1  package util
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"errors"
     7  	"io"
     8  	"os"
     9  	"os/exec"
    10  	"strings"
    11  	"time"
    12  
    13  	"bosun.org/slog"
    14  )
    15  
    16  var (
    17  	// ErrPath is returned by Command if the program is not in the PATH.
    18  	ErrPath = errors.New("program not in PATH")
    19  	// ErrTimeout is returned by Command if the program timed out.
    20  	ErrTimeout = errors.New("program killed after timeout")
    21  
    22  	// Debug enables debug logging.
    23  	Debug = false
    24  )
    25  
    26  // Command executes the named program with the given arguments. If it does not
    27  // exit within timeout, it is sent SIGINT (if supported by Go). After
    28  // another timeout, it is killed.
    29  func Command(timeout time.Duration, stdin io.Reader, name string, arg ...string) (io.Reader, error) {
    30  	if _, err := exec.LookPath(name); err != nil {
    31  		return nil, ErrPath
    32  	}
    33  	if Debug {
    34  		slog.Infof("executing command: %v %v", name, arg)
    35  	}
    36  	c := exec.Command(name, arg...)
    37  	b := &bytes.Buffer{}
    38  	c.Stdout = b
    39  	c.Stdin = stdin
    40  	if err := c.Start(); err != nil {
    41  		return nil, err
    42  	}
    43  	timedOut := false
    44  	intTimer := time.AfterFunc(timeout, func() {
    45  		slog.Errorf("Process taking too long. Interrupting: %s %s", name, strings.Join(arg, " "))
    46  		c.Process.Signal(os.Interrupt)
    47  		timedOut = true
    48  	})
    49  	killTimer := time.AfterFunc(timeout*2, func() {
    50  		slog.Errorf("Process taking too long. Killing: %s %s", name, strings.Join(arg, " "))
    51  		c.Process.Signal(os.Kill)
    52  		timedOut = true
    53  	})
    54  	err := c.Wait()
    55  	intTimer.Stop()
    56  	killTimer.Stop()
    57  	if timedOut {
    58  		return nil, ErrTimeout
    59  	}
    60  	return b, err
    61  }
    62  
    63  // ReadCommand runs command name with args and calls line for each line from its
    64  // stdout. Command is interrupted (if supported by Go) after 10 seconds and
    65  // killed after 20 seconds.
    66  func ReadCommand(line func(string) error, name string, arg ...string) error {
    67  	return ReadCommandTimeout(time.Second*10, line, nil, name, arg...)
    68  }
    69  
    70  // ReadCommandTimeout is the same as ReadCommand with a specifiable timeout.
    71  // It can also take a []byte as input (useful for chaining commands).
    72  func ReadCommandTimeout(timeout time.Duration, line func(string) error, stdin io.Reader, name string, arg ...string) error {
    73  	b, err := Command(timeout, stdin, name, arg...)
    74  	if err != nil {
    75  		return err
    76  	}
    77  	scanner := bufio.NewScanner(b)
    78  	for scanner.Scan() {
    79  		if err := line(scanner.Text()); err != nil {
    80  			return err
    81  		}
    82  	}
    83  	if err := scanner.Err(); err != nil {
    84  		slog.Infof("%v: %v\n", name, err)
    85  	}
    86  	return nil
    87  }