ntool/ncmd/cmd.go

package ncmd

import (
	"errors"
	"fmt"
	"git.noahlan.cn/noahlan/ntool/ndef"
	"git.noahlan.cn/noahlan/ntool/nlog"
	"io"
	"os/exec"
	"sync"
	"syscall"
	"time"
)

var (
	ErrNotStarted   = errors.New("command not running")
	ErrMarshalerNil = errors.New("marshaller is nil but data not []byte type")
)

type (
	Options struct {
		Marshaler ndef.Marshaler // 序列化

		// 缓存 STDOUT 和 STDERR，最终写入 Status 内
		// 通过 Cmd.Status 获取
		Buffered bool

		// STDOUT 和 STDERR 写入到 Status.Stdout 中，类似 linux 中的 2>&1
		// 注：当Streaming开启时，此特性无效
		CombinedOutput bool

		// 流式处理，会将所有数据即时输出
		Streaming bool

		// LineBufferSize 设置 OutputStream 行缓冲区的大小，默认值 DEFAULT_LINE_BUFFER_SIZE
		// 但如果发生 ErrLineBufferOverflow 错误，请尝试增加大小
		LineBufferSize uint

		// DevMode 开发模式，打印log
		DevMode bool
	}
	Option func(opt *Options)
)

// Cmd 处理器
type Cmd struct {
	*Options

	// Runtime
	Cmd   *exec.Cmd      // CMD
	Name  string         // name of binary(command) to run
	Args  []string       // arguments
	Env   []string       // *optional environment
	Dir   string         // working dir
	stdIn io.WriteCloser // 标准输入通道

	stdoutStream *OutputStream // 标准输出通道
	stderrStream *OutputStream // 标准错误输出通道

	// flags
	started    bool          // cmd.Start called, no error
	stopped    bool          // cmd.Stop called
	done       bool          // run() done
	final      bool          // status finalized in Status
	doneChan   chan struct{} // closed when done running
	status     Status
	statusChan chan Status // nil until Start() called
	startTime  time.Time

	// messages
	Stdout    chan string   // Stdout sets streaming STDOUT if enabled, else nil (see Options).
	Stderr    chan string   // Stderr sets streaming STDERR if enabled, else nil (see Options).
	stdoutBuf *OutputBuffer // 标准输出缓冲区
	stderrBuf *OutputBuffer // 标准错误缓冲区
	chSend    chan []byte   // 发送的数据缓冲区

	*sync.Mutex
}

type Status struct {
	Cmd      string
	PID      int
	Complete bool     // 如果被主动停止或kill，该值为false
	Exit     int      // exitCode
	Error    error    // go error
	StartTs  int64    // 启动时间（纳秒） 0表示未启动(started)
	StopTs   int64    // 停止时间（纳秒） 0表示未启动(started)或未运行(running)
	Runtime  float64  // 运行时长
	Stdout   []string // buffered STDOUT
	Stderr   []string // buffered STDIN
}

func NewCmd(opts ...Option) *Cmd {
	tmp := &Cmd{
		Options: &Options{
			LineBufferSize: DEFAULT_LINE_BUFFER_SIZE,
			Buffered:       true,
		},
		// flags
		started:  false,
		stopped:  false,
		done:     false,
		final:    false,
		doneChan: make(chan struct{}),
		status: Status{
			Exit: -1,
		},
		statusChan: nil,
		// messages
		chSend: make(chan []byte, 1024),
		Mutex:  &sync.Mutex{},
	}

	for _, opt := range opts {
		opt(tmp.Options)
	}

	options := tmp.Options
	if options.Buffered {
		tmp.stdoutBuf = NewOutputBuffer()
		tmp.stderrBuf = NewOutputBuffer()
	}
	if options.CombinedOutput {
		tmp.stdoutBuf = NewOutputBuffer()
		tmp.stderrBuf = nil
	}
	if options.Streaming {
		tmp.Stdout = make(chan string, DEFAULT_STREAM_CHAN_SIZE)
		tmp.stdoutStream = NewOutputStream(tmp.Stdout)
		tmp.stdoutStream.SetLineBufferSize(int(options.LineBufferSize))

		tmp.Stderr = make(chan string, DEFAULT_STREAM_CHAN_SIZE)
		tmp.stderrStream = NewOutputStream(tmp.Stderr)
		tmp.stderrStream.SetLineBufferSize(int(options.LineBufferSize))
	}

	return tmp
}

