Merge pull request #610 from achanda/isolators

Add support for CPU and memory isolators

Author: dadgar

client/driver/rkt.go

@@ -14,6 +14,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/hashicorp/go-version"
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/config"
 	cstructs "github.com/hashicorp/nomad/client/driver/structs"
@@ -28,6 +29,13 @@ var (
 	reAppcVersion = regexp.MustCompile(`appc version (\d[.\d]+)`)
 )
 
+const (
+	// rkt added support for CPU and memory isolators in 0.14.0. We cannot support
+	// an earlier version to maintain an uniform interface across all drivers
+	minRktVersion    = "0.14.0"
+	conversionFactor = 1024 * 1024
+)
+
 // 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
@@ -85,6 +93,13 @@ func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, e
 	node.Attributes["driver.rkt.version"] = rktMatches[1]
 	node.Attributes["driver.rkt.appc.version"] = appcMatches[1]
 
+	minVersion, _ := version.NewVersion(minRktVersion)
+	currentVersion, _ := version.NewVersion(node.Attributes["driver.rkt.version"])
+	if currentVersion.LessThan(minVersion) {
+		// Do not allow rkt < 0.14.0
+		d.logger.Printf("[WARN] driver.rkt: please upgrade rkt to a version >= %s", minVersion)
+		node.Attributes["driver.rkt"] = "0"
+	}
 	return true, nil
 }
 
@@ -109,21 +124,21 @@ func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, e
 	taskLocal := filepath.Join(taskDir, allocdir.TaskLocal)
 
 	// Add the given trust prefix
-	trust_prefix, trust_cmd := task.Config["trust_prefix"]
-	if trust_cmd {
+	trustPrefix, trustCmd := task.Config["trust_prefix"]
+	if trustCmd {
 		var outBuf, errBuf bytes.Buffer
-		cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trust_prefix))
+		cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trustPrefix))
 		cmd.Stdout = &outBuf
 		cmd.Stderr = &errBuf
 		if err := cmd.Run(); err != nil {
 			return nil, fmt.Errorf("Error running rkt trust: %s\n\nOutput: %s\n\nError: %s",
 				err, outBuf.String(), errBuf.String())
 		}
-		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trust_prefix)
+		d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trustPrefix)
 	}
 
 	// Build the command.
-	var cmd_args []string
+	var cmdArgs []string
 
 	// Inject the environment variables.
 	envVars := TaskEnvironmentVariables(ctx, task)
@@ -133,33 +148,46 @@ func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, e
 	envVars.ClearAllocDir()
 
 	for k, v := range envVars.Map() {
-		cmd_args = append(cmd_args, fmt.Sprintf("--set-env=%v=%v", k, v))
+		cmdArgs = append(cmdArgs, fmt.Sprintf("--set-env=%v=%v", k, v))
 	}
 
 	// Disble signature verification if the trust command was not run.
-	if !trust_cmd {
-		cmd_args = append(cmd_args, "--insecure-skip-verify")
+	if !trustCmd {
+		cmdArgs = append(cmdArgs, "--insecure-skip-verify")
 	}
 
 	// Append the run command.
-	cmd_args = append(cmd_args, "run", "--mds-register=false", img)
+	cmdArgs = append(cmdArgs, "run", "--mds-register=false", img)
 
 	// Check if the user has overriden the exec command.
-	if exec_cmd, ok := task.Config["command"]; ok {
-		cmd_args = append(cmd_args, fmt.Sprintf("--exec=%v", exec_cmd))
+	if execCmd, ok := task.Config["command"]; ok {
+		cmdArgs = append(cmdArgs, fmt.Sprintf("--exec=%v", execCmd))
+	}
+
+	if task.Resources.MemoryMB == 0 {
+		return nil, fmt.Errorf("Memory limit cannot be zero")
+	}
+	if task.Resources.CPU == 0 {
+		return nil, fmt.Errorf("CPU limit cannot be zero")
 	}
 
+	// Add memory isolator
+	cmdArgs = append(cmdArgs, fmt.Sprintf("--memory=%vM", int64(task.Resources.MemoryMB)*conversionFactor))
+
+	// Add CPU isolator
+	cmdArgs = append(cmdArgs, fmt.Sprintf("--cpu=%vm", int64(task.Resources.CPU)))
+
 	// Add user passed arguments.
 	if len(driverConfig.Args) != 0 {
 		parsed := args.ParseAndReplace(driverConfig.Args, envVars.Map())
 
 		// Need to start arguments with "--"
 		if len(parsed) > 0 {
-			cmd_args = append(cmd_args, "--")
+			cmdArgs = append(cmdArgs, "--")
 		}
 
 		for _, arg := range parsed {
-			cmd_args = append(cmd_args, fmt.Sprintf("%v", arg))
+			cmdArgs = append(cmdArgs, fmt.Sprintf("%v", arg))
 		}
 	}
 
@@ -177,7 +205,7 @@ func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, e
 		return nil, fmt.Errorf("Error opening file to redirect stderr: %v", err)
 	}
 
-	cmd := exec.Command("rkt", cmd_args...)
+	cmd := exec.Command("rkt", cmdArgs...)
 	cmd.Stdout = stdo
 	cmd.Stderr = stde
 

client/driver/rkt_test.go

@@ -81,6 +81,10 @@ func TestRktDriver_Start(t *testing.T) {
 			"image":        "coreos.com/etcd:v2.0.4",
 			"command":      "/etcd",
 		},
+		Resources: &structs.Resources{
+			MemoryMB: 256,
+			CPU:      512,
+		},
 	}
 
 	driverCtx := testDriverContext(task.Name)
@@ -121,6 +125,10 @@ func TestRktDriver_Start_Wait(t *testing.T) {
 			"command":      "/etcd",
 			"args":         []string{"--version"},
 		},
+		Resources: &structs.Resources{
+			MemoryMB: 256,
+			CPU:      512,
+		},
 	}
 
 	driverCtx := testDriverContext(task.Name)
@@ -162,6 +170,10 @@ func TestRktDriver_Start_Wait_Skip_Trust(t *testing.T) {
 			"command": "/etcd",
 			"args":    []string{"--version"},
 		},
+		Resources: &structs.Resources{
+			MemoryMB: 256,
+			CPU:      512,
+		},
 	}
 
 	driverCtx := testDriverContext(task.Name)
@@ -204,6 +216,10 @@ func TestRktDriver_Start_Wait_Logs(t *testing.T) {
 			"command":      "/etcd",
 			"args":         []string{"--version"},
 		},
+		Resources: &structs.Resources{
+			MemoryMB: 256,
+			CPU:      512,
+		},
 	}
 
 	driverCtx := testDriverContext(task.Name)