Go1.8实现watchdog功能,实现依赖启动服务程序

package main

import (
"flag"
"fmt"
"log"
"os"
"os/signal"
"os/user"
"path"
"strconv"
"strings"
"sync"
"syscall"
"time"

"github.com/czxichen/command/watchdog"
conf "github.com/dlintw/goconf"
)

//+build windows,linux

const logDir = "./watchdog"

func newProc(svc *Service, null, pw *os.File) *os.ProcAttr {
return &os.ProcAttr{Dir: svc.path, Files: []*os.File{null, pw, pw}}
}

func setPriority(pid, priority uintptr) syscall.Errno {
return 0
}

var (
logpath    = flag.String("log_path", "", "Specify log path")
configFile = flag.String("config", "watchdog.ini", "Watchdog configuration file")
)

func cfgOpt(cfg *conf.ConfigFile, section, option string) string {
if !cfg.HasOption(section, option) {
return ""
}
s, err := cfg.GetString(section, option)
if err != nil {
log.Fatalf("Failed to get %s for %s: %v", option, section, err)
}
return s
}

func svcOpt(cfg *conf.ConfigFile, service, option string, required bool) string {
opt := cfgOpt(cfg, service, option)
if opt == "" && required {
log.Fatalf("Service %s has missing %s option", service, option)
}
return opt
}

var signalNames = map[syscall.Signal]string{
syscall.SIGINT:  "SIGINT",
syscall.SIGQUIT: "SIGQUIT",
syscall.SIGTERM: "SIGTERM",
}

func signalName(s syscall.Signal) string {
if name, ok := signalNames[s]; ok {
return name
}
return fmt.Sprintf("SIG %d", s)
}

type Shutdowner interface {
Shutdown()
}

func shutdownHandler(server Shutdowner) {
sigc := make(chan os.Signal, 3)
signal.Notify(sigc, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM)
go func() {
for s := range sigc {
name := s.String()
if sig, ok := s.(syscall.Signal); ok {
name = signalName(sig)
}
log.Printf("Received %v, initiating shutdown...", name)
server.Shutdown()
}
}()
}

var (
restartDelay      = 2 * time.Second
restartBackoff    = 5 * time.Second
restartBackoffMax = 60 * time.Second
)

type Watchdog struct {
services map[string]*Service
shutdown chan bool
}

func NewWatchdog() *Watchdog {
return &Watchdog{
services: make(map[string]*Service),
shutdown: make(chan bool),
}
}

//关闭服务
func (w *Watchdog) Shutdown() {
select {
case w.shutdown <- true:
default:
}
}

//添加服务,如果存在
func (w *Watchdog) AddService(name, binary string) (*Service, error) {
if _, ok := w.services[name]; ok {
return nil, fmt.Errorf("Service %q already exists", name)
}

svc := newService(name, binary)
w.services[name] = svc

return svc, nil
}

//启动服务
func (w *Watchdog) Walk() {
log.Printf("Seesaw watchdog starting...")

w.mapDependencies()

for _, svc := range w.services {
go svc.run()
}
<-w.shutdown
for _, svc := range w.services {
go svc.stop()
}
for _, svc := range w.services {
stopped := <-svc.stopped
svc.stopped <- stopped
}
}

//设置依赖关系
func (w *Watchdog) mapDependencies() {
for name := range w.services {
svc := w.services[name]
for depName := range svc.dependencies {
dep, ok := w.services[depName]
if !ok {
log.Fatalf("Failed to find dependency %q for service %q", depName, name)
}
svc.dependencies[depName] = dep //依赖谁,依赖启动后才会启动自身
dep.dependents[svc.name] = svc  //谁依赖它,依赖它的服务退出后,才退出本身
}
}
}

//默认的优先级为0
const prioProcess = 0

//定义服务的类型.
type Service struct {
name   string
binary string
path   string
args   []string

uid      uint32
gid      uint32
priority int

dependencies map[string]*Service
dependents   map[string]*Service

termTimeout time.Duration

lock    sync.Mutex
process *os.Process

done     chan bool
shutdown chan bool
started  chan bool
stopped  chan bool

failures uint64
restarts uint64

lastFailure time.Time
lastRestart time.Time
}

//初始化一个Service.
func newService(name, binary string) *Service {
return &Service{
name:         name,
binary:       binary,
args:         make([]string, 0),
dependencies: make(map[string]*Service),
dependents:   make(map[string]*Service),

done:     make(chan bool),
shutdown: make(chan bool, 1),
started:  make(chan bool, 1),
stopped:  make(chan bool, 1),

termTimeout: 5 * time.Second,
}
}

//给这个服务添加依赖.
func (svc *Service) AddDependency(name string) {
svc.dependencies[name] = nil
}

//为服务添加启动参数.
func (svc *Service) AddArgs(args string) {
svc.args = strings.Fields(args)
}

//为进程设置优先级,Windows下面无效.
func (svc *Service) SetPriority(priority int) error {
if priority < -20 || priority > 19 {
return fmt.Errorf("Invalid priority %d - must be between -20 and 19", priority)
}
svc.priority = priority
return nil
}

func (svc *Service) SetTermTimeout(tt time.Duration) {
svc.termTimeout = tt
}

func (svc *Service) SetUser(username string) error {
u, err := user.Lookup(username)
if err != nil {
return err
}
uid, err := strconv.Atoi(u.Uid)
if err != nil {
return err
}
gid, err := strconv.Atoi(u.Gid)
if err != nil {
return err
}
svc.uid = uint32(uid)
svc.gid = uint32(gid)
return nil
}

