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