open-nomad/client/allocrunner/taskrunner/driver_handle.go
Michael Schurter e62795798d core: propagate remote task handles
Add a new driver capability: RemoteTasks.

When a task is run by a driver with RemoteTasks set, its TaskHandle will
be propagated to the server in its allocation's TaskState. If the task
is replaced due to a down node or draining, its TaskHandle will be
propagated to its replacement allocation.

This allows tasks to be scheduled in remote systems whose lifecycles are
disconnected from the Nomad node's lifecycle.

See https://github.com/hashicorp/nomad-driver-ecs for an example ECS
remote task driver.
2021-04-27 15:07:03 -07:00

112 lines
3 KiB
Go

package taskrunner
import (
"context"
"fmt"
"time"
cstructs "github.com/hashicorp/nomad/client/structs"
"github.com/hashicorp/nomad/nomad/structs"
"github.com/hashicorp/nomad/plugins/drivers"
)
// NewDriverHandle returns a handle for task operations on a specific task
func NewDriverHandle(driver drivers.DriverPlugin, taskID string, task *structs.Task, net *drivers.DriverNetwork) *DriverHandle {
return &DriverHandle{
driver: driver,
net: net,
taskID: taskID,
killSignal: task.KillSignal,
killTimeout: task.KillTimeout,
}
}
// DriverHandle encapsulates a driver plugin client and task identifier and exposes
// an api to perform driver operations on the task
type DriverHandle struct {
driver drivers.DriverPlugin
net *drivers.DriverNetwork
taskID string
killSignal string
killTimeout time.Duration
}
func (h *DriverHandle) ID() string {
return h.taskID
}
func (h *DriverHandle) WaitCh(ctx context.Context) (<-chan *drivers.ExitResult, error) {
return h.driver.WaitTask(ctx, h.taskID)
}
// SetKillSignal allows overriding the signal sent to kill the task.
func (h *DriverHandle) SetKillSignal(signal string) {
h.killSignal = signal
}
func (h *DriverHandle) Kill() error {
return h.driver.StopTask(h.taskID, h.killTimeout, h.killSignal)
}
func (h *DriverHandle) Stats(ctx context.Context, interval time.Duration) (<-chan *cstructs.TaskResourceUsage, error) {
return h.driver.TaskStats(ctx, h.taskID, interval)
}
func (h *DriverHandle) Signal(s string) error {
return h.driver.SignalTask(h.taskID, s)
}
// Exec is the handled used by client endpoint handler to invoke the appropriate task driver exec.
func (h *DriverHandle) Exec(timeout time.Duration, cmd string, args []string) ([]byte, int, error) {
command := append([]string{cmd}, args...)
res, err := h.driver.ExecTask(h.taskID, command, timeout)
if err != nil {
return nil, 0, err
}
return res.Stdout, res.ExitResult.ExitCode, res.ExitResult.Err
}
// ExecStreaming is the handled used by client endpoint handler to invoke the appropriate task driver exec.
// while allowing to stream input and output
func (h *DriverHandle) ExecStreaming(ctx context.Context,
command []string,
tty bool,
stream drivers.ExecTaskStream) error {
if impl, ok := h.driver.(drivers.ExecTaskStreamingRawDriver); ok {
return impl.ExecTaskStreamingRaw(ctx, h.taskID, command, tty, stream)
}
d, ok := h.driver.(drivers.ExecTaskStreamingDriver)
if !ok {
return fmt.Errorf("task driver does not support exec")
}
execOpts, doneCh := drivers.StreamToExecOptions(
ctx, command, tty, stream)
result, err := d.ExecTaskStreaming(ctx, h.taskID, execOpts)
if err != nil {
return err
}
execOpts.Stdout.Close()
execOpts.Stderr.Close()
select {
case err = <-doneCh:
case <-ctx.Done():
err = fmt.Errorf("exec task timed out: %v", ctx.Err())
}
if err != nil {
return err
}
return stream.Send(drivers.NewExecStreamingResponseExit(result.ExitCode))
}
func (h *DriverHandle) Network() *drivers.DriverNetwork {
return h.net
}