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

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