func (c *Cmd) Start(name string, args ...string) <-chan Status {
	c.Name = name
	c.Args = args
	c.status.Cmd = name

	c.Lock()
	defer c.Unlock()

	if c.statusChan != nil {
		return c.statusChan
	}
	c.statusChan = make(chan Status, 1)

	go c.run()

	return c.statusChan
}

func (c *Cmd) Stop() error {
	c.Lock()
	defer c.Unlock()

	if c.statusChan == nil || !c.started {
		return ErrNotStarted
	}

	// 如果 done 为true，表示程序已经执行完毕，无需再Stop
	if c.done {
		return nil
	}

	c.stopped = true

	// Signal the process group (-pid), not just the process, so that the process
	// and all its children are signaled. Else, child procs can keep running and
	// keep the stdout/stderr fd open and cause cmd.Wait to hang.
	return terminateProcess(c.Cmd.Process.Pid)
}

func (c *Cmd) Status() Status {
	c.Lock()
	defer c.Unlock()

	// Return default status if cmd hasn't been started
	if c.statusChan == nil || !c.started {
		return c.status
	}
	if c.done {
		// No longer running
		if !c.final {
			if c.stdoutBuf != nil {
				c.status.Stdout = c.stdoutBuf.Lines()
				c.stdoutBuf = nil // release buffers

			}
			if c.stderrBuf != nil {
				c.status.Stderr = c.stderrBuf.Lines()
				c.stderrBuf = nil // release buffers
			}
			c.final = true
		}
	} else {
		// Still running
		c.status.Runtime = time.Now().Sub(c.startTime).Seconds()
		if c.stdoutBuf != nil {
			c.status.Stdout = c.stdoutBuf.Lines()

		}
		if c.stderrBuf != nil {
			c.status.Stderr = c.stderrBuf.Lines()
		}
	}

	return c.status
}

// Done 返回一个在Cmd停止运行时关闭的通道。
// 等待所有的 goroutine 完成执行
// 命令完成后调用 Status 来获取其最终状态
func (c *Cmd) Done() <-chan struct{} {
	return c.doneChan
}

