driver/raw_exec: initial raw_exec implementation
This commit is contained in:
Nick Ethier
Michael Schurter
parent
bcc5c4a8bd
commit
5617f3615b
|
|
@ -0,0 +1,515 @@
|
||||||
|
package raw_exec
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strconv"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/hashicorp/consul-template/signals"
|
||||||
|
hclog "github.com/hashicorp/go-hclog"
|
||||||
|
plugin "github.com/hashicorp/go-plugin"
|
||||||
|
"github.com/hashicorp/nomad/client/driver/executor"
|
||||||
|
dstructs "github.com/hashicorp/nomad/client/driver/structs"
|
||||||
|
bbase "github.com/hashicorp/nomad/plugins/base"
|
||||||
|
"github.com/hashicorp/nomad/plugins/drivers/base"
|
||||||
|
"github.com/hashicorp/nomad/plugins/drivers/utils"
|
||||||
|
"github.com/hashicorp/nomad/plugins/shared/hclspec"
|
||||||
|
"golang.org/x/net/context"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
// pluginName is the name of the plugin
|
||||||
|
pluginName = "raw_exec"
|
||||||
|
|
||||||
|
// fingerprintPeriod is the interval at which the driver will send fingerprint responses
|
||||||
|
fingerprintPeriod = 30 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
pluginInfo = &bbase.PluginInfoResponse{
|
||||||
|
Type: bbase.PluginTypeDriver,
|
||||||
|
PluginApiVersion: "0.0.1",
|
||||||
|
PluginVersion: "0.1.0",
|
||||||
|
Name: pluginName,
|
||||||
|
}
|
||||||
|
|
||||||
|
configSpec = hclspec.NewObject(map[string]*hclspec.Spec{
|
||||||
|
"enabled": hclspec.NewDefault(
|
||||||
|
hclspec.NewAttr("enabled", "bool", false),
|
||||||
|
hclspec.NewLiteral("false"),
|
||||||
|
),
|
||||||
|
"no_cgroups": hclspec.NewDefault(
|
||||||
|
hclspec.NewAttr("no_cgroups", "bool", false),
|
||||||
|
hclspec.NewLiteral("false"),
|
||||||
|
),
|
||||||
|
})
|
||||||
|
|
||||||
|
taskConfigSpec = hclspec.NewObject(map[string]*hclspec.Spec{
|
||||||
|
"command": hclspec.NewAttr("command", "string", true),
|
||||||
|
"args": hclspec.NewAttr("command", "list(string)", false),
|
||||||
|
})
|
||||||
|
|
||||||
|
capabilities = &base.Capabilities{
|
||||||
|
SendSignals: true,
|
||||||
|
Exec: true,
|
||||||
|
FSIsolation: base.FSIsolationNone,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
// The RawExecDriver is a privileged version of the exec driver. It provides no
|
||||||
|
// resource isolation and just fork/execs. The Exec driver should be preferred
|
||||||
|
// and this should only be used when explicitly needed.
|
||||||
|
type RawExecDriver struct {
|
||||||
|
*utils.Eventer
|
||||||
|
config *Config
|
||||||
|
tasks *taskStore
|
||||||
|
|
||||||
|
// fingerprintCh is a channel which other funcs can send fingerprints to
|
||||||
|
// that will immediatly be sent
|
||||||
|
fingerprintCh chan *base.Fingerprint
|
||||||
|
|
||||||
|
stopCh chan struct{}
|
||||||
|
|
||||||
|
logger hclog.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
type Config struct {
|
||||||
|
// noCgroups tracks whether we should use a cgroup to manage the process
|
||||||
|
// tree
|
||||||
|
noCgroups bool `codec:"no_cgroups"`
|
||||||
|
|
||||||
|
// enabled is set to true to enable the raw_exec driver
|
||||||
|
enabled bool `codec:"enabled"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type TaskConfig struct {
|
||||||
|
Command string `codec:"command"`
|
||||||
|
Args []string `codec:"args"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewRawExecDriver(logger hclog.Logger) base.DriverPlugin {
|
||||||
|
stopCh := make(chan struct{})
|
||||||
|
return &RawExecDriver{
|
||||||
|
Eventer: utils.NewEventer(stopCh),
|
||||||
|
config: &Config{},
|
||||||
|
tasks: newTaskStore(),
|
||||||
|
fingerprintCh: make(chan *base.Fingerprint),
|
||||||
|
stopCh: stopCh,
|
||||||
|
logger: logger.Named(pluginName),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) PluginInfo() (*bbase.PluginInfoResponse, error) {
|
||||||
|
return pluginInfo, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) ConfigSchema() (*hclspec.Spec, error) {
|
||||||
|
return configSpec, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) SetConfig(data []byte) error {
|
||||||
|
var config Config
|
||||||
|
if err := bbase.MsgPackDecode(data, &config); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
r.config = &config
|
||||||
|
go r.fingerprintNow()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) Shutdown(ctx context.Context) error {
|
||||||
|
close(r.stopCh)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) TaskConfigSchema() (*hclspec.Spec, error) {
|
||||||
|
return taskConfigSpec, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) Capabilities() (*base.Capabilities, error) {
|
||||||
|
return capabilities, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) Fingerprint(ctx context.Context) (<-chan *base.Fingerprint, error) {
|
||||||
|
ch := make(chan *base.Fingerprint)
|
||||||
|
go r.fingerprint(ctx, ch)
|
||||||
|
return r.fingerprintCh, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) fingerprint(ctx context.Context, ch chan *base.Fingerprint) {
|
||||||
|
defer close(r.fingerprintCh)
|
||||||
|
|
||||||
|
ticker := time.NewTimer(0)
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
ticker.Reset(fingerprintPeriod)
|
||||||
|
go r.fingerprintNow()
|
||||||
|
case f := <-r.fingerprintCh:
|
||||||
|
ch <- f
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) fingerprintNow() {
|
||||||
|
if r.fingerprintCh == nil {
|
||||||
|
r.logger.Debug("fingerprint channel was nil, skipping fingerprint")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var health base.HealthState
|
||||||
|
var desc string
|
||||||
|
if r.config.enabled {
|
||||||
|
health = base.HealthStateHealthy
|
||||||
|
desc = "raw_exec enabled"
|
||||||
|
} else {
|
||||||
|
health = base.HealthStateUnhealthy
|
||||||
|
desc = "raw_exec not enabled"
|
||||||
|
}
|
||||||
|
r.fingerprintCh <- &base.Fingerprint{
|
||||||
|
Attributes: map[string]string{},
|
||||||
|
Health: health,
|
||||||
|
HealthDescription: desc,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) RecoverTask(handle *base.TaskHandle) error {
|
||||||
|
var taskState RawExecTaskState
|
||||||
|
|
||||||
|
err := handle.GetDriverState(&taskState)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
pluginConfig := &plugin.ClientConfig{
|
||||||
|
Reattach: taskState.ReattachConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
exec, pluginClient, err := utils.CreateExecutorWithConfig(pluginConfig, os.Stderr)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
h := &rawExecTaskHandle{
|
||||||
|
exec: exec,
|
||||||
|
pid: taskState.Pid,
|
||||||
|
pluginClient: pluginClient,
|
||||||
|
task: taskState.TaskConfig,
|
||||||
|
procState: base.TaskStateRunning,
|
||||||
|
startedAt: taskState.StartedAt,
|
||||||
|
}
|
||||||
|
|
||||||
|
r.tasks.Set(taskState.TaskConfig.ID, h)
|
||||||
|
|
||||||
|
go h.run()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) StartTask(cfg *base.TaskConfig) (*base.TaskHandle, error) {
|
||||||
|
if _, ok := r.tasks.Get(cfg.ID); ok {
|
||||||
|
return nil, fmt.Errorf("task with ID '%s' already started", cfg.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
var driverConfig TaskConfig
|
||||||
|
if err := cfg.DecodeDriverConfig(&driverConfig); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
handle := base.NewTaskHandle(pluginName)
|
||||||
|
handle.Config = cfg
|
||||||
|
|
||||||
|
// Get the command to be ran
|
||||||
|
command := driverConfig.Command
|
||||||
|
if err := utils.ValidateCommand(command, "args"); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
pluginLogFile := filepath.Join(cfg.TaskDir().Dir, "executor.out")
|
||||||
|
executorConfig := &dstructs.ExecutorConfig{
|
||||||
|
LogFile: pluginLogFile,
|
||||||
|
LogLevel: "debug",
|
||||||
|
}
|
||||||
|
|
||||||
|
exec, pluginClient, err := utils.CreateExecutor(os.Stderr, hclog.Debug, 14000, 14512, executorConfig)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
execCmd := &executor.ExecCommand{
|
||||||
|
Cmd: driverConfig.Command,
|
||||||
|
Args: driverConfig.Args,
|
||||||
|
Env: cfg.EnvList(),
|
||||||
|
User: cfg.User,
|
||||||
|
//TaskKillSignal: os.Interrupt,
|
||||||
|
BasicProcessCgroup: !r.config.noCgroups,
|
||||||
|
}
|
||||||
|
|
||||||
|
ps, err := exec.Launch(execCmd)
|
||||||
|
if err != nil {
|
||||||
|
pluginClient.Kill()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
h := &rawExecTaskHandle{
|
||||||
|
exec: exec,
|
||||||
|
pid: ps.Pid,
|
||||||
|
pluginClient: pluginClient,
|
||||||
|
task: cfg,
|
||||||
|
procState: base.TaskStateRunning,
|
||||||
|
startedAt: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
r.tasks.Set(cfg.ID, h)
|
||||||
|
|
||||||
|
handle.SetDriverState(RawExecTaskState{
|
||||||
|
ReattachConfig: pluginClient.ReattachConfig(),
|
||||||
|
Pid: ps.Pid,
|
||||||
|
TaskConfig: cfg,
|
||||||
|
StartedAt: h.startedAt,
|
||||||
|
})
|
||||||
|
|
||||||
|
go h.run()
|
||||||
|
return handle, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) WaitTask(ctx context.Context, taskID string) (<-chan *base.ExitResult, error) {
|
||||||
|
ch := make(chan *base.ExitResult)
|
||||||
|
handle, ok := r.tasks.Get(taskID)
|
||||||
|
if !ok {
|
||||||
|
return nil, base.ErrTaskNotFound
|
||||||
|
}
|
||||||
|
go r.handleWait(ctx, handle, ch)
|
||||||
|
|
||||||
|
return ch, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) handleWait(ctx context.Context, handle *rawExecTaskHandle, ch chan *base.ExitResult) {
|
||||||
|
defer close(ch)
|
||||||
|
var result *base.ExitResult
|
||||||
|
ps, err := handle.exec.Wait()
|
||||||
|
if err != nil {
|
||||||
|
result = &base.ExitResult{
|
||||||
|
Err: fmt.Errorf("executor: error waiting on process: %v", err),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
result = &base.ExitResult{
|
||||||
|
ExitCode: ps.ExitCode,
|
||||||
|
Signal: ps.Signal,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return
|
||||||
|
case ch <- result:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) StopTask(taskID string, timeout time.Duration, signal string) error {
|
||||||
|
/*
|
||||||
|
handle, ok := r.tasks.Get(taskID)
|
||||||
|
if !ok {
|
||||||
|
return base.ErrTaskNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
//TODO executor only supports shutting down with the initial signal provided
|
||||||
|
if err := handle.exec.Shutdown(); err != nil {
|
||||||
|
if handle.pluginClient.Exited() {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return fmt.Errorf("executor Shutdown failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-d.WaitTask(taskID):
|
||||||
|
return nil
|
||||||
|
case <-time.After(timeout):
|
||||||
|
if handle.pluginClient.Exited() {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err := handle.exec.Exit(); err != nil {
|
||||||
|
return fmt.Errorf("executor Exit failed: %v", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}*/
|
||||||
|
panic("not implemented")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) DestroyTask(taskID string, force bool) error {
|
||||||
|
handle, ok := r.tasks.Get(taskID)
|
||||||
|
if !ok {
|
||||||
|
return base.ErrTaskNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
if !handle.pluginClient.Exited() {
|
||||||
|
if err := handle.exec.Destroy(); err != nil {
|
||||||
|
handle.logger.Error("destroying executor failed", "err", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
handle.pluginClient.Kill()
|
||||||
|
r.tasks.Delete(taskID)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) InspectTask(taskID string) (*base.TaskStatus, error) {
|
||||||
|
handle, ok := r.tasks.Get(taskID)
|
||||||
|
if !ok {
|
||||||
|
return nil, base.ErrTaskNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
status := &base.TaskStatus{
|
||||||
|
ID: handle.task.ID,
|
||||||
|
Name: handle.task.Name,
|
||||||
|
State: handle.procState,
|
||||||
|
SizeOnDiskMB: 0,
|
||||||
|
StartedAt: handle.startedAt,
|
||||||
|
CompletedAt: handle.completedAt,
|
||||||
|
ExitResult: handle.exitResult,
|
||||||
|
DriverAttributes: map[string]string{
|
||||||
|
"pid": strconv.Itoa(handle.pid),
|
||||||
|
},
|
||||||
|
NetworkOverride: &base.NetworkOverride{},
|
||||||
|
}
|
||||||
|
|
||||||
|
return status, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) TaskStats(taskID string) (*base.TaskStats, error) {
|
||||||
|
handle, ok := r.tasks.Get(taskID)
|
||||||
|
if !ok {
|
||||||
|
return nil, base.ErrTaskNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
stats, err := handle.exec.Stats()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return &base.TaskStats{
|
||||||
|
ID: handle.task.ID,
|
||||||
|
Timestamp: stats.Timestamp,
|
||||||
|
AggResourceUsage: stats.ResourceUsage,
|
||||||
|
ResourceUsageByPid: stats.Pids,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) SignalTask(taskID string, signal string) error {
|
||||||
|
handle, ok := r.tasks.Get(taskID)
|
||||||
|
if !ok {
|
||||||
|
return base.ErrTaskNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
sig := os.Interrupt
|
||||||
|
if s, ok := signals.SignalLookup[signal]; ok {
|
||||||
|
sig = s
|
||||||
|
}
|
||||||
|
return handle.exec.Signal(sig)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *RawExecDriver) ExecTask(taskID string, cmd []string, timeout time.Duration) (*base.ExecTaskResult, error) {
|
||||||
|
if len(cmd) == 0 {
|
||||||
|
return nil, fmt.Errorf("error cmd must have atleast one value")
|
||||||
|
}
|
||||||
|
handle, ok := r.tasks.Get(taskID)
|
||||||
|
if !ok {
|
||||||
|
return nil, base.ErrTaskNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
args := []string{}
|
||||||
|
if len(cmd) > 1 {
|
||||||
|
args = cmd[1:]
|
||||||
|
}
|
||||||
|
|
||||||
|
out, exitCode, err := handle.exec.Exec(time.Now().Add(timeout), cmd[0], args)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return &base.ExecTaskResult{
|
||||||
|
Stdout: out,
|
||||||
|
ExitResult: &base.ExitResult{
|
||||||
|
ExitCode: exitCode,
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type taskStore struct {
|
||||||
|
store map[string]*rawExecTaskHandle
|
||||||
|
lock sync.RWMutex
|
||||||
|
}
|
||||||
|
|
||||||
|
func newTaskStore() *taskStore {
|
||||||
|
return &taskStore{store: map[string]*rawExecTaskHandle{}}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (ts *taskStore) Set(id string, handle *rawExecTaskHandle) {
|
||||||
|
ts.lock.Lock()
|
||||||
|
defer ts.lock.Unlock()
|
||||||
|
ts.store[id] = handle
|
||||||
|
}
|
||||||
|
|
||||||
|
func (ts *taskStore) Get(id string) (*rawExecTaskHandle, bool) {
|
||||||
|
ts.lock.RLock()
|
||||||
|
defer ts.lock.RUnlock()
|
||||||
|
t, ok := ts.store[id]
|
||||||
|
return t, ok
|
||||||
|
}
|
||||||
|
|
||||||
|
func (ts *taskStore) Delete(id string) {
|
||||||
|
ts.lock.Lock()
|
||||||
|
defer ts.lock.Unlock()
|
||||||
|
delete(ts.store, id)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (ts *taskStore) Range(f func(id string, handle *rawExecTaskHandle) bool) {
|
||||||
|
ts.lock.RLock()
|
||||||
|
defer ts.lock.RUnlock()
|
||||||
|
for k, v := range ts.store {
|
||||||
|
if f(k, v) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type RawExecTaskState struct {
|
||||||
|
ReattachConfig *plugin.ReattachConfig
|
||||||
|
TaskConfig *base.TaskConfig
|
||||||
|
Pid int
|
||||||
|
StartedAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
type rawExecTaskHandle struct {
|
||||||
|
exec executor.Executor
|
||||||
|
pid int
|
||||||
|
pluginClient *plugin.Client
|
||||||
|
task *base.TaskConfig
|
||||||
|
procState base.TaskState
|
||||||
|
startedAt time.Time
|
||||||
|
completedAt time.Time
|
||||||
|
exitResult *base.ExitResult
|
||||||
|
logger hclog.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *rawExecTaskHandle) run() {
|
||||||
|
if h.exitResult == nil {
|
||||||
|
h.exitResult = &base.ExitResult{}
|
||||||
|
}
|
||||||
|
|
||||||
|
ps, err := h.exec.Wait()
|
||||||
|
if err != nil {
|
||||||
|
h.exitResult.Err = err
|
||||||
|
h.procState = base.TaskStateUnknown
|
||||||
|
h.completedAt = time.Now()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
h.procState = base.TaskStateExited
|
||||||
|
h.exitResult.ExitCode = ps.ExitCode
|
||||||
|
h.exitResult.Signal = ps.Signal
|
||||||
|
h.completedAt = ps.Time
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,37 @@
|
||||||
|
package raw_exec
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/hashicorp/nomad/helper/testlog"
|
||||||
|
"github.com/hashicorp/nomad/helper/uuid"
|
||||||
|
"github.com/hashicorp/nomad/plugins/drivers/base"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDriverStartTask(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
require := require.New(t)
|
||||||
|
|
||||||
|
d := NewRawExecDriver(testlog.HCLogger(t))
|
||||||
|
harness := base.NewDriverHarness(t, d)
|
||||||
|
task := &base.TaskConfig{
|
||||||
|
ID: uuid.Generate(),
|
||||||
|
Name: "test",
|
||||||
|
}
|
||||||
|
task.EncodeDriverConfig(&TaskConfig{
|
||||||
|
Command: "go",
|
||||||
|
Args: []string{"version"},
|
||||||
|
})
|
||||||
|
cleanup := harness.MkAllocDir(task)
|
||||||
|
defer cleanup()
|
||||||
|
|
||||||
|
handle, err := harness.StartTask(task)
|
||||||
|
require.NoError(err)
|
||||||
|
|
||||||
|
ch, err := harness.WaitTask(context.Background(), handle.Config.ID)
|
||||||
|
require.NoError(err)
|
||||||
|
result := <-ch
|
||||||
|
require.Zero(result.ExitCode)
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,101 @@
|
||||||
|
package utils
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/hashicorp/nomad/plugins/drivers/base"
|
||||||
|
"golang.org/x/net/context"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
//DefaultSendEventTimeout is the timeout used when publishing events to consumers
|
||||||
|
DefaultSendEventTimeout = 2 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Eventer is a utility to control broadcast of TaskEvents to multiple consumers.
|
||||||
|
// It also implements the TaskStats func in the DriverPlugin interface so that
|
||||||
|
// it can be embedded in a implementing driver struct.
|
||||||
|
type Eventer struct {
|
||||||
|
sync.RWMutex
|
||||||
|
|
||||||
|
// events is a channel were events to be broadcasted are sent
|
||||||
|
events chan *base.TaskEvent
|
||||||
|
|
||||||
|
// streamers is a slice of consumers to broadcast events to
|
||||||
|
// access is gaurded by RWMutex
|
||||||
|
streamers []*eventStreamer
|
||||||
|
|
||||||
|
// stop chan to allow control of event loop shutdown
|
||||||
|
stop chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewEventer returns an Eventer with a running event loop that can be stopped
|
||||||
|
// by closing the given stop channel
|
||||||
|
func NewEventer(stop chan struct{}) *Eventer {
|
||||||
|
e := &Eventer{
|
||||||
|
events: make(chan *base.TaskEvent),
|
||||||
|
stop: stop,
|
||||||
|
}
|
||||||
|
go e.eventLoop()
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
|
||||||
|
// eventLoop is the main logic which pulls events from the channel and broadcasts
|
||||||
|
// them to all consumers
|
||||||
|
func (e *Eventer) eventLoop() {
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-e.stop:
|
||||||
|
for _, stream := range e.streamers {
|
||||||
|
close(stream.ch)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
case event := <-e.events:
|
||||||
|
e.RLock()
|
||||||
|
for _, stream := range e.streamers {
|
||||||
|
stream.send(event)
|
||||||
|
}
|
||||||
|
e.RUnlock()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type eventStreamer struct {
|
||||||
|
timeout time.Duration
|
||||||
|
ctx context.Context
|
||||||
|
ch chan *base.TaskEvent
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *eventStreamer) send(event *base.TaskEvent) {
|
||||||
|
select {
|
||||||
|
case <-time.After(s.timeout):
|
||||||
|
case <-s.ctx.Done():
|
||||||
|
case s.ch <- event:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *Eventer) newStream(ctx context.Context) <-chan *base.TaskEvent {
|
||||||
|
e.Lock()
|
||||||
|
defer e.Unlock()
|
||||||
|
|
||||||
|
stream := &eventStreamer{
|
||||||
|
ch: make(chan *base.TaskEvent),
|
||||||
|
ctx: ctx,
|
||||||
|
timeout: DefaultSendEventTimeout,
|
||||||
|
}
|
||||||
|
e.streamers = append(e.streamers, stream)
|
||||||
|
|
||||||
|
return stream.ch
|
||||||
|
}
|
||||||
|
|
||||||
|
// TaskEvents is an implementation of the DriverPlugin.TaskEvents function
|
||||||
|
func (e *Eventer) TaskEvents(ctx context.Context) (<-chan *base.TaskEvent, error) {
|
||||||
|
stream := e.newStream(ctx)
|
||||||
|
return stream, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// EmitEvent can be used to broadcast a new event
|
||||||
|
func (e *Eventer) EmitEvent(event *base.TaskEvent) {
|
||||||
|
e.events <- event
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,72 @@
|
||||||
|
package utils
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/hashicorp/nomad/plugins/drivers/base"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestEventer(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
require := require.New(t)
|
||||||
|
|
||||||
|
stop := make(chan struct{})
|
||||||
|
e := NewEventer(stop)
|
||||||
|
|
||||||
|
events := []*base.TaskEvent{
|
||||||
|
{
|
||||||
|
TaskID: "a",
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
TaskID: "b",
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
TaskID: "c",
|
||||||
|
Timestamp: time.Now(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
consumer1, err := e.TaskEvents(context.Background())
|
||||||
|
require.NoError(err)
|
||||||
|
consumer2, err := e.TaskEvents(context.Background())
|
||||||
|
require.NoError(err)
|
||||||
|
|
||||||
|
var buffer1, buffer2 []*base.TaskEvent
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(2)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
for {
|
||||||
|
event, ok := <-consumer1
|
||||||
|
if !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
buffer1 = append(buffer1, event)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
for {
|
||||||
|
event, ok := <-consumer2
|
||||||
|
if !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
buffer2 = append(buffer2, event)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
for _, event := range events {
|
||||||
|
e.EmitEvent(event)
|
||||||
|
}
|
||||||
|
|
||||||
|
close(stop)
|
||||||
|
wg.Wait()
|
||||||
|
require.Exactly(events, buffer1)
|
||||||
|
require.Exactly(events, buffer2)
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue