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// +build darwin linux
16
17package shell
18
19import (
20	"bytes"
21	"fmt"
22	"log"
23	"os"
24	"os/exec"
25	"os/signal"
26	"strconv"
27	"syscall"
28	"time"
29
30	"../cause"
31)
32
33func init() {
34	// As we are going to be running a number of tests concurrently, we need to
35	// limit the amount of virtual memory each test uses, otherwise memory
36	// hungry tests can bring the whole system down into a swapping apocalypse.
37	//
38	// Linux has the setrlimit() function to limit a process (and child's)
39	// virtual memory usage - but we cannot call this from the regres process
40	// as this process may need more memory than the limit allows.
41	//
42	// Unfortunately golang has no native support for setting rlimits for child
43	// processes (https://github.com/golang/go/issues/6603), so we instead wrap
44	// the exec to the test executable with another child regres process using a
45	// special --exec mode:
46	//
47	// [regres] -> [regres --exec <test-exe N args...>] -> [test-exe]
48	//               ^^^^
49	//          (calls rlimit() with memory limit of N bytes)
50
51	if len(os.Args) > 3 && os.Args[1] == "--exec" {
52		exe := os.Args[2]
53		limit, err := strconv.ParseUint(os.Args[3], 10, 64)
54		if err != nil {
55			log.Fatalf("Expected memory limit as 3rd argument. %v\n", err)
56		}
57		if limit > 0 {
58			if err := syscall.Setrlimit(syscall.RLIMIT_AS, &syscall.Rlimit{Cur: limit, Max: limit}); err != nil {
59				log.Fatalln(cause.Wrap(err, "Setrlimit").Error())
60			}
61		}
62		cmd := exec.Command(exe, os.Args[4:]...)
63		cmd.Stdin = os.Stdin
64		cmd.Stdout = os.Stdout
65		cmd.Stderr = os.Stderr
66		if err := cmd.Start(); err != nil {
67			os.Stderr.WriteString(err.Error())
68			os.Exit(1)
69		}
70		// Forward signals to the child process
71		c := make(chan os.Signal, 1)
72		signal.Notify(c, os.Interrupt)
73		go func() {
74			for sig := range c {
75				cmd.Process.Signal(sig)
76			}
77		}()
78		cmd.Wait()
79		close(c)
80		os.Exit(cmd.ProcessState.ExitCode())
81	}
82}
83
84// Exec runs the executable exe with the given arguments, in the working
85// directory wd, with the custom environment flags.
86// If the process does not finish within timeout a errTimeout will be returned.
87func Exec(timeout time.Duration, exe, wd string, env []string, args ...string) ([]byte, error) {
88	// Shell via regres: --exec N <exe> <args...>
89	// See main() for details.
90	args = append([]string{"--exec", exe, fmt.Sprintf("%v", MaxProcMemory)}, args...)
91	b := bytes.Buffer{}
92	c := exec.Command(os.Args[0], args...)
93	c.Dir = wd
94	c.Env = env
95	c.Stdout = &b
96	c.Stderr = &b
97
98	if err := c.Start(); err != nil {
99		return nil, err
100	}
101
102	res := make(chan error)
103	go func() { res <- c.Wait() }()
104
105	select {
106	case <-time.NewTimer(timeout).C:
107		c.Process.Signal(syscall.SIGINT)
108		time.Sleep(time.Second * 3)
109		if c.ProcessState == nil || !c.ProcessState.Exited() {
110			log.Printf("Process %v still has not exited, killing\n", c.Process.Pid)
111			syscall.Kill(-c.Process.Pid, syscall.SIGKILL)
112		}
113		return b.Bytes(), ErrTimeout{exe, timeout}
114	case err := <-res:
115		return b.Bytes(), err
116	}
117}
118