func (c *Cmd) run() {
	defer func() {
		c.statusChan <- c.Status()
		close(c.doneChan)
	}()

	cmd := exec.Command(c.Name, c.Args...)
	c.Cmd = cmd

	setProcessGroupID(cmd)

	switch {
	case c.stdoutBuf != nil && c.stderrBuf != nil && c.stdoutStream != nil: // buffer and stream
		cmd.Stdout = io.MultiWriter(c.stdoutStream, c.stdoutBuf)
		cmd.Stderr = io.MultiWriter(c.stderrStream, c.stderrBuf)
	case c.stdoutBuf != nil && c.stderrBuf == nil && c.stdoutStream != nil: // combined buffer and stream
		cmd.Stdout = io.MultiWriter(c.stdoutStream, c.stdoutBuf)
		cmd.Stderr = io.MultiWriter(c.stderrStream, c.stdoutBuf)
	case c.stdoutBuf != nil && c.stderrBuf != nil: // buffer only
		cmd.Stdout = c.stdoutBuf
		cmd.Stderr = c.stderrBuf
	case c.stdoutBuf != nil && c.stderrBuf == nil: // buffer combining stderr into stdout
		cmd.Stdout = c.stdoutBuf
		cmd.Stderr = c.stdoutBuf
	case c.stdoutStream != nil: // stream only
		cmd.Stdout = c.stdoutStream
		cmd.Stderr = c.stderrStream
	default: // no output (cmd >/dev/null 2>&1)
		cmd.Stdout = nil
		cmd.Stderr = nil
	}
	c.stdIn, _ = cmd.StdinPipe()

	c.Env = cmd.Env
	c.Dir = cmd.Dir

	// Always close output streams. Do not do this after Wait because if Start
	// fails and we return without closing these, it could deadlock the caller
	// who's waiting for us to close them.
	if c.stdoutStream != nil {
		defer func() {
			c.stdoutStream.Flush()
			c.stderrStream.Flush()
			// exec.Cmd.Wait has already waited for all output:
			//   Otherwise, during the execution of the command a separate goroutine
			//   reads from the process over a pipe and delivers that data to the
			//   corresponding Writer. In this case, Wait does not complete until the
			//   goroutine reaches EOF or encounters an error.
			// from https://golang.org/pkg/os/exec/#Cmd
			close(c.Stdout)
			close(c.Stderr)
		}()
	}

	now := time.Now()
	if err := cmd.Start(); err != nil {
		c.Lock()
		c.status.Error = err
		c.status.StartTs = now.UnixNano()
		c.status.StopTs = time.Now().UnixNano()
		c.done = true
		c.Unlock()
		return
	}

	// Set initial status
	c.Lock()
	c.startTime = now // running
	c.status.PID = cmd.Process.Pid
	c.status.StartTs = now.UnixNano()
	c.started = true
	c.Unlock()

	// inner write
	go c.writeLoop()

	err := cmd.Wait()
	now = time.Now()

	// Get exit code of the command. According to the manual, Wait() returns:
	// "If the command fails to run or doesn't complete successfully, the error
	// is of type *ExitError. Other error types may be returned for I/O problems."
	exitCode := 0
	signaled := false
	if err != nil && fmt.Sprintf("%T", err) == "*exec.ExitError" {
		// This is the normal case which is not really an error. It's string
		// representation is only "*exec.ExitError". It only means the cmd
		// did not exit zero and caller should see ExitError.Stderr, which
		// we already have. So first we'll have this as the real/underlying
		// type, then discard err so status.Error doesn't contain a useless
		// "*exec.ExitError". With the real type we can get the non-zero
		// exit code and determine if the process was signaled, which yields
		// a more specific error message, so we set err again in that case.
		exiterr := err.(*exec.ExitError)
		err = nil
		if waitStatus, ok := exiterr.Sys().(syscall.WaitStatus); ok {
			exitCode = waitStatus.ExitStatus() // -1 if signaled
			if waitStatus.Signaled() {
				signaled = true
				err = errors.New(exiterr.Error()) // "signal: terminated"
			}
		}
	}
	// Set final status
	c.Lock()
	if !c.stopped && !signaled {
		c.status.Complete = true
	}
	c.status.Runtime = now.Sub(c.startTime).Seconds()
	c.status.StopTs = now.UnixNano()
	c.status.Exit = exitCode
	c.status.Error = err
	c.done = true
	c.Unlock()
}

// Send 发送指令
func (c *Cmd) Send(data any) (err error) {
	var bytes []byte
	if c.Marshaler != nil {
		if bs, ok := data.([]byte); ok {
			bytes = bs
		} else {
			return ErrMarshalerNil
		}
	}
	bytes, err = c.Marshaler.Marshal(data)
	if err != nil {
		nlog.Errorf("marshal data [%v] err %v", data, err)
		return err
	}
	c.chSend <- bytes

	return nil
}

func (c *Cmd) writeLoop() {
	defer func() {
		close(c.chSend)
		_ = c.stdIn.Close()
	}()

	for {
		select {
		case <-c.Done():
			return
		case data := <-c.chSend:
			// 实际写入数据
			if c.DevMode {
				nlog.Debugf("发送数据: [%s]", string(data))
			}

			data = append(data, '\n')
			if _, err := c.stdIn.Write(data); err != nil {
				return
			}
		}
	}
}