github.com/huiliang/nomad@v0.2.1-0.20151124023127-7a8b664699ff/client/driver/exec.go (about)

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