github.com/jmitchell/nomad@v0.1.3-0.20151007230021-7ab84c2862d8/client/driver/rkt.go

package driver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	reRktVersion  = regexp.MustCompile("rkt version ([\\d\\.]+).+")
	reAppcVersion = regexp.MustCompile("appc version ([\\d\\.]+).+")
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
	DriverContext
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
	proc   *os.Process
	name   string
	logger *log.Logger
	waitCh chan error
	doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
	Pid  int
	Name string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
	return &RktDriver{*ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	// Only enable if we are root when running on non-windows systems.
	if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
		d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
		return false, nil
	}

	outBytes, err := exec.Command("rkt", "version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	rktMatches := reRktVersion.FindStringSubmatch(out)
	appcMatches := reRktVersion.FindStringSubmatch(out)
	if len(rktMatches) != 2 || len(appcMatches) != 2 {
		return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
	}

	node.Attributes["driver.rkt"] = "true"
	node.Attributes["driver.rkt.version"] = rktMatches[0]
	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]

	return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	trust_prefix, ok := task.Config["trust_prefix"]
	if !ok || trust_prefix == "" {
		return nil, fmt.Errorf("Missing trust prefix for rkt")
	}

	// Add the given trust prefix
	var outBuf, errBuf bytes.Buffer
	cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trust_prefix))
	cmd.Stdout = &outBuf
	cmd.Stderr = &errBuf
	d.logger.Printf("[DEBUG] driver.rkt: starting rkt command: %q", cmd.Args)
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf(
			"Error running rkt: %s\n\nOutput: %s\n\nError: %s",
			err, outBuf.String(), errBuf.String())
	}
	d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trust_prefix)

	name, ok := task.Config["name"]
	if !ok || name == "" {
		return nil, fmt.Errorf("Missing ACI name for rkt")
	}

	exec_cmd, ok := task.Config["exec"]
	if !ok || exec_cmd == "" {
		d.logger.Printf("[WARN] driver.rkt: could not find a command to execute in the ACI, the default command will be executed")
	}

	// Run the ACI
	var aoutBuf, aerrBuf bytes.Buffer
	run_cmd := []string{
		"rkt",
		"run",
		"--mds-register=false",
		name,
	}
	if exec_cmd != "" {
		splitted := strings.Fields(exec_cmd)
		run_cmd = append(run_cmd, "--exec=", splitted[0], "--")
		run_cmd = append(run_cmd, splitted[1:]...)
		run_cmd = append(run_cmd, "---")
	}
	acmd := exec.Command(run_cmd[0], run_cmd[1:]...)
	acmd.Stdout = &aoutBuf
	acmd.Stderr = &aerrBuf
	d.logger.Printf("[DEBUG] driver:rkt: starting rkt command: %q", acmd.Args)
	if err := acmd.Start(); err != nil {
		return nil, fmt.Errorf(
			"Error running rkt: %s\n\nOutput: %s\n\nError: %s",
			err, aoutBuf.String(), aerrBuf.String())
	}
	d.logger.Printf("[DEBUG] driver.rkt: started ACI: %q", name)
	h := &rktHandle{
		proc:   acmd.Process,
		name:   name,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan error, 1),
	}
	go h.run()
	return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
	qpid := &rktPID{}
	if err := json.Unmarshal(pidBytes, qpid); err != nil {
		return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
	}

	// Find the process
	proc, err := os.FindProcess(qpid.Pid)
	if proc == nil || err != nil {
		return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
	}

	// Return a driver handle
	h := &rktHandle{
		proc:   proc,
		name:   qpid.Name,
		logger: d.logger,
		doneCh: make(chan struct{}),
		waitCh: make(chan error, 1),
	}

	go h.run()
	return h, nil
}

func (h *rktHandle) ID() string {
	// Return a handle to the PID
	pid := &rktPID{
		Pid:  h.proc.Pid,
		Name: h.name,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		h.logger.Printf("[ERR] driver.rkt: failed to marshal rkt PID to JSON: %s", err)
	}
	return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan error {
	return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
func (h *rktHandle) Kill() error {
	h.proc.Signal(os.Interrupt)
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(5 * time.Second):
		return h.proc.Kill()
	}
}

func (h *rktHandle) run() {
	ps, err := h.proc.Wait()
	close(h.doneCh)
	if err != nil {
		h.waitCh <- err
	} else if !ps.Success() {
		h.waitCh <- fmt.Errorf("task exited with error")
	}
	close(h.waitCh)
}