From 4304da2b5800d5d4907956b359d4972c5a8cf6ef Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcgstyle.net>
Date: Mon, 15 May 2017 17:38:12 -0700
Subject: [PATCH] Delete rootfs when container is deleted

Signed-off-by: Derek McGowan <derek@mcgstyle.net>
---
 cmd/ctr/delete.go | 17 +++++++++++++++--
 cmd/ctr/run.go    | 11 +++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/cmd/ctr/delete.go b/cmd/ctr/delete.go
index 99988e6e8..570325f2e 100644
--- a/cmd/ctr/delete.go
+++ b/cmd/ctr/delete.go
@@ -17,13 +17,26 @@ var deleteCommand = cli.Command{
 		if err != nil {
 			return err
 		}
+		snapshotter, err := getSnapshotter(context)
+		if err != nil {
+			return err
+		}
 		id := context.Args().First()
 		if id == "" {
 			return errors.New("container id must be provided")
 		}
-		_, err = containers.Delete(gocontext.Background(), &execution.DeleteRequest{
+		ctx := gocontext.TODO()
+		_, err = containers.Delete(ctx, &execution.DeleteRequest{
 			ID: id,
 		})
-		return err
+		if err != nil {
+			return errors.Wrap(err, "failed to delete container")
+		}
+
+		if err := snapshotter.Remove(ctx, id); err != nil {
+			return errors.Wrapf(err, "failed to remove snapshot %q", id)
+		}
+
+		return nil
 	},
 }
diff --git a/cmd/ctr/run.go b/cmd/ctr/run.go
index 5b8908a23..53bf6054a 100644
--- a/cmd/ctr/run.go
+++ b/cmd/ctr/run.go
@@ -127,6 +127,14 @@ var runCommand = cli.Command{
 				if err != nil {
 					return err
 				}
+			} else {
+				defer func() {
+					if id != "" {
+						if err := snapshotter.Remove(ctx, id); err != nil {
+							logrus.WithError(err).Errorf("failed to remove snapshot %q", id)
+						}
+					}
+				}()
 			}
 
 			ic, err := image.Config(ctx, content)
@@ -206,6 +214,9 @@ var runCommand = cli.Command{
 			}); err != nil {
 				return err
 			}
+		} else {
+			// Don't remove snapshot
+			id = ""
 		}
 		if status != 0 {
 			return cli.NewExitError("", int(status))