diff --git a/process.go b/process.go index 30e9e31..ee1ad06 100644 --- a/process.go +++ b/process.go @@ -1,4 +1,4 @@ -// Copyright 2017 CoreSwitch +// Copyright 2017 CoreSwitch. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -15,28 +15,60 @@ package process import ( + "fmt" + "os" + "os/exec" "sync" + "time" + + "golang.org/x/net/context" ) type Process struct { - Name string - Vrf string - Args []string - File string + Name string + Vrf string + Args []string + File string + ErrLookup string + ErrStart string + ErrWait string + ExitFunc func() + State int + Cmd *exec.Cmd + StartTimer int + RetryTimer int } -type ProcessList map[string]*Process +type ProcessSlice []*Process -var ( - ProcessVrfMap = map[string]*ProcessList{} - ProcessMutex sync.RWMutex +const ( + PROCESS_STARTING = iota + PROCESS_RUNNING + PROCESS_RETRY + PROCESS_EXIT_CALLED + PROCESS_STOP_WAIT + PROCESS_STOP ) -func NewProcess(name string, vrf string, args ...string) *Process { +var ( + ProcessList = ProcessSlice{} + ProcessMutex sync.RWMutex +) + +var ProcessStateStr = map[int]string{ + PROCESS_STARTING: "Starting", + PROCESS_RUNNING: "Running", + PROCESS_RETRY: "Retry", + PROCESS_EXIT_CALLED: "Exit Called", + PROCESS_STOP_WAIT: "Stop Wait", + PROCESS_STOP: "Stop", +} + +func NewProcess(name string, args ...string) *Process { proc := &Process{ - Name: name, - Vrf: vrf, - Args: []string(args), + Name: name, + Args: []string(args), + RetryTimer: 1, } return proc } @@ -44,25 +76,169 @@ func NewProcess(name string, vrf string, args ...string) *Process { func ProcessRegister(proc *Process) { ProcessMutex.Lock() defer ProcessMutex.Unlock() + + ProcessList = append(ProcessList, proc) + proc.Start() } func ProcessUnregister(proc *Process) { ProcessMutex.Lock() defer ProcessMutex.Unlock() + + proc.Stop() + + procList := ProcessSlice{} + for _, p := range ProcessList { + if p != proc { + procList = append(procList, p) + } + } + ProcessList = procList } -func ProcessUnregisterByCommand(name string) { +func ProcessCount() int { ProcessMutex.Lock() defer ProcessMutex.Unlock() + + return len(ProcessList) } -func ProcessUnregisterByVrf() { +func ProcessStart(index int) { ProcessMutex.Lock() defer ProcessMutex.Unlock() + + if index <= 0 { + return + } + if len(ProcessList) < index { + return + } + proc := ProcessList[index-1] + if proc == nil { + return + } + proc.Start() +} + +func ProcessStop(index int) { + ProcessMutex.Lock() + defer ProcessMutex.Unlock() + + if index <= 0 { + return + } + if len(ProcessList) < index { + return + } + proc := ProcessList[index-1] + if proc == nil { + return + } + proc.Stop() } func (proc *Process) Start() { + if proc.ExitFunc != nil { + return + } + + proc.State = PROCESS_STOP + binary, err := exec.LookPath(proc.Name) + if err != nil { + proc.ErrLookup = err.Error() + return + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + var wg sync.WaitGroup + + wg.Add(1) + go func() { + defer wg.Done() + + for { + proc.State = PROCESS_STARTING + if proc.File != "" { + os.OpenFile(proc.File, os.O_RDWR|os.O_CREATE, 0644) + } + + cmd := exec.CommandContext(ctx, binary, proc.Args...) + + env := os.Environ() + if proc.Vrf != "" { + env = append(env, fmt.Sprintf("VRF=%s", proc.Vrf)) + env = append(env, "LD_PRELOAD=/usr/bin/vrf_socket.so") + } + cmd.Env = env + proc.Cmd = cmd + + if proc.StartTimer != 0 { + time.Sleep(time.Duration(proc.StartTimer) * time.Second) + } + + // fmt.Println("process:", cmd.Path, cmd.Args) + err = cmd.Start() + if err != nil { + proc.ErrStart = err.Error() + } + + proc.State = PROCESS_RUNNING + err = cmd.Wait() + if err != nil { + proc.ErrWait = err.Error() + } + + proc.State = PROCESS_RETRY + retryTimer := time.NewTimer(time.Duration(proc.RetryTimer) * time.Second) + select { + case <-retryTimer.C: + case <-done: + retryTimer.Stop() + return + } + } + }() + + proc.ExitFunc = func() { + proc.State = PROCESS_EXIT_CALLED + close(done) + cancel() + proc.State = PROCESS_STOP_WAIT + wg.Wait() + proc.State = PROCESS_STOP + } } func (proc *Process) Stop() { + if proc.ExitFunc != nil { + proc.ExitFunc() + proc.ExitFunc = nil + } +} + +func ProcessListShow() string { + str := "" + for pos, proc := range ProcessList { + str += fmt.Sprintf("%d %s", pos+1, proc.Name) + if proc.Vrf != "" { + str += fmt.Sprintf("@%s", proc.Vrf) + } + str += fmt.Sprintf(": %s", ProcessStateStr[proc.State]) + if proc.State == PROCESS_RUNNING && proc.Cmd != nil && proc.Cmd.Process != nil { + str += fmt.Sprintf(" (pid %d)", proc.Cmd.Process.Pid) + } + str += "\n" + if proc.ErrLookup != "" { + str += fmt.Sprintf(" Last Lookup Error: %s\n", proc.ErrLookup) + } + if proc.ErrStart != "" { + str += fmt.Sprintf(" Last Start Error: %s\n", proc.ErrStart) + } + if proc.ErrWait != "" { + str += fmt.Sprintf(" Last Wait Error: %s\n", proc.ErrWait) + } + str += fmt.Sprintf(" %s\n", proc.Args) + } + return str } diff --git a/process_test.go b/process_test.go index a82a5d7..bd859d1 100644 --- a/process_test.go +++ b/process_test.go @@ -15,11 +15,24 @@ package process import ( + "fmt" "testing" ) func TestProcessStart(t *testing.T) { args := []string{"-d", "-f"} - proc := NewProcess("dhcp", "", args...) - ProcessRegister(proc) + proc1 := NewProcess("dhcp", args...) + ProcessRegister(proc1) + proc2 := NewProcess("dhcp", args...) + ProcessRegister(proc2) + fmt.Println(ProcessCount()) + + ProcessUnregister(proc2) + fmt.Println(ProcessCount()) + proc3 := NewProcess("dhcp", args...) + ProcessUnregister(proc3) + fmt.Println(ProcessCount()) + + ProcessUnregister(proc1) + fmt.Println(ProcessCount()) }