pmt/db/proc.go

package db

import (
	"errors"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
)

const (
	StatusUnknown  uint8 = iota // 未知状态
	StatusStarting              // 正在启动
	StatusRunning               // 正在运行
	StatusStopping              // 正在停止
	StatusStopped               // 已经停止
)

type Proc struct {
	sync.RWMutex
	tempDir    string      // 数据缓存目录
	stdinFile  string      // 数据输入文件
	stdoutFile string      // 数据输出文件
	stderrFile string      // 错误输出文件
	pidfile    string      // 记录PID的文件
	starts     int         // 重启次数
	status     uint8       // 程序状态
	pid        int         // 进程PID
	app        *App        // 应用信息
	process    *os.Process // 系统进程
}

func NewProc(app *App) *Proc {
	tempDir, err := TempDir(app.Name)
	if err != nil {
		panic(err)
	}

	proc := &Proc{
		tempDir:    tempDir,
		stdinFile:  filepath.Join(tempDir, "stdin.txt"),
		stdoutFile: filepath.Join(tempDir, "stdout.txt"),
		stderrFile: filepath.Join(tempDir, "stderr.txt"),
		pidfile:    filepath.Join(tempDir, "pid.txt"),
		starts:     app.StartCount,
		status:     app.Status,
		pid:        app.PID,
		app:        app,
		process:    nil,
	}

	inner, err := os.FindProcess(app.PID)
	if err == nil && inner.Signal(syscall.Signal(0)) == nil {
		proc.process = inner
	}

	return proc
}

// Start 启动进程
func (p *Proc) Start() error {
	p.RLock()
	if p.status != StatusUnknown && p.status != StatusStopped {
		p.RUnlock()
		return errors.New("process was running")
	}
	stdinFile := p.stdinFile
	stdoutFile := p.stdoutFile
	stderrFile := p.stderrFile
	app := p.app
	p.setStatus(StatusStarting)
	p.RUnlock()

	// 确定数据输入文件
	stdin, err := GetFile(stdinFile)
	if err != nil {
		return err
	}
	if len(app.Stdin) > 0 {
		err = WriteFile(stdinFile, []byte(app.Stdin))
		if err != nil {
			return err
		}
	}

	// 确定数据输出文件
	// TODO 监听内容输入，记录到数据库
	stdout, err := GetFile(stdoutFile)
	if err != nil {
		return err
	}

	// 确定错误输出文件
	// TODO 监听内容输入，记录到数据库
	stderr, err := GetFile(stderrFile)
	if err != nil {
		return err
	}

	// 确定工作目录
	cwd := app.Cwd
	if len(cwd) == 0 {
		cwd, _ = os.Getwd()
	}

	// 确定进程属性
	attr := &os.ProcAttr{
		Dir:   cwd,
		Env:   os.Environ(),
		Files: []*os.File{stdin, stdout, stderr},
		//Sys:   &syscall.SysProcAttr{HideWindow: true},
	}

	// 确定启动参数
	var args []string
	if len(app.InterpreterArgs) > 0 {
		args = append(args, app.InterpreterArgs...)
	}
	if len(app.Script) > 0 {
		args = append(args, app.Script)
	}
	if len(app.Command) > 0 {
		args = append(args, app.Command...)
	}
	if len(app.Arguments) > 0 {
		args = append(args, app.Arguments...)
	}

	// 启动进程
	process, err := os.StartProcess(app.Interpreter, args, attr)
	if err != nil {
		return err
	}

	p.Lock()
	defer p.Unlock()

	p.process = process
	p.pid = process.Pid
	p.starts++

	err = WriteFile(p.pidfile, []byte(strconv.Itoa(p.pid)))
	if err != nil {
		go p.Stop(true)
		return err
	}

	err = UpdateApp(app.ID,
		ConfigStatus(StatusRunning),
		ConfigPid(process.Pid),
		IncreaseStartCounter)
	if err != nil {
		go p.Stop(true)
		return err
	}

	return nil
}

// Stop 停止进程
func (p *Proc) Stop(force ...bool) error {
	p.Lock()
	if p.process == nil {
		p.Unlock()
		return errors.New("process does not exist")
	}
	p.setStatus(StatusStopping)
	p.Unlock()

	// 确定是否优雅关停
	gracefully := true
	if len(force) > 0 {
		gracefully = force[0] == false
	}

	p.Lock()
	defer p.Unlock()

	var err error
	if gracefully {
		err = p.process.Signal(syscall.SIGTERM)
	} else {
		err = p.process.Signal(syscall.SIGKILL)
		p.process.Release()
	}

	p.setStatus(StatusStopped)
	UpdateApp(p.app.ID, ConfigPid(-1), ConfigStatus(StatusStopped))
	DeleteFile(p.pidfile)

	return err
}

// Restart 重启进程
func (p *Proc) Restart(force ...bool) error {
	// 如果进程可用则停止
	if p.IsAlive() {
		err := p.Stop(force...)
		if err != nil {
			return err
		}
	}

	return p.Start()
}

// Wait 等待进程结束
func (p *Proc) Wait() (*os.ProcessState, error) {
	p.RLock()
	defer p.RUnlock()
	if p.process == nil {
		return nil, errors.New("process does not started")
	}
	return p.process.Wait()
}

func (p *Proc) setStatus(status uint8) {
	UpdateApp(p.app.ID, ConfigStatus(status))
}

// IsAlive 判断进程是否激活
func (p *Proc) IsAlive() bool {
	p.RLock()
	pid := p.pid
	p.RUnlock()

	if proc, err := os.FindProcess(pid); err != nil {
		return false
	} else {
		return proc.Signal(syscall.Signal(0)) == nil
	}
}

// Release 释放进程资源
func (p *Proc) Release() error {
	p.Lock()
	defer p.Unlock()
	var err error
	if p.process != nil {
		err = p.process.Release()
	}
	if de := DeleteFile(p.pidfile); de != nil {
		return de
	}
	if de := DeleteFile(p.stdinFile); de != nil {
		return de
	}
	if de := DeleteFile(p.stdoutFile); de != nil {
		return de
	}
	if de := DeleteFile(p.stderrFile); de != nil {
		return de
	}
	if de := DeleteFile(p.tempDir); de != nil {
		return err
	}
	return err
}

// Status 获取运行状态
func (p *Proc) Status() uint8 {
	p.RLock()
	defer p.RUnlock()
	return p.status
}

// Pid 获取进程PID
func (p *Proc) Pid() int {
	p.RLock()
	defer p.RUnlock()
	return p.pid
}

func (p *Proc) NotifyStopped() {
	p.Lock()
	defer p.Unlock()
	p.pid = -1
}