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  }