github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/allocrunner/taskrunner/driver_handle.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	cstructs "github.com/hashicorp/nomad/client/structs"
     9  	"github.com/hashicorp/nomad/nomad/structs"
    10  	"github.com/hashicorp/nomad/plugins/drivers"
    11  )
    12  
    13  // NewDriverHandle returns a handle for task operations on a specific task
    14  func NewDriverHandle(driver drivers.DriverPlugin, taskID string, task *structs.Task, net *drivers.DriverNetwork) *DriverHandle {
    15  	return &DriverHandle{
    16  		driver: driver,
    17  		net:    net,
    18  		taskID: taskID,
    19  		task:   task,
    20  	}
    21  }
    22  
    23  // DriverHandle encapsulates a driver plugin client and task identifier and exposes
    24  // an api to perform driver operations on the task
    25  type DriverHandle struct {
    26  	driver drivers.DriverPlugin
    27  	net    *drivers.DriverNetwork
    28  	task   *structs.Task
    29  	taskID string
    30  }
    31  
    32  func (h *DriverHandle) ID() string {
    33  	return h.taskID
    34  }
    35  
    36  func (h *DriverHandle) WaitCh(ctx context.Context) (<-chan *drivers.ExitResult, error) {
    37  	return h.driver.WaitTask(ctx, h.taskID)
    38  }
    39  
    40  func (h *DriverHandle) Update(task *structs.Task) error {
    41  	return nil
    42  }
    43  
    44  func (h *DriverHandle) Kill() error {
    45  	return h.driver.StopTask(h.taskID, h.task.KillTimeout, h.task.KillSignal)
    46  }
    47  
    48  func (h *DriverHandle) Stats(ctx context.Context, interval time.Duration) (<-chan *cstructs.TaskResourceUsage, error) {
    49  	return h.driver.TaskStats(ctx, h.taskID, interval)
    50  }
    51  
    52  func (h *DriverHandle) Signal(s string) error {
    53  	return h.driver.SignalTask(h.taskID, s)
    54  }
    55  
    56  // Exec is the handled used by client endpoint handler to invoke the appropriate task driver exec.
    57  func (h *DriverHandle) Exec(timeout time.Duration, cmd string, args []string) ([]byte, int, error) {
    58  	command := append([]string{cmd}, args...)
    59  	res, err := h.driver.ExecTask(h.taskID, command, timeout)
    60  	if err != nil {
    61  		return nil, 0, err
    62  	}
    63  	return res.Stdout, res.ExitResult.ExitCode, res.ExitResult.Err
    64  }
    65  
    66  // ExecStreaming is the handled used by client endpoint handler to invoke the appropriate task driver exec.
    67  // while allowing to stream input and output
    68  func (h *DriverHandle) ExecStreaming(ctx context.Context,
    69  	command []string,
    70  	tty bool,
    71  	stream drivers.ExecTaskStream) error {
    72  
    73  	if impl, ok := h.driver.(drivers.ExecTaskStreamingRawDriver); ok {
    74  		return impl.ExecTaskStreamingRaw(ctx, h.taskID, command, tty, stream)
    75  	}
    76  
    77  	d, ok := h.driver.(drivers.ExecTaskStreamingDriver)
    78  	if !ok {
    79  		return fmt.Errorf("task driver does not support exec")
    80  	}
    81  
    82  	execOpts, doneCh := drivers.StreamToExecOptions(
    83  		ctx, command, tty, stream)
    84  
    85  	result, err := d.ExecTaskStreaming(ctx, h.taskID, execOpts)
    86  	if err != nil {
    87  		return err
    88  	}
    89  
    90  	execOpts.Stdout.Close()
    91  	execOpts.Stderr.Close()
    92  
    93  	select {
    94  	case err = <-doneCh:
    95  	case <-ctx.Done():
    96  		err = fmt.Errorf("exec task timed out: %v", ctx.Err())
    97  	}
    98  
    99  	if err != nil {
   100  		return err
   101  	}
   102  
   103  	return stream.Send(drivers.NewExecStreamingResponseExit(result.ExitCode))
   104  }
   105  
   106  func (h *DriverHandle) Network() *drivers.DriverNetwork {
   107  	return h.net
   108  }