github.com/quite/nomad@v0.8.6/helper/testtask/testtask.go (about)

     1  // Package testtask implements a portable set of commands useful as stand-ins
     2  // for user tasks.
     3  package testtask
     4  
     5  import (
     6  	"fmt"
     7  	"io/ioutil"
     8  	"os"
     9  	"os/exec"
    10  	"strconv"
    11  	"syscall"
    12  	"time"
    13  
    14  	"github.com/hashicorp/nomad/nomad/structs"
    15  	"github.com/kardianos/osext"
    16  )
    17  
    18  // Path returns the path to the currently running executable.
    19  func Path() string {
    20  	path, err := osext.Executable()
    21  	if err != nil {
    22  		panic(err)
    23  	}
    24  	return path
    25  }
    26  
    27  // SetCmdEnv configures the environment of cmd so that Run executes a testtask
    28  // script when called from within cmd.
    29  func SetCmdEnv(cmd *exec.Cmd) {
    30  	cmd.Env = append(os.Environ(), "TEST_TASK=execute")
    31  }
    32  
    33  // SetTaskEnv configures the environment of t so that Run executes a testtask
    34  // script when called from within t.
    35  func SetTaskEnv(t *structs.Task) {
    36  	if t.Env == nil {
    37  		t.Env = map[string]string{}
    38  	}
    39  	t.Env["TEST_TASK"] = "execute"
    40  }
    41  
    42  // Run interprets os.Args as a testtask script if the current program was
    43  // launched with an environment configured by SetCmdEnv or SetTaskEnv. It
    44  // returns false if the environment was not set by this package.
    45  func Run() bool {
    46  	switch tm := os.Getenv("TEST_TASK"); tm {
    47  	case "":
    48  		return false
    49  	case "execute":
    50  		execute()
    51  		return true
    52  	default:
    53  		fmt.Fprintf(os.Stderr, "unexpected value for TEST_TASK, \"%s\"\n", tm)
    54  		os.Exit(1)
    55  		return true
    56  	}
    57  }
    58  
    59  func execute() {
    60  	if len(os.Args) < 2 {
    61  		fmt.Fprintln(os.Stderr, "no command provided")
    62  		os.Exit(1)
    63  	}
    64  
    65  	args := os.Args[1:]
    66  
    67  	// popArg removes the first argument from args and returns it.
    68  	popArg := func() string {
    69  		s := args[0]
    70  		args = args[1:]
    71  		return s
    72  	}
    73  
    74  	// execute a sequence of operations from args
    75  	for len(args) > 0 {
    76  		switch cmd := popArg(); cmd {
    77  
    78  		case "sleep":
    79  			// sleep <dur>: sleep for a duration indicated by the first
    80  			// argument
    81  			if len(args) < 1 {
    82  				fmt.Fprintln(os.Stderr, "expected arg for sleep")
    83  				os.Exit(1)
    84  			}
    85  			dur, err := time.ParseDuration(popArg())
    86  			if err != nil {
    87  				fmt.Fprintf(os.Stderr, "could not parse sleep time: %v", err)
    88  				os.Exit(1)
    89  			}
    90  			time.Sleep(dur)
    91  
    92  		case "echo":
    93  			// echo <msg>: write the msg followed by a newline to stdout.
    94  			fmt.Println(popArg())
    95  
    96  		case "write":
    97  			// write <msg> <file>: write a message to a file. The first
    98  			// argument is the msg. The second argument is the path to the
    99  			// target file.
   100  			if len(args) < 2 {
   101  				fmt.Fprintln(os.Stderr, "expected two args for write")
   102  				os.Exit(1)
   103  			}
   104  			msg := popArg()
   105  			file := popArg()
   106  			ioutil.WriteFile(file, []byte(msg), 0666)
   107  
   108  		case "pgrp":
   109  			// pgrp <group_int> puts the pid in a new process group
   110  			if len(args) < 1 {
   111  				fmt.Fprintln(os.Stderr, "expected process group number for pgrp")
   112  				os.Exit(1)
   113  			}
   114  			num := popArg()
   115  			grp, err := strconv.Atoi(num)
   116  			if err != nil {
   117  				fmt.Fprintf(os.Stderr, "failed to convert process group number %q: %v\n", num, err)
   118  				os.Exit(1)
   119  			}
   120  			if err := syscall.Setpgid(0, grp); err != nil {
   121  				fmt.Fprintf(os.Stderr, "failed to set process group: %v\n", err)
   122  				os.Exit(1)
   123  			}
   124  
   125  		case "fork/exec":
   126  			// fork/exec <pid_file> <args> forks execs the helper process
   127  			if len(args) < 2 {
   128  				fmt.Fprintln(os.Stderr, "expect pid file and remaining args to fork exec")
   129  				os.Exit(1)
   130  			}
   131  
   132  			pidFile := popArg()
   133  
   134  			cmd := exec.Command(Path(), args...)
   135  			SetCmdEnv(cmd)
   136  			if err := cmd.Start(); err != nil {
   137  				fmt.Fprintf(os.Stderr, "failed to fork/exec: %v\n", err)
   138  				os.Exit(1)
   139  			}
   140  
   141  			if err := ioutil.WriteFile(pidFile, []byte(fmt.Sprintf("%d", cmd.Process.Pid)), 777); err != nil {
   142  				fmt.Fprintf(os.Stderr, "failed to write pid file: %v\n", err)
   143  				os.Exit(1)
   144  			}
   145  
   146  			if err := cmd.Wait(); err != nil {
   147  				fmt.Fprintf(os.Stderr, "wait failed: %v\n", err)
   148  				os.Exit(1)
   149  			}
   150  			return
   151  
   152  		default:
   153  			fmt.Fprintln(os.Stderr, "unknown command:", cmd)
   154  			os.Exit(1)
   155  		}
   156  	}
   157  }