binlog-parser/vendor/b612.me/staros/process.go

package staros

import (
	"bytes"
	"context"
	"errors"
	"io"
	"os"
	"os/exec"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

//StarCmd Is Here

type StarCmd struct {
	CMD     *exec.Cmd
	outfile io.ReadCloser
	infile  io.WriteCloser
	errfile io.ReadCloser
	running int32
	//Store AlL of the Standed Outputs
	stdout []byte
	//Store All of the Standed Errors
	errout       []byte
	runerr       error
	exitcode     int
	stdoutBuf    *bytes.Buffer
	stderrBuf    *bytes.Buffer
	stdoutpoint  int
	stderrpoint  int
	lock         sync.Mutex
	prewrite     []string
	prewritetime time.Duration
	stopctxfunc  context.CancelFunc
	stopctx      context.Context
}

func Command(command string, args ...string) (*StarCmd, error) {
	var err error
	shell := new(StarCmd)
	shell.running = 0
	shell.prewritetime = time.Millisecond * 200
	shell.stdoutBuf = bytes.NewBuffer(make([]byte, 0))
	shell.stderrBuf = bytes.NewBuffer(make([]byte, 0))
	shell.stopctx, shell.stopctxfunc = context.WithCancel(context.Background())
	cmd := exec.Command(command, args...)
	shell.CMD = cmd
	shell.infile, err = shell.CMD.StdinPipe()
	if err != nil {
		return shell, err
	}
	shell.errfile, err = shell.CMD.StderrPipe()
	if err != nil {
		return shell, err
	}
	shell.outfile, err = shell.CMD.StdoutPipe()
	if err != nil {
		return shell, err
	}
	shell.runerr = nil
	shell.exitcode = -999
	return shell, nil
}
func CommandContext(ctx context.Context, command string, args ...string) (*StarCmd, error) {
	var err error
	shell := new(StarCmd)
	shell.running = 0
	shell.stdoutBuf = bytes.NewBuffer(make([]byte, 0))
	shell.stderrBuf = bytes.NewBuffer(make([]byte, 0))
	shell.prewritetime = time.Millisecond * 200
	shell.stopctx, shell.stopctxfunc = context.WithCancel(context.Background())
	cmd := exec.CommandContext(ctx, command, args...)
	shell.CMD = cmd
	shell.infile, err = shell.CMD.StdinPipe()
	if err != nil {
		return shell, err
	}
	shell.errfile, err = shell.CMD.StderrPipe()
	if err != nil {
		return shell, err
	}
	shell.outfile, err = shell.CMD.StdoutPipe()
	if err != nil {
		return shell, err
	}
	shell.runerr = nil
	shell.exitcode = -999
	return shell, nil
}

func (starcli *StarCmd) queryStdout(ctx context.Context) {
	for starcli.IsRunning() && starcli.CMD != nil {
		select {
		case <-ctx.Done():
			return
		default:
		}
		out := make([]byte, 65535)
		n, err := starcli.outfile.Read(out)
		if n != 0 {
			starcli.lock.Lock()
			starcli.stdoutBuf.Write(out[:n])
			starcli.lock.Unlock()
			for _, v := range out[:n] {
				starcli.stdout = append(starcli.stdout, v)
			}
		}
		if err != nil {
			if err == io.EOF {
				break
			} else {
				if !strings.Contains(err.Error(), "file already closed") {
					starcli.runerr = err
				}
				return
			}
		}
	}
}

func (starcli *StarCmd) queryStderr(ctx context.Context) {
	for starcli.IsRunning() && starcli.CMD != nil {
		select {
		case <-ctx.Done():
			return
		default:
		}
		out := make([]byte, 65535)
		n, err := starcli.errfile.Read(out)
		if n != 0 {
			starcli.lock.Lock()
			starcli.stderrBuf.Write(out[:n])
			starcli.lock.Unlock()
			for _, v := range out[:n] {
				starcli.errout = append(starcli.errout, v)
			}
		}
		if err != nil {
			if err == io.EOF {
				break
			} else {
				if !strings.Contains(err.Error(), "file already closed") {
					starcli.runerr = err
				}
				return
			}
		}
	}
	return
}
func (starcli *StarCmd) NowLineOutput() (string, error) {
	starcli.lock.Lock()
	buf, _ := starcli.stdoutBuf.ReadBytes('\n')
	buferr, _ := starcli.stderrBuf.ReadBytes(byte('\n'))
	starcli.lock.Unlock()
	if len(buferr) != 0 {
		return string(buf), errors.New(string(buferr))
	}
	return string(buf), nil
}

func (starcli *StarCmd) NowLineStdOut() string {
	starcli.lock.Lock()
	defer starcli.lock.Unlock()
	buf, _ := starcli.stdoutBuf.ReadBytes('\n')
	return string(buf)
}

func (starcli *StarCmd) NowLineStdErr() error {
	starcli.lock.Lock()
	defer starcli.lock.Unlock()
	buferr, _ := starcli.stderrBuf.ReadBytes(byte('\n'))
	if len(buferr) != 0 {
		return errors.New(string(buferr))
	}
	return nil
}

func (starcli *StarCmd) NowAllOutput() (string, error) {
	var outstr string
	starcli.lock.Lock()
	buf := make([]byte, starcli.stdoutBuf.Len())
	n, _ := starcli.stdoutBuf.Read(buf)
	starcli.lock.Unlock()
	if n != 0 {
		outstr = string(buf[:n])
	}
	if starcli.runerr != nil {
		return outstr, starcli.runerr
	}
	starcli.lock.Lock()
	buf = make([]byte, starcli.stderrBuf.Len())
	n, _ = starcli.stderrBuf.Read(buf)
	starcli.lock.Unlock()
	if n != 0 {
		return outstr, errors.New(string(buf[:n]))
	}
	return outstr, nil
}

func (starcli *StarCmd) NowStdOut() string {
	var outstr string
	starcli.lock.Lock()
	buf := make([]byte, starcli.stdoutBuf.Len())
	n, _ := starcli.stdoutBuf.Read(buf)
	starcli.lock.Unlock()
	if n != 0 {
		outstr = string(buf[:n])
	}
	return outstr
}

func (starcli *StarCmd) NowStdErr() error {
	starcli.lock.Lock()
	buf := make([]byte, starcli.stderrBuf.Len())
	n, _ := starcli.stderrBuf.Read(buf)
	starcli.lock.Unlock()
	if n != 0 {
		return errors.New(string(buf[:n]))
	}
	return nil
}

func (starcli *StarCmd) AllOutPut() (string, error) {
	err := starcli.runerr
	if err == nil && len(starcli.errout) != 0 {
		err = errors.New(string(starcli.errout))
	}
	return string(starcli.stdout), err
}

func (starcli *StarCmd) AllStdOut() string {
	return string(starcli.stdout)
}

func (starcli *StarCmd) AllStdErr() error {
	err := starcli.runerr
	if err == nil && len(starcli.errout) != 0 {
		err = errors.New(string(starcli.errout))
	}
	return err
}

func (starcli *StarCmd) setRunning(alive bool) {
	if alive {
		val := atomic.LoadInt32(&starcli.running)
		if val == 0 {
			atomic.AddInt32(&starcli.running, 1)
		} else {
			atomic.AddInt32(&starcli.running, 1-val)
		}
		return
	}
	val := atomic.LoadInt32(&starcli.running)
	if val == 1 {
		atomic.AddInt32(&starcli.running, -1)
	} else {
		atomic.AddInt32(&starcli.running, -val)
	}
}
func (starcli *StarCmd) Start() error {
	if err := starcli.CMD.Start(); err != nil {
		return err
	}
	starcli.setRunning(true)
	go func() {
		err := starcli.CMD.Wait()
		if err != nil {
			starcli.runerr = err
		}
		starcli.stopctxfunc()
		starcli.setRunning(false)
		if starcli.CMD.ProcessState != nil {
			starcli.exitcode = starcli.CMD.ProcessState.Sys().(syscall.WaitStatus).ExitStatus()
		}
	}()
	go starcli.queryStdout(starcli.stopctx)
	go starcli.queryStderr(starcli.stopctx)
	go func(ctx context.Context) {
		if len(starcli.prewrite) != 0 {
			for _, v := range starcli.prewrite {
				select {
				case <-ctx.Done():
					return
				default:
					break
				}
				starcli.WriteCmd(v)
				time.Sleep(starcli.prewritetime)
			}
		}
	}(starcli.stopctx)
	return nil
}

func (starcli *StarCmd) IsRunning() bool {
	return 0 != atomic.LoadInt32(&starcli.running)
}

func (starcli *StarCmd) Stoped() <-chan struct{} {
	return starcli.stopctx.Done()
}

func (starcli *StarCmd) Exec(cmd string, wait int) (string, error) {
	starcli.infile.Write([]byte(cmd + "\n"))
	time.Sleep(time.Millisecond * time.Duration(wait))
	return starcli.NowAllOutput()
}

func (starcli *StarCmd) WriteCmd(cmdstr string) {
	starcli.infile.Write([]byte(cmdstr + "\n"))
}

func (starcli *StarCmd) PreWrite(cmd ...string) {
	for _, v := range cmd {
		starcli.prewrite = append(starcli.prewrite, v)
	}
}

func (starcli *StarCmd) PreWriteInterval(dt time.Duration) {
	starcli.prewritetime = dt
}

func (starcli *StarCmd) ExitCode() int {
	return starcli.exitcode
}

func (starcli *StarCmd) Kill() error {
	err := starcli.CMD.Process.Kill()
	if err != nil {
		return err
	}
	starcli.setRunning(false)
	return nil
}

func (starcli *StarCmd) GetPid() int {
	return starcli.CMD.Process.Pid
}

func (starcli *StarCmd) Signal(sig os.Signal) error {
	return starcli.CMD.Process.Signal(sig)
}