github.com/ryanslade/nomad@v0.2.4-0.20160128061903-fc95782f2089/client/driver/raw_exec.go (about)

     1  package driver
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"log"
     7  	"path/filepath"
     8  	"time"
     9  
    10  	"github.com/hashicorp/nomad/client/allocdir"
    11  	"github.com/hashicorp/nomad/client/config"
    12  	"github.com/hashicorp/nomad/client/driver/executor"
    13  	cstructs "github.com/hashicorp/nomad/client/driver/structs"
    14  	"github.com/hashicorp/nomad/client/fingerprint"
    15  	"github.com/hashicorp/nomad/client/getter"
    16  	"github.com/hashicorp/nomad/nomad/structs"
    17  	"github.com/mitchellh/mapstructure"
    18  )
    19  
    20  const (
    21  	// The option that enables this driver in the Config.Options map.
    22  	rawExecConfigOption = "driver.raw_exec.enable"
    23  )
    24  
    25  // The RawExecDriver is a privileged version of the exec driver. It provides no
    26  // resource isolation and just fork/execs. The Exec driver should be preferred
    27  // and this should only be used when explicitly needed.
    28  type RawExecDriver struct {
    29  	DriverContext
    30  	fingerprint.StaticFingerprinter
    31  }
    32  
    33  // rawExecHandle is returned from Start/Open as a handle to the PID
    34  type rawExecHandle struct {
    35  	cmd         executor.Executor
    36  	killTimeout time.Duration
    37  	logger      *log.Logger
    38  	waitCh      chan *cstructs.WaitResult
    39  	doneCh      chan struct{}
    40  }
    41  
    42  // NewRawExecDriver is used to create a new raw exec driver
    43  func NewRawExecDriver(ctx *DriverContext) Driver {
    44  	return &RawExecDriver{DriverContext: *ctx}
    45  }
    46  
    47  func (d *RawExecDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
    48  	// Check that the user has explicitly enabled this executor.
    49  	enabled := cfg.ReadBoolDefault(rawExecConfigOption, false)
    50  
    51  	if enabled {
    52  		d.logger.Printf("[WARN] driver.raw_exec: raw exec is enabled. Only enable if needed")
    53  		node.Attributes["driver.raw_exec"] = "1"
    54  		return true, nil
    55  	}
    56  
    57  	return false, nil
    58  }
    59  
    60  func (d *RawExecDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
    61  	var driverConfig ExecDriverConfig
    62  	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
    63  		return nil, err
    64  	}
    65  	// Get the tasks local directory.
    66  	taskName := d.DriverContext.taskName
    67  	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
    68  	if !ok {
    69  		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
    70  	}
    71  
    72  	// Get the command to be ran
    73  	command := driverConfig.Command
    74  	if command == "" {
    75  		return nil, fmt.Errorf("missing command for Raw Exec driver")
    76  	}
    77  
    78  	// Check if an artificat is specified and attempt to download it
    79  	source, ok := task.Config["artifact_source"]
    80  	if ok && source != "" {
    81  		// Proceed to download an artifact to be executed.
    82  		_, err := getter.GetArtifact(
    83  			filepath.Join(taskDir, allocdir.TaskLocal),
    84  			driverConfig.ArtifactSource,
    85  			driverConfig.Checksum,
    86  			d.logger,
    87  		)
    88  		if err != nil {
    89  			return nil, err
    90  		}
    91  	}
    92  
    93  	// Setup the command
    94  	execCtx := executor.NewExecutorContext(d.taskEnv)
    95  	cmd := executor.NewBasicExecutor(execCtx)
    96  	executor.SetCommand(cmd, command, driverConfig.Args)
    97  	if err := cmd.Limit(task.Resources); err != nil {
    98  		return nil, fmt.Errorf("failed to constrain resources: %s", err)
    99  	}
   100  
   101  	// Populate environment variables
   102  	cmd.Command().Env = d.taskEnv.EnvList()
   103  
   104  	if err := cmd.ConfigureTaskDir(d.taskName, ctx.AllocDir); err != nil {
   105  		return nil, fmt.Errorf("failed to configure task directory: %v", err)
   106  	}
   107  
   108  	if err := cmd.Start(); err != nil {
   109  		return nil, fmt.Errorf("failed to start command: %v", err)
   110  	}
   111  
   112  	// Return a driver handle
   113  	h := &execHandle{
   114  		cmd:         cmd,
   115  		killTimeout: d.DriverContext.KillTimeout(task),
   116  		logger:      d.logger,
   117  		doneCh:      make(chan struct{}),
   118  		waitCh:      make(chan *cstructs.WaitResult, 1),
   119  	}
   120  	go h.run()
   121  	return h, nil
   122  }
   123  
   124  type rawExecId struct {
   125  	ExecutorId  string
   126  	KillTimeout time.Duration
   127  }
   128  
   129  func (d *RawExecDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
   130  	id := &rawExecId{}
   131  	if err := json.Unmarshal([]byte(handleID), id); err != nil {
   132  		return nil, fmt.Errorf("Failed to parse handle '%s': %v", handleID, err)
   133  	}
   134  
   135  	// Find the process
   136  	execCtx := executor.NewExecutorContext(d.taskEnv)
   137  	cmd := executor.NewBasicExecutor(execCtx)
   138  	if err := cmd.Open(id.ExecutorId); err != nil {
   139  		return nil, fmt.Errorf("failed to open ID %v: %v", id.ExecutorId, err)
   140  	}
   141  
   142  	// Return a driver handle
   143  	h := &execHandle{
   144  		cmd:         cmd,
   145  		logger:      d.logger,
   146  		killTimeout: id.KillTimeout,
   147  		doneCh:      make(chan struct{}),
   148  		waitCh:      make(chan *cstructs.WaitResult, 1),
   149  	}
   150  	go h.run()
   151  	return h, nil
   152  }
   153  
   154  func (h *rawExecHandle) ID() string {
   155  	executorId, _ := h.cmd.ID()
   156  	id := rawExecId{
   157  		ExecutorId:  executorId,
   158  		KillTimeout: h.killTimeout,
   159  	}
   160  
   161  	data, err := json.Marshal(id)
   162  	if err != nil {
   163  		h.logger.Printf("[ERR] driver.raw_exec: failed to marshal ID to JSON: %s", err)
   164  	}
   165  	return string(data)
   166  }
   167  
   168  func (h *rawExecHandle) WaitCh() chan *cstructs.WaitResult {
   169  	return h.waitCh
   170  }
   171  
   172  func (h *rawExecHandle) Update(task *structs.Task) error {
   173  	// Update is not possible
   174  	return nil
   175  }
   176  
   177  func (h *rawExecHandle) Kill() error {
   178  	h.cmd.Shutdown()
   179  	select {
   180  	case <-h.doneCh:
   181  		return nil
   182  	case <-time.After(h.killTimeout):
   183  		return h.cmd.ForceStop()
   184  	}
   185  }
   186  
   187  func (h *rawExecHandle) run() {
   188  	res := h.cmd.Wait()
   189  	close(h.doneCh)
   190  	h.waitCh <- res
   191  	close(h.waitCh)
   192  }