github.com/janma/nomad@v0.11.3/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  }