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