From ed6b8fb0aa1ffcb5c14408f6f5bb3cfadb752277 Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Mon, 28 Aug 2017 11:10:15 -0400
Subject: [PATCH] Add KillOpts for killing all processes

Fixes #1431

This adds KillOpts so that a client can specify when they want to kill a
single process or all the processes inside a container.

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 cmd/ctr/run.go          |  2 +-
 container_linux_test.go | 65 +++++++++++++++++++++++++++++++++++++++++
 process.go              | 11 +++++--
 task.go                 |  9 +++++-
 task_opts.go            | 14 +++++++++
 5 files changed, 97 insertions(+), 4 deletions(-)

diff --git a/cmd/ctr/run.go b/cmd/ctr/run.go
index 880000876..5194dffa4 100644
--- a/cmd/ctr/run.go
+++ b/cmd/ctr/run.go
@@ -21,7 +21,7 @@ type resizer interface {
 }
 
 type killer interface {
-	Kill(gocontext.Context, syscall.Signal) error
+	Kill(gocontext.Context, syscall.Signal, ...containerd.KillOpts) error
 }
 
 func withEnv(context *cli.Context) containerd.SpecOpts {
diff --git a/container_linux_test.go b/container_linux_test.go
index c77063256..06a01e72e 100644
--- a/container_linux_test.go
+++ b/container_linux_test.go
@@ -647,3 +647,68 @@ func TestContainerUserID(t *testing.T) {
 		t.Errorf("expected uid:gid to be 3:4, but received %q", output)
 	}
 }
+
+func TestContainerKillAll(t *testing.T) {
+	t.Parallel()
+
+	client, err := newClient(t, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	var (
+		image       Image
+		ctx, cancel = testContext()
+		id          = t.Name()
+	)
+	defer cancel()
+
+	image, err = client.GetImage(ctx, testImage)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	container, err := client.NewContainer(ctx, id,
+		withNewSnapshot(id, image),
+		WithNewSpec(withImageConfig(image),
+			withProcessArgs("sh", "-c", "top"),
+			WithHostNamespace(specs.PIDNamespace),
+		),
+	)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	defer container.Delete(ctx, WithSnapshotCleanup)
+
+	stdout := bytes.NewBuffer(nil)
+	task, err := container.NewTask(ctx, NewIO(bytes.NewBuffer(nil), stdout, bytes.NewBuffer(nil)))
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	defer task.Delete(ctx)
+
+	statusC, err := task.Wait(ctx)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	if err := task.Start(ctx); err != nil {
+		t.Error(err)
+		return
+	}
+
+	if err := task.Kill(ctx, syscall.SIGKILL, WithKillAll); err != nil {
+		t.Error(err)
+	}
+
+	<-statusC
+	if _, err := task.Delete(ctx); err != nil {
+		t.Error(err)
+		return
+	}
+}
diff --git a/process.go b/process.go
index 6cea4c47a..536d8a85e 100644
--- a/process.go
+++ b/process.go
@@ -23,7 +23,7 @@ type Process interface {
 	// Delete removes the process and any resources allocated returning the exit status
 	Delete(context.Context, ...ProcessDeleteOpts) (*ExitStatus, error)
 	// Kill sends the provided signal to the process
-	Kill(context.Context, syscall.Signal) error
+	Kill(context.Context, syscall.Signal, ...KillOpts) error
 	// Wait asynchronously waits for the process to exit, and sends the exit code to the returned channel
 	Wait(context.Context) (<-chan ExitStatus, error)
 	// CloseIO allows various pipes to be closed on the process
@@ -104,11 +104,18 @@ func (p *process) Start(ctx context.Context) error {
 	return nil
 }
 
-func (p *process) Kill(ctx context.Context, s syscall.Signal) error {
+func (p *process) Kill(ctx context.Context, s syscall.Signal, opts ...KillOpts) error {
+	var i KillInfo
+	for _, o := range opts {
+		if err := o(ctx, p, &i); err != nil {
+			return err
+		}
+	}
 	_, err := p.task.client.TaskService().Kill(ctx, &tasks.KillRequest{
 		Signal:      uint32(s),
 		ContainerID: p.task.id,
 		ExecID:      p.id,
+		All:         i.All,
 	})
 	return errdefs.FromGRPC(err)
 }
diff --git a/task.go b/task.go
index 787db09fc..93a57e545 100644
--- a/task.go
+++ b/task.go
@@ -163,10 +163,17 @@ func (t *task) Start(ctx context.Context) error {
 	return errdefs.FromGRPC(err)
 }
 
-func (t *task) Kill(ctx context.Context, s syscall.Signal) error {
+func (t *task) Kill(ctx context.Context, s syscall.Signal, opts ...KillOpts) error {
+	var i KillInfo
+	for _, o := range opts {
+		if err := o(ctx, t, &i); err != nil {
+			return err
+		}
+	}
 	_, err := t.client.TaskService().Kill(ctx, &tasks.KillRequest{
 		Signal:      uint32(s),
 		ContainerID: t.id,
+		All:         i.All,
 	})
 	if err != nil {
 		return errdefs.FromGRPC(err)
diff --git a/task_opts.go b/task_opts.go
index 3dcab170f..77f67458e 100644
--- a/task_opts.go
+++ b/task_opts.go
@@ -51,3 +51,17 @@ func WithProcessKill(ctx context.Context, p Process) error {
 	<-s
 	return nil
 }
+
+type KillInfo struct {
+	// All kills all processes inside the task
+	// only valid on tasks, ignored on processes
+	All bool
+}
+
+type KillOpts func(context.Context, Process, *KillInfo) error
+
+// WithKillAll kills all processes for a task
+func WithKillAll(ctx context.Context, p Process, i *KillInfo) error {
+	i.All = true
+	return nil
+}