tests/regres/shell/shell.go

   1 // Copyright 2019 The SwiftShader Authors. All Rights Reserved.
   2 //
   3 // Licensed under the Apache License, Version 2.0 (the "License");
   4 // you may not use this file except in compliance with the License.
   5 // You may obtain a copy of the License at
   6 //
   7 //    http://www.apache.org/licenses/LICENSE-2.0
   8 //
   9 // Unless required by applicable law or agreed to in writing, software
  10 // distributed under the License is distributed on an "AS IS" BASIS,
  11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 // See the License for the specific language governing permissions and
  13 // limitations under the License.
  14
  15 // Package shell provides functions for running sub-processes.
  16 package shell
  17
  18 import (
  19         "bytes"
  20         "fmt"
  21         "log"
  22         "os"
  23         "os/exec"
  24         "os/signal"
  25         "strconv"
  26         "syscall"
  27         "time"
  28
  29         "../cause"
  30 )
  31
  32 // MaxProcMemory is the maximum virtual memory per child process
  33 var MaxProcMemory uint64 = 2 * 1024 * 1024 * 1024 // 2GB
  34
  35 func init() {
  36         // As we are going to be running a number of tests concurrently, we need to
  37         // limit the amount of virtual memory each test uses, otherwise memory
  38         // hungry tests can bring the whole system down into a swapping apocalypse.
  39         //
  40         // Linux has the setrlimit() function to limit a process (and child's)
  41         // virtual memory usage - but we cannot call this from the regres process
  42         // as this process may need more memory than the limit allows.
  43         //
  44         // Unfortunately golang has no native support for setting rlimits for child
  45         // processes (https://github.com/golang/go/issues/6603), so we instead wrap
  46         // the exec to the test executable with another child regres process using a
  47         // special --exec mode:
  48         //
  49         // [regres] -> [regres --exec <test-exe N args...>] -> [test-exe]
  50         //               ^^^^
  51         //          (calls rlimit() with memory limit of N bytes)
  52
  53         if len(os.Args) > 3 && os.Args[1] == "--exec" {
  54                 exe := os.Args[2]
  55                 limit, err := strconv.ParseUint(os.Args[3], 10, 64)
  56                 if err != nil {
  57                         log.Fatalf("Expected memory limit as 3rd argument. %v\n", err)
  58                 }
  59                 if limit > 0 {
  60                         if err := syscall.Setrlimit(syscall.RLIMIT_AS, &syscall.Rlimit{Cur: limit, Max: limit}); err != nil {
  61                                 log.Fatalln(cause.Wrap(err, "Setrlimit").Error())
  62                         }
  63                 }
  64                 cmd := exec.Command(exe, os.Args[4:]...)
  65                 cmd.Stdin = os.Stdin
  66                 cmd.Stdout = os.Stdout
  67                 cmd.Stderr = os.Stderr
  68                 if err := cmd.Start(); err != nil {
  69                         os.Stderr.WriteString(err.Error())
  70                         os.Exit(1)
  71                 }
  72                 // Forward signals to the child process
  73                 c := make(chan os.Signal, 1)
  74                 signal.Notify(c, os.Interrupt)
  75                 go func() {
  76                         for sig := range c {
  77                                 cmd.Process.Signal(sig)
  78                         }
  79                 }()
  80                 cmd.Wait()
  81                 close(c)
  82                 os.Exit(cmd.ProcessState.ExitCode())
  83         }
  84 }
  85
  86 // Shell runs the executable exe with the given arguments, in the working
  87 // directory wd.
  88 // If the process does not finish within timeout a errTimeout will be returned.
  89 func Shell(timeout time.Duration, exe, wd string, args ...string) error {
  90         if out, err := Exec(timeout, exe, wd, nil, args...); err != nil {
  91                 return cause.Wrap(err, "%s", out)
  92         }
  93         return nil
  94 }
  95
  96 // Exec runs the executable exe with the given arguments, in the working
  97 // directory wd, with the custom environment flags.
  98 // If the process does not finish within timeout a errTimeout will be returned.
  99 func Exec(timeout time.Duration, exe, wd string, env []string, args ...string) ([]byte, error) {
 100         // Shell via regres: --exec N <exe> <args...>
 101         // See main() for details.
 102         args = append([]string{"--exec", exe, fmt.Sprintf("%v", MaxProcMemory)}, args...)
 103         b := bytes.Buffer{}
 104         c := exec.Command(os.Args[0], args...)
 105         c.Dir = wd
 106         c.Env = env
 107         c.Stdout = &b
 108         c.Stderr = &b
 109
 110         if err := c.Start(); err != nil {
 111                 return nil, err
 112         }
 113
 114         res := make(chan error)
 115         go func() { res <- c.Wait() }()
 116
 117         select {
 118         case <-time.NewTimer(timeout).C:
 119                 log.Printf("Timeout for process %v\n", c.Process.Pid)
 120                 c.Process.Signal(syscall.SIGINT)
 121                 time.Sleep(time.Second * 5)
 122                 if !c.ProcessState.Exited() {
 123                         log.Printf("Process %v still has not exited, killing\n", c.Process.Pid)
 124                         syscall.Kill(-c.Process.Pid, syscall.SIGKILL)
 125                 }
 126                 return b.Bytes(), ErrTimeout{exe, timeout}
 127         case err := <-res:
 128                 return b.Bytes(), err
 129         }
 130 }
 131
 132 // ErrTimeout is the error returned when a process does not finish with its
 133 // permitted time.
 134 type ErrTimeout struct {
 135         process string
 136         timeout time.Duration
 137 }
 138
 139 func (e ErrTimeout) Error() string {
 140         return fmt.Sprintf("'%v' did not return after %v", e.process, e.timeout)
 141 }