1
0

manager.go 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. // Copyright 2014 The Gogs Authors. All rights reserved.
  2. // Use of this source code is governed by a MIT-style
  3. // license that can be found in the LICENSE file.
  4. package process
  5. import (
  6. "bytes"
  7. "errors"
  8. "fmt"
  9. "os/exec"
  10. "sync"
  11. "time"
  12. log "unknwon.dev/clog/v2"
  13. )
  14. var ErrExecTimeout = errors.New("process execution timeout")
  15. const DEFAULT_TIMEOUT = 60 * time.Second
  16. // Process represents a running process calls shell command.
  17. type Process struct {
  18. PID int64
  19. Description string
  20. Start time.Time
  21. Cmd *exec.Cmd
  22. }
  23. type pidCounter struct {
  24. sync.Mutex
  25. // The current number of pid, initial is 0, and increase 1 every time it's been used.
  26. pid int64
  27. }
  28. func (c *pidCounter) PID() int64 {
  29. c.pid++
  30. return c.pid
  31. }
  32. var (
  33. counter = new(pidCounter)
  34. Processes []*Process
  35. )
  36. // Add adds a process to global list and returns its PID.
  37. func Add(desc string, cmd *exec.Cmd) int64 {
  38. counter.Lock()
  39. defer counter.Unlock()
  40. pid := counter.PID()
  41. Processes = append(Processes, &Process{
  42. PID: pid,
  43. Description: desc,
  44. Start: time.Now(),
  45. Cmd: cmd,
  46. })
  47. return pid
  48. }
  49. // Remove removes a process from global list.
  50. // It returns true if the process is found and removed by given pid.
  51. func Remove(pid int64) bool {
  52. counter.Lock()
  53. defer counter.Unlock()
  54. for i := range Processes {
  55. if Processes[i].PID == pid {
  56. Processes = append(Processes[:i], Processes[i+1:]...)
  57. return true
  58. }
  59. }
  60. return false
  61. }
  62. // Exec starts executing a shell command in given path, it tracks corresponding process and timeout.
  63. func ExecDir(timeout time.Duration, dir, desc, cmdName string, args ...string) (string, string, error) {
  64. if timeout == -1 {
  65. timeout = DEFAULT_TIMEOUT
  66. }
  67. bufOut := new(bytes.Buffer)
  68. bufErr := new(bytes.Buffer)
  69. cmd := exec.Command(cmdName, args...)
  70. cmd.Dir = dir
  71. cmd.Stdout = bufOut
  72. cmd.Stderr = bufErr
  73. if err := cmd.Start(); err != nil {
  74. return "", err.Error(), err
  75. }
  76. pid := Add(desc, cmd)
  77. done := make(chan error)
  78. go func() {
  79. done <- cmd.Wait()
  80. }()
  81. var err error
  82. select {
  83. case <-time.After(timeout):
  84. if errKill := Kill(pid); errKill != nil {
  85. log.Error("Failed to kill timeout process [pid: %d, desc: %s]: %v", pid, desc, errKill)
  86. }
  87. <-done
  88. return "", ErrExecTimeout.Error(), ErrExecTimeout
  89. case err = <-done:
  90. }
  91. Remove(pid)
  92. return bufOut.String(), bufErr.String(), err
  93. }
  94. // Exec starts executing a shell command, it tracks corresponding process and timeout.
  95. func ExecTimeout(timeout time.Duration, desc, cmdName string, args ...string) (string, string, error) {
  96. return ExecDir(timeout, "", desc, cmdName, args...)
  97. }
  98. // Exec starts executing a shell command, it tracks corresponding its process and use default timeout.
  99. func Exec(desc, cmdName string, args ...string) (string, string, error) {
  100. return ExecDir(-1, "", desc, cmdName, args...)
  101. }
  102. // Kill kills and removes a process from global list.
  103. func Kill(pid int64) error {
  104. for _, proc := range Processes {
  105. if proc.PID == pid {
  106. if proc.Cmd != nil && proc.Cmd.Process != nil &&
  107. proc.Cmd.ProcessState != nil && !proc.Cmd.ProcessState.Exited() {
  108. if err := proc.Cmd.Process.Kill(); err != nil {
  109. return fmt.Errorf("fail to kill process [pid: %d, desc: %s]: %v", proc.PID, proc.Description, err)
  110. }
  111. }
  112. Remove(pid)
  113. return nil
  114. }
  115. }
  116. return nil
  117. }