github.com/ronaksoft/rony@v0.16.26-0.20230807065236-1743dbfe6959/tools/commands.go (about)

     1  package tools
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"os/exec"
     7  )
     8  
     9  var (
    10  	ErrNoCommands = errors.New("pipe: there were no commands provided")
    11  )
    12  
    13  func NewPipeCommands(cmds ...*exec.Cmd) (*cmdPipe, error) {
    14  	if len(cmds) == 0 {
    15  		return nil, ErrNoCommands
    16  	}
    17  
    18  	// Link each command's Stdout w/ the next command's
    19  	// Stdin
    20  	writeClosers := make([]io.WriteCloser, len(cmds))
    21  	for i := range cmds {
    22  		if i == 0 {
    23  			continue
    24  		}
    25  		r, w := io.Pipe()
    26  		writeClosers[i-1] = w
    27  		cmds[i-1].Stdout = w
    28  		cmds[i].Stdin = r
    29  	}
    30  
    31  	return &cmdPipe{
    32  		writeClosers: writeClosers,
    33  		commands:     cmds,
    34  	}, nil
    35  }
    36  
    37  type cmdPipe struct {
    38  	// Change these *BEFORE* calling start
    39  	Stdin        io.Reader
    40  	Stdout       io.Writer
    41  	writeClosers []io.WriteCloser
    42  	commands     []*exec.Cmd
    43  }
    44  
    45  func (c *cmdPipe) Start() error {
    46  	c.commands[0].Stdin = c.Stdin
    47  	c.commands[len(c.commands)-1].Stdout = c.Stdout
    48  	for i := range c.commands {
    49  		if err := c.commands[i].Start(); err != nil {
    50  			return err
    51  		}
    52  	}
    53  
    54  	return nil
    55  }
    56  
    57  func (c *cmdPipe) Wait() error {
    58  	for i := range c.commands {
    59  		if err := c.commands[i].Wait(); err != nil {
    60  			return err
    61  		}
    62  		if i != len(c.commands)-1 {
    63  			if err := c.writeClosers[i].Close(); err != nil {
    64  				return err
    65  			}
    66  		}
    67  	}
    68  
    69  	return nil
    70  }