From d3a8055e2d79f2d2a6112b73e07bd1c15387e776 Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Tue, 20 Feb 2018 15:10:50 -0500
Subject: [PATCH] Add --pid-file to `ctr`

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 cmd/ctr/commands/commands.go    | 20 ++++++++++++++++++++
 cmd/ctr/commands/run/run.go     |  9 +++++++++
 cmd/ctr/commands/tasks/start.go | 10 +++++++++-
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/cmd/ctr/commands/commands.go b/cmd/ctr/commands/commands.go
index 8922cf34f..1a0c3b6b3 100644
--- a/cmd/ctr/commands/commands.go
+++ b/cmd/ctr/commands/commands.go
@@ -20,6 +20,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"path/filepath"
 	"strings"
 
 	"github.com/containerd/containerd"
@@ -99,3 +100,22 @@ func PrintAsJSON(x interface{}) {
 	}
 	fmt.Println(string(b))
 }
+
+// WritePidFile writes the pid atomically to a file
+func WritePidFile(path string, pid int) error {
+	path, err := filepath.Abs(path)
+	if err != nil {
+		return err
+	}
+	tempPath := filepath.Join(filepath.Dir(path), fmt.Sprintf(".%s", filepath.Base(path)))
+	f, err := os.OpenFile(tempPath, os.O_RDWR|os.O_CREATE|os.O_EXCL|os.O_SYNC, 0666)
+	if err != nil {
+		return err
+	}
+	_, err = fmt.Fprintf(f, "%d", pid)
+	f.Close()
+	if err != nil {
+		return err
+	}
+	return os.Rename(tempPath, path)
+}
diff --git a/cmd/ctr/commands/run/run.go b/cmd/ctr/commands/run/run.go
index 0cc3b1f4a..f9a429dfd 100644
--- a/cmd/ctr/commands/run/run.go
+++ b/cmd/ctr/commands/run/run.go
@@ -85,6 +85,10 @@ var ContainerFlags = []cli.Flag{
 		Name:  "with-ns",
 		Usage: "specify existing Linux namespaces to join at container runtime (format '<nstype>:<path>')",
 	},
+	cli.StringFlag{
+		Name:  "pid-file",
+		Usage: "file path to write the task's pid",
+	},
 }
 
 func loadSpec(path string, s *specs.Spec) error {
@@ -211,6 +215,11 @@ var Command = cli.Command{
 				return err
 			}
 		}
+		if context.IsSet("pid-file") {
+			if err := commands.WritePidFile(context.String("pid-file"), int(task.Pid())); err != nil {
+				return err
+			}
+		}
 		var con console.Console
 		if tty {
 			con = console.Current()
diff --git a/cmd/ctr/commands/tasks/start.go b/cmd/ctr/commands/tasks/start.go
index 68777b56c..63694a2e1 100644
--- a/cmd/ctr/commands/tasks/start.go
+++ b/cmd/ctr/commands/tasks/start.go
@@ -38,6 +38,10 @@ var startCommand = cli.Command{
 			Name:  "fifo-dir",
 			Usage: "directory used for storing IO FIFOs",
 		},
+		cli.StringFlag{
+			Name:  "pid-file",
+			Usage: "file path to write the task's pid",
+		},
 	},
 	Action: func(context *cli.Context) error {
 		var (
@@ -72,7 +76,11 @@ var startCommand = cli.Command{
 			return err
 		}
 		defer task.Delete(ctx)
-
+		if context.IsSet("pid-file") {
+			if err := commands.WritePidFile(context.String("pid-file"), int(task.Pid())); err != nil {
+				return err
+			}
+		}
 		statusC, err := task.Wait(ctx)
 		if err != nil {
 			return err