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