bitbucket.org/ai69/amoy@v0.2.3/exec_start.go

package amoy

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sync"
	"time"

	"github.com/1set/gut/ystring"
)

// Command represents an external command being running or exited.
type Command struct {
	sync.Mutex
	cmd     *exec.Cmd
	bufOut  bytes.Buffer
	bufErr  bytes.Buffer
	err     error
	pid     int
	startAt time.Time
	stopAt  time.Time
	done    chan struct{}
}

// Done returns a channel that's closed when the command exits.
func (c *Command) Done() <-chan struct{} {
	return c.done
}

// Stdout returns a slice holding the content of the standard output of the command.
func (c *Command) Stdout() []byte {
	return c.bufOut.Bytes()
}

// Stderr returns a slice holding the content of the standard error of the command.
func (c *Command) Stderr() []byte {
	return c.bufErr.Bytes()
}

// Error returns the error after the command exits if it exists.
func (c *Command) Error() error {
	return c.err
}

// ProcessID indicates PID of the command.
func (c *Command) ProcessID() int {
	return c.pid
}

// StartedAt indicates the moment that the command started.
func (c *Command) StartedAt() time.Time {
	return c.startAt
}

// StoppedAt indicates the moment that the command stopped or zero if it's still running.
func (c *Command) StoppedAt() time.Time {
	return c.stopAt
}

// Exited reports whether the command has exited.
func (c *Command) Exited() bool {
	if c.cmd.ProcessState != nil {
		return c.cmd.ProcessState.Exited()
	}
	return false
}

// Kill causes the command to exit immediately, and does not wait until it has actually exited.
func (c *Command) Kill() error {
	// TODO: use sync.Once to wrapper it, Status() func to get the status of the command: running, exited, crashed, killed
	if c.cmd.Process != nil {
		return c.cmd.Process.Kill()
	}
	return errMissingProcInfo
}

// CommandOptions represents custom options to execute external command.
type CommandOptions struct {
	// WorkDir is the working directory of the command.
	WorkDir string
	// EnvVar appends the environment variables of the command.
	EnvVar map[string]string
	// Stdout is the writer to write standard output to. Use os.Pipe() to create a pipe for this, if you want to handle the result synchronously.
	Stdout io.Writer
	// Stderr is the writer to write standard error to. Use os.Pipe() to create a pipe for this, if you want to handle the result synchronously.
	Stderr io.Writer
	// DisableResult indicates whether to disable the buffered result of the command, especially important for long-running commands.
	DisableResult bool
	// Stdin is the input to the command's standard input.
	Stdin io.Reader
	// TODO: not implemented
	Timeout time.Time
}

// StartSimpleCommand starts the specified command but does not wait for it to complete, and simultaneously writes its standard output and standard error to given writers.
func StartSimpleCommand(command string, writerStdout, writerStderr io.Writer) (*Command, error) {
	return StartCommand(command, &CommandOptions{
		Stdout: writerStdout,
		Stderr: writerStderr,
	})
}

// StartCommand starts the specified command with given options but does not wait for it to complete.
func StartCommand(command string, opts ...*CommandOptions) (*Command, error) {
	cmd, err := parseSingleRawCommand(command)
	if err != nil {
		return nil, err
	}

	var (
		// handler for cmd
		handler = Command{
			cmd:  cmd,
			done: make(chan struct{}),
		}

		// pipes that will be connected
		pipeStdout, _ = cmd.StdoutPipe()
		pipeStderr, _ = cmd.StderrPipe()

		// combined writers that duplicate given writers and buffers
		comStdout io.Writer
		comStderr io.Writer
	)

	// use the first option if it exists
	if len(opts) > 0 {
		opt := opts[0]

		// for env var
		if len(opt.EnvVar) > 0 {
			cmd.Env = os.Environ()
			for k, v := range opt.EnvVar {
				cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v))
			}
		}

		// for standard input stream
		if opt.Stdin != nil {
			cmd.Stdin = opt.Stdin
		}

		// set working directory
		if ystring.IsNotBlank(opt.WorkDir) {
			cmd.Dir = opt.WorkDir
		}

		// handle standard output
		if opt.Stdout != nil {
			// if the given writer exists
			if opt.DisableResult {
				// if the buffered result is disabled, just redirects to the given writer
				comStdout = opt.Stdout
			} else {
				// otherwise, writes both to buffered result and given writer
				comStdout = io.MultiWriter(opt.Stdout, &handler.bufOut)
			}
		} else {
			// if the given writer doesn't exist
			if opt.DisableResult {
				// if the buffered result is disabled, writes to /dev/null
				comStdout = DiscardWriter
			} else {
				// otherwise, writes to buffered result
				comStdout = &handler.bufOut
			}
		}

		// handle standard error
		if opt.Stderr != nil {
			// if the given writer exists
			if opt.DisableResult {
				// if the buffered result is disabled, just redirects to the given writer
				comStderr = opt.Stderr
			} else {
				// otherwise, writes both to buffered result and given writer
				comStderr = io.MultiWriter(opt.Stderr, &handler.bufErr)
			}
		} else {
			// if the given writer doesn't exist
			if opt.DisableResult {
				// if the buffered result is disabled, writes to /dev/null
				comStderr = DiscardWriter
			} else {
				// otherwise, writes to buffered result
				comStderr = &handler.bufErr
			}
		}
	}

	// let's start!
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("fail to start exec: %w", err)
	}

	if cmd.Process != nil {
		handler.startAt = time.Now()
		handler.pid = cmd.Process.Pid
	}

	var (
		wg        sync.WaitGroup
		errStdout error
		errStderr error
	)
	wg.Add(1)
	go func() {
		defer wg.Done()
		_, errStdout = io.Copy(comStdout, pipeStdout)
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		_, errStderr = io.Copy(comStderr, pipeStderr)
	}()

	go func() {
		defer func() {
			close(handler.done)
		}()

		wg.Wait()
		handler.stopAt = time.Now()

		if err := cmd.Wait(); err != nil {
			handler.err = fmt.Errorf("fail to run exec: %w", err)
			return
		}
		if errStdout != nil {
			handler.err = fmt.Errorf("fail to capture stdout: %w", errStdout)
			return
		}
		if errStderr != nil {
			handler.err = fmt.Errorf("fail to capture stderr: %w", errStderr)
			return
		}
	}()

	return &handler, nil
}