github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/drivers/mock/command.go (about) 1 package mock 2 3 import ( 4 "errors" 5 "io" 6 "sync" 7 "time" 8 9 hclog "github.com/hashicorp/go-hclog" 10 bstructs "github.com/hashicorp/nomad/plugins/base/structs" 11 "github.com/hashicorp/nomad/plugins/drivers" 12 ) 13 14 func runCommand(c Command, stdout, stderr io.WriteCloser, cancelCh <-chan struct{}, pluginExitTimer <-chan time.Time, logger hclog.Logger) *drivers.ExitResult { 15 errCh := make(chan error, 1) 16 17 var wg sync.WaitGroup 18 wg.Add(1) 19 go func() { 20 defer wg.Done() 21 runCommandOutput(stdout, c.StdoutString, c.StdoutRepeat, c.stdoutRepeatDuration, cancelCh, logger, errCh) 22 }() 23 24 wg.Add(1) 25 go func() { 26 defer wg.Done() 27 runCommandOutput(stderr, c.StderrString, c.StderrRepeat, c.stderrRepeatDuration, cancelCh, logger, errCh) 28 }() 29 30 timer := time.NewTimer(c.runForDuration) 31 defer timer.Stop() 32 33 select { 34 case <-timer.C: 35 logger.Debug("run_for time elapsed; exiting", "run_for", c.RunFor) 36 case <-cancelCh: 37 logger.Debug("killed; exiting") 38 case <-pluginExitTimer: 39 logger.Debug("exiting plugin") 40 return &drivers.ExitResult{ 41 Err: bstructs.ErrPluginShutdown, 42 } 43 case err := <-errCh: 44 logger.Error("error running mock task; exiting", "error", err) 45 return &drivers.ExitResult{ 46 Err: err, 47 } 48 } 49 50 wg.Wait() 51 52 var exitErr error 53 if c.ExitErrMsg != "" { 54 exitErr = errors.New(c.ExitErrMsg) 55 } 56 57 return &drivers.ExitResult{ 58 ExitCode: c.ExitCode, 59 Signal: c.ExitSignal, 60 Err: exitErr, 61 } 62 } 63 64 func runCommandOutput(writer io.WriteCloser, 65 output string, outputRepeat int, repeatDuration time.Duration, 66 cancelCh <-chan struct{}, logger hclog.Logger, errCh chan error) { 67 68 defer writer.Close() 69 70 if output == "" { 71 return 72 } 73 74 if _, err := io.WriteString(writer, output); err != nil { 75 logger.Error("failed to write to stdout", "error", err) 76 errCh <- err 77 return 78 } 79 80 for i := 0; i < outputRepeat; i++ { 81 select { 82 case <-cancelCh: 83 logger.Warn("exiting before done writing output", "i", i, "total", outputRepeat) 84 return 85 case <-time.After(repeatDuration): 86 if _, err := io.WriteString(writer, output); err != nil { 87 logger.Error("failed to write to stdout", "error", err) 88 errCh <- err 89 return 90 } 91 } 92 } 93 }