Merge pull request #809 from remind101/run-constraints

Set constraints for one off tasks

CHANGELOG.md

@@ -12,6 +12,7 @@
 * Stdout and Stdin from interactive run sessions can now be sent to CloudWatch Logs for longterm storage and auditing [#757](https://github.com/remind101/empire/pull/757).
 * Add `Environment` and `Release` to Deploy Events. `--environment` will likely be used for tagging resources later. [#758](https://github.com/remind101/empire/pull/758)
 * Add constraint changes to scale events [#773](https://github.com/remind101/empire/pull/773)
+* You can now specify the CPU and memory constraints for attached one-off tasks with the `-s` flag to `emp run` [#809](https://github.com/remind101/empire/pull/809)
 
 **Bugs**
 
@@ -20,6 +21,7 @@
 * Fixed a bug where it was previously possible to create a signed access token with an empty username [#780](https://github.com/remind101/empire/pull/780)
 * ECR authentication now supports multiple regions, and works independently of ECS region [#784](https://github.com/remind101/empire/pull/784)
 * Provisioned ELB's are only destroyed when the entire app is removed [#801](https://github.com/remind101/empire/pull/801)
+* Docker containers started by attached runs now have labels, cpu and memory constraints applied to them [#809](https://github.com/remind101/empire/pull/809)
 
 **Performance**
 

empire.go

@@ -362,6 +362,9 @@ type RunOpts struct {
 
 	// Extra environment variables to set.
 	Env map[string]string
+
+	// Optional memory/cpu/nproc constraints.
+	Constraints *Constraints
 }
 
 func (opts RunOpts) Event() RunEvent {

pkg/runner/runner.go

@@ -29,6 +29,13 @@ type RunOpts struct {
 	// Environment variables to set.
 	Env map[string]string
 
+	// Labels to set
+	Labels map[string]string
+
+	// Memory/CPUShares.
+	Memory    int64
+	CPUShares int64
+
 	// Streams fo Stdout, Stderr and Stdin.
 	Input  io.Reader
 	Output io.Writer
@@ -86,9 +93,12 @@ func (r *Runner) create(ctx context.Context, opts RunOpts) (*docker.Container, e
 			AttachStdout: true,
 			AttachStderr: true,
 			OpenStdin:    true,
+			Memory:       opts.Memory,
+			CPUShares:    opts.CPUShares,
 			Image:        opts.Image.String(),
 			Cmd:          opts.Command,
 			Env:          envKeys(opts.Env),
+			Labels:       opts.Labels,
 		},
 		HostConfig: &docker.HostConfig{
 			LogConfig: docker.LogConfig{

processes.go

@@ -26,6 +26,16 @@ func ParseCommand(command string) (Command, error) {
 	return shellwords.Parse(command)
 }
 
+// MustParseCommand parses the string into a Command, panicing if there's an
+// error. This method should only be used in tests for convenience.
+func MustParseCommand(command string) Command {
+	c, err := ParseCommand(command)
+	if err != nil {
+		panic(err)
+	}
+	return c
+}
+
 // Scan implements the sql.Scanner interface.
 func (c *Command) Scan(src interface{}) error {
 	bytes, ok := src.([]byte)

runner.go

@@ -62,12 +62,19 @@ func (r *runnerService) Run(ctx context.Context, opts RunOpts) error {
 		return err
 	}
 
+	proc := Process{Command: opts.Command, Quantity: 1}
+
+	// Set the size of the process.
+	constraints := DefaultConstraints
+	if opts.Constraints != nil {
+		constraints = *opts.Constraints
+	}
+	proc.SetConstraints(constraints)
+
 	a := newServiceApp(release)
-	p := newServiceProcess(release, "run", Process{
-		Command:  opts.Command,
-		Quantity: 1,
-	})
+	p := newServiceProcess(release, "run", proc)
 
+	// Add additional environment variables to the process.
 	for k, v := range opts.Env {
 		p.Env[k] = v
 	}

scheduler/runner.go

@@ -18,11 +18,14 @@ func (m *AttachedRunner) Run(ctx context.Context, app *App, p *Process, in io.Re
 	// If an output stream is provided, run using the docker runner.
 	if out != nil {
 		return m.Runner.Run(ctx, runner.RunOpts{
-			Image:   p.Image,
-			Command: p.Command,
-			Env:     p.Env,
-			Input:   in,
-			Output:  out,
+			Image:     p.Image,
+			Command:   p.Command,
+			Env:       p.Env,
+			Memory:    int64(p.MemoryLimit),
+			CPUShares: int64(p.CPUShares),
+			Labels:    p.Labels,
+			Input:     in,
+			Output:    out,
 		})
 	}
 

server/heroku/processes.go

@@ -58,10 +58,10 @@ func (h *GetProcesses) ServeHTTPContext(ctx context.Context, w http.ResponseWrit
 }
 
 type PostProcessForm struct {
-	Command string            `json:"command"`
-	Attach  bool              `json:"attach"`
-	Env     map[string]string `json:"env"`
-	Size    string            `json:"size"`
+	Command string              `json:"command"`
+	Attach  bool                `json:"attach"`
+	Env     map[string]string   `json:"env"`
+	Size    *empire.Constraints `json:"size"`
 }
 
 type PostProcess struct {
@@ -91,11 +91,12 @@ func (h *PostProcess) ServeHTTPContext(ctx context.Context, w http.ResponseWrite
 	}
 
 	opts := empire.RunOpts{
-		User:    UserFromContext(ctx),
-		App:     a,
-		Command: command,
-		Env:     form.Env,
-		Message: m,
+		User:        UserFromContext(ctx),
+		App:         a,
+		Command:     command,
+		Env:         form.Env,
+		Constraints: form.Size,
+		Message:     m,
 	}
 
 	if form.Attach {

tests/empire/empire_test.go

@@ -235,6 +235,149 @@ func TestEmpire_Deploy_Concurrent(t *testing.T) {
 	s.AssertExpectations(t)
 }
 
+func TestEmpire_Run(t *testing.T) {
+	e := empiretest.NewEmpire(t)
+
+	user := &empire.User{Name: "ejholmes"}
+
+	app, err := e.Create(context.Background(), empire.CreateOpts{
+		User: user,
+		Name: "acme-inc",
+	})
+	assert.NoError(t, err)
+
+	img := image.Image{Repository: "remind101/acme-inc"}
+	_, err = e.Deploy(context.Background(), empire.DeploymentsCreateOpts{
+		App:    app,
+		User:   user,
+		Output: ioutil.Discard,
+		Image:  img,
+	})
+	assert.NoError(t, err)
+
+	s := new(mockScheduler)
+	e.Scheduler = s
+
+	s.On("Run", &scheduler.App{
+		ID:   app.ID,
+		Name: "acme-inc",
+	},
+		&scheduler.Process{
+			Type:        "run",
+			Image:       img,
+			Command:     []string{"bundle", "exec", "rake", "db:migrate"},
+			Instances:   1,
+			MemoryLimit: 536870912,
+			CPUShares:   256,
+			Nproc:       256,
+			Env: map[string]string{
+				"EMPIRE_APPID":      app.ID,
+				"EMPIRE_APPNAME":    "acme-inc",
+				"EMPIRE_PROCESS":    "run",
+				"EMPIRE_RELEASE":    "v1",
+				"SOURCE":            "acme-inc.run.v1",
+				"EMPIRE_CREATED_AT": "2015-01-01T01:01:01Z",
+				"TERM":              "xterm",
+			},
+			Labels: map[string]string{
+				"empire.app.name":    "acme-inc",
+				"empire.app.id":      app.ID,
+				"empire.app.process": "run",
+				"empire.app.release": "v1",
+			},
+		}, nil, nil).Return(nil)
+
+	err = e.Run(context.Background(), empire.RunOpts{
+		User:    user,
+		App:     app,
+		Command: empire.MustParseCommand("bundle exec rake db:migrate"),
+
+		// Detached Process
+		Output: nil,
+		Input:  nil,
+
+		Env: map[string]string{
+			"TERM": "xterm",
+		},
+	})
+	assert.NoError(t, err)
+
+	s.AssertExpectations(t)
+}
+
+func TestEmpire_Run_WithConstraints(t *testing.T) {
+	e := empiretest.NewEmpire(t)
+
+	user := &empire.User{Name: "ejholmes"}
+
+	app, err := e.Create(context.Background(), empire.CreateOpts{
+		User: user,
+		Name: "acme-inc",
+	})
+	assert.NoError(t, err)
+
+	img := image.Image{Repository: "remind101/acme-inc"}
+	_, err = e.Deploy(context.Background(), empire.DeploymentsCreateOpts{
+		App:    app,
+		User:   user,
+		Output: ioutil.Discard,
+		Image:  img,
+	})
+	assert.NoError(t, err)
+
+	s := new(mockScheduler)
+	e.Scheduler = s
+
+	s.On("Run", &scheduler.App{
+		ID:   app.ID,
+		Name: "acme-inc",
+	},
+		&scheduler.Process{
+			Type:        "run",
+			Image:       img,
+			Command:     []string{"bundle", "exec", "rake", "db:migrate"},
+			Instances:   1,
+			MemoryLimit: 1073741824,
+			CPUShares:   512,
+			Nproc:       512,
+			Env: map[string]string{
+				"EMPIRE_APPID":      app.ID,
+				"EMPIRE_APPNAME":    "acme-inc",
+				"EMPIRE_PROCESS":    "run",
+				"EMPIRE_RELEASE":    "v1",
+				"SOURCE":            "acme-inc.run.v1",
+				"EMPIRE_CREATED_AT": "2015-01-01T01:01:01Z",
+				"TERM":              "xterm",
+			},
+			Labels: map[string]string{
+				"empire.app.name":    "acme-inc",
+				"empire.app.id":      app.ID,
+				"empire.app.process": "run",
+				"empire.app.release": "v1",
+			},
+		}, nil, nil).Return(nil)
+
+	constraints := empire.NamedConstraints["2X"]
+	err = e.Run(context.Background(), empire.RunOpts{
+		User:    user,
+		App:     app,
+		Command: empire.MustParseCommand("bundle exec rake db:migrate"),
+
+		// Detached Process
+		Output: nil,
+		Input:  nil,
+
+		Env: map[string]string{
+			"TERM": "xterm",
+		},
+
+		Constraints: &constraints,
+	})
+	assert.NoError(t, err)
+
+	s.AssertExpectations(t)
+}
+
 func TestEmpire_Set(t *testing.T) {
 	e := empiretest.NewEmpire(t)
 	s := new(mockScheduler)
@@ -359,3 +502,9 @@ func (m *mockScheduler) Submit(_ context.Context, app *scheduler.App) error {
 	args := m.Called(app)
 	return args.Error(0)
 }
+
+func (m *mockScheduler) Run(_ context.Context, app *scheduler.App, process *scheduler.Process, in io.Reader, out io.Writer) error {
+	app.Processes = nil // This is bogus and doesn't actually matter for Runs.
+	args := m.Called(app, process, in, out)
+	return args.Error(0)
+}