func (svc *Service) run() {
//如果存在依赖,要等依赖全部启动完毕之后才会自动自身.
for _, dep := range svc.dependencies {
log.Printf("Service %s waiting for %s to start", svc.name, dep.name)
select {
case started := <-dep.started:
dep.started <- started
case <-svc.shutdown:
goto done
}
}

for {
//如果启动失败,怎等待时间会延长,最大不超过restartBackoffMax时间
//程序启动必须是阻塞的,不然会重复运行
if svc.failures > 0 {
delay := time.Duration(svc.failures) * restartBackoff
if delay > restartBackoffMax {
delay = restartBackoffMax
}
log.Printf("Service %s has failed %d times - delaying %s before restart",
svc.name, svc.failures, delay)

select {
case <-time.After(delay):
case <-svc.shutdown:
goto done
}
}

svc.restarts++
svc.lastRestart = time.Now()
svc.runOnce()

select {
case <-time.After(restartDelay):
case <-svc.shutdown:
goto done
}
}
done:
svc.done <- true
}

//为服务创建日志文件
func (svc *Service) logFile() (*os.File, error) {
logName := svc.name + ".log"

if err := os.MkdirAll(logDir, 0666); err != nil {
if !os.IsExist(err) {
return nil, err
}
}
f, err := os.Create(path.Join(logDir, logName))
if err != nil {
return nil, err
}
fmt.Fprintf(f, "Log file for %s (stdout/stderr)\n", svc.name)
fmt.Fprintf(f, "Created at: %s\n", time.Now().Format("2006/01/02 15:04:05"))
return f, nil
}

//运行程序
func (svc *Service) runOnce() {
args := make([]string, len(svc.args)+1)
args[0] = svc.name
copy(args[1:], svc.args)

fmt.Println("Args:", args)
null, err := os.Open(os.DevNull)
if err != nil {
log.Printf("Service %s - failed to open %s: %v", svc.name, os.DevNull, err)
return
}

lfile, err := svc.logFile()
if err != nil {
log.Printf("Service %s - failed to create log file: %v", svc.name, err)
null.Close()
return
}

attr := newProc(svc, null, lfile)

log.Printf("Starting service %s...", svc.name)
proc, err := os.StartProcess(svc.binary, args, attr)
if err != nil {
log.Printf("Service %s failed to start: %v", svc.name, err)
svc.lastFailure = time.Now()
svc.failures++
null.Close()
return
}

null.Close()
lfile.Close()
svc.lock.Lock()
svc.process = proc
svc.lock.Unlock()

if err := setPriority(uintptr(proc.Pid), uintptr(svc.priority)); err != 0 {
log.Printf("Failed to set priority to %d for service %s: %v", svc.priority, svc.name, err)
}
select {
case svc.started <- true:
default:
}

state, err := svc.process.Wait()
if err != nil {
log.Printf("Service %s wait failed with %v", svc.name, err)
svc.lastFailure = time.Now()
svc.failures++
return
}
if !state.Success() {
log.Printf("Service %s exited with %v", svc.name, state)
svc.lastFailure = time.Now()
svc.failures++
return
}

svc.failures = 0
log.Printf("Service %s exited normally.", svc.name)
}

//给进程发送信号
func (svc *Service) signal(sig os.Signal) error {
svc.lock.Lock()
defer svc.lock.Unlock()
if svc.process == nil {
return nil
}
return svc.process.Signal(sig)
}

//停止服务
func (svc *Service) stop() {
log.Printf("Stopping service %s...", svc.name)
//等待依赖它的进程退出完毕之后再退出自己.
for _, dep := range svc.dependents {
log.Printf("Service %s waiting for %s to stop", svc.name, dep.name)
stopped := <-dep.stopped
dep.stopped <- stopped
}

svc.shutdown <- true
//首先给进程发送退出信号,如果超时没有退出,则直接发送Kill信号.
svc.signal(syscall.SIGTERM)
select {
case <-svc.done:
case <-time.After(svc.termTimeout):
svc.signal(syscall.SIGKILL)
<-svc.done
}
log.Printf("Service %s stopped", svc.name)
svc.stopped <- true
}

func main() {
flag.Parse()
if *logpath == "" {
*logpath = os.Args[0] + ".log"
}
logFile, err := os.Create(*logpath)
if err != nil {
log.Fatalf("Create log file error:%s\n", err.Error())
}
defer logFile.Close()
log.SetOutput(logFile)

cfg, err := conf.ReadConfigFile(*configFile)
if err != nil {
log.Fatalf("Failed to read config file %q: %v", *configFile, err)
}

fido := watchdog.NewWatchdog()

shutdownHandler(fido)
for _, name := range cfg.GetSections() {
if name == "default" {
continue
}

binary := svcOpt(cfg, name, "binary", true)
args := svcOpt(cfg, name, "args", false)

svc, err := fido.AddService(name, binary)
if err != nil {
log.Fatalf("Failed to add service %q: %v", name, err)
}
svc.AddArgs(args)
if dep := svcOpt(cfg, name, "dependency", false); dep != "" {
svc.AddDependency(dep)
}
if opt := svcOpt(cfg, name, "priority", false); opt != "" {
prio, err := strconv.Atoi(opt)
if err != nil {
log.Fatalf("Service %s has invalid priority %q: %v", name, opt, err)
}
if err := svc.SetPriority(prio); err != nil {
log.Fatalf("Failed to set priority for service %s: %v", name, err)
}
}
if opt := svcOpt(cfg, name, "term_timeout", false); opt != "" {
tt, err := time.ParseDuration(opt)
if err != nil {
log.Fatalf("Service %s has invalid term_timeout %q: %v", name, opt, err)
}
svc.SetTermTimeout(tt)
}

if user := svcOpt(cfg, name, "user", false); user != "" {
if err := svc.SetUser(user); err != nil {
log.Fatalf("Failed to set user for service %s: %v", name, err)
}
}
}

fido.Walk()
}