services/crunch-run/background.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "encoding/json"
   9         "fmt"
  10         "io"
  11         "io/ioutil"
  12         "os"
  13         "os/exec"
  14         "path/filepath"
  15         "strings"
  16         "syscall"
  17         "time"
  18 )
  19
  20 var (
  21         lockdir    = "/var/lock"
  22         lockprefix = "crunch-run-"
  23         locksuffix = ".lock"
  24 )
  25
  26 // procinfo is saved in each process's lockfile.
  27 type procinfo struct {
  28         UUID   string
  29         PID    int
  30         Stdout string
  31         Stderr string
  32 }
  33
  34 // Detach acquires a lock for the given uuid, and starts the current
  35 // program as a child process (with -no-detach prepended to the given
  36 // arguments so the child knows not to detach again). The lock is
  37 // passed along to the child process.
  38 func Detach(uuid string, args []string, stdout, stderr io.Writer) int {
  39         return exitcode(stderr, detach(uuid, args, stdout, stderr))
  40 }
  41 func detach(uuid string, args []string, stdout, stderr io.Writer) error {
  42         lockfile, err := os.OpenFile(filepath.Join(lockdir, lockprefix+uuid+locksuffix), os.O_CREATE|os.O_RDWR, 0700)
  43         if err != nil {
  44                 return err
  45         }
  46         defer lockfile.Close()
  47         err = syscall.Flock(int(lockfile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
  48         if err != nil {
  49                 return err
  50         }
  51         lockfile.Truncate(0)
  52
  53         outfile, err := ioutil.TempFile("", "crunch-run-"+uuid+"-stdout-")
  54         if err != nil {
  55                 return err
  56         }
  57         defer outfile.Close()
  58         errfile, err := ioutil.TempFile("", "crunch-run-"+uuid+"-stderr-")
  59         if err != nil {
  60                 os.Remove(outfile.Name())
  61                 return err
  62         }
  63         defer errfile.Close()
  64
  65         cmd := exec.Command(args[0], append([]string{"-no-detach"}, args[1:]...)...)
  66         cmd.Stdout = outfile
  67         cmd.Stderr = errfile
  68         // Child inherits lockfile.
  69         cmd.ExtraFiles = []*os.File{lockfile}
  70         // Ensure child isn't interrupted even if we receive signals
  71         // from parent (sshd) while sending lockfile content to
  72         // caller.
  73         cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
  74         err = cmd.Start()
  75         if err != nil {
  76                 os.Remove(outfile.Name())
  77                 os.Remove(errfile.Name())
  78                 return err
  79         }
  80
  81         w := io.MultiWriter(stdout, lockfile)
  82         err = json.NewEncoder(w).Encode(procinfo{
  83                 UUID:   uuid,
  84                 PID:    cmd.Process.Pid,
  85                 Stdout: outfile.Name(),
  86                 Stderr: errfile.Name(),
  87         })
  88         if err != nil {
  89                 os.Remove(outfile.Name())
  90                 os.Remove(errfile.Name())
  91                 return err
  92         }
  93         return nil
  94 }
  95
  96 // KillProcess finds the crunch-run process corresponding to the given
  97 // uuid, and sends the given signal to it. It then waits up to 1
  98 // second for the process to die. It returns 0 if the process is
  99 // successfully killed or didn't exist in the first place.
 100 func KillProcess(uuid string, signal syscall.Signal, stdout, stderr io.Writer) int {
 101         return exitcode(stderr, kill(uuid, signal, stdout, stderr))
 102 }
 103
 104 func kill(uuid string, signal syscall.Signal, stdout, stderr io.Writer) error {
 105         path := filepath.Join(lockdir, lockprefix+uuid+locksuffix)
 106         f, err := os.Open(path)
 107         if os.IsNotExist(err) {
 108                 return nil
 109         } else if err != nil {
 110                 return err
 111         }
 112         defer f.Close()
 113
 114         var pi procinfo
 115         err = json.NewDecoder(f).Decode(&pi)
 116         if err != nil {
 117                 return fmt.Errorf("%s: %s\n", path, err)
 118         }
 119
 120         if pi.UUID != uuid || pi.PID == 0 {
 121                 return fmt.Errorf("%s: bogus procinfo: %+v", path, pi)
 122         }
 123
 124         proc, err := os.FindProcess(pi.PID)
 125         if err != nil {
 126                 return err
 127         }
 128
 129         err = proc.Signal(signal)
 130         for deadline := time.Now().Add(time.Second); err == nil && time.Now().Before(deadline); time.Sleep(time.Second / 100) {
 131                 err = proc.Signal(syscall.Signal(0))
 132         }
 133         if err == nil {
 134                 return fmt.Errorf("pid %d: sent signal %d (%s) but process is still alive", pi.PID, signal, signal)
 135         }
 136         fmt.Fprintf(stderr, "pid %d: %s\n", pi.PID, err)
 137         return nil
 138 }
 139
 140 // List UUIDs of active crunch-run processes.
 141 func ListProcesses(stdout, stderr io.Writer) int {
 142         return exitcode(stderr, filepath.Walk(lockdir, func(path string, info os.FileInfo, err error) error {
 143                 if info.IsDir() {
 144                         return filepath.SkipDir
 145                 }
 146                 if name := info.Name(); !strings.HasPrefix(name, lockprefix) || !strings.HasSuffix(name, locksuffix) {
 147                         return nil
 148                 }
 149                 if info.Size() == 0 {
 150                         // race: process has opened/locked but hasn't yet written pid/uuid
 151                         return nil
 152                 }
 153
 154                 f, err := os.Open(path)
 155                 if err != nil {
 156                         return nil
 157                 }
 158                 defer f.Close()
 159
 160                 // TODO: Do this check without risk of disrupting lock
 161                 // acquisition during races, e.g., by connecting to a
 162                 // unix socket or checking /proc/$pid/fd/$n ->
 163                 // lockfile.
 164                 err = syscall.Flock(int(f.Fd()), syscall.LOCK_SH|syscall.LOCK_NB)
 165                 if err == nil {
 166                         // lockfile is stale
 167                         err := os.Remove(path)
 168                         if err != nil {
 169                                 fmt.Fprintln(stderr, err)
 170                         }
 171                         return nil
 172                 }
 173
 174                 var pi procinfo
 175                 err = json.NewDecoder(f).Decode(&pi)
 176                 if err != nil {
 177                         fmt.Fprintf(stderr, "%s: %s\n", path, err)
 178                         return nil
 179                 }
 180                 if pi.UUID == "" || pi.PID == 0 {
 181                         fmt.Fprintf(stderr, "%s: bogus procinfo: %+v", path, pi)
 182                         return nil
 183                 }
 184
 185                 fmt.Fprintln(stdout, pi.UUID)
 186                 return nil
 187         }))
 188 }
 189
 190 // If err is nil, return 0 ("success"); otherwise, print err to stderr
 191 // and return 1.
 192 func exitcode(stderr io.Writer, err error) int {
 193         if err != nil {
 194                 fmt.Fprintln(stderr, err)
 195                 return 1
 196         }
 197         return 0
 198 }