From 5b7969448ad752ae66504fe7f2fba47aba2c766a Mon Sep 17 00:00:00 2001
From: Shishir Mahajan <smahajan@roblox.com>
Date: Fri, 9 Apr 2021 10:52:42 -0700
Subject: [PATCH] Add support for image_pull_timeout.

---
 README.md                |  1 +
 containerd/containerd.go |  9 +++++++--
 containerd/driver.go     | 43 ++++++++++++++++++++++------------------
 3 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index bb2f794..ad378ec 100644
--- a/README.md
+++ b/README.md
@@ -93,6 +93,7 @@ More detailed instructions are in the [`example README.md`](https://github.com/R
 | Option | Type | Required | Description |
 | :---: | :---: | :---: | :--- |
 | **image** | string | yes | OCI image (docker is also OCI compatible) for your container. |
+| **image_pull_timeout** | string | no | A time duration that controls how long `containerd-driver` will wait before cancelling an in-progress pull of the OCI image as specified in `image`. Defaults to `"5m"`. |
 | **command** | string | no | Command to override command defined in the image. |
 | **args** | []string | no | Arguments to the command. |
 | **entrypoint** | []string | no | A string list overriding the image's entrypoint. |
diff --git a/containerd/containerd.go b/containerd/containerd.go
index 4350688..a55785b 100644
--- a/containerd/containerd.go
+++ b/containerd/containerd.go
@@ -61,8 +61,13 @@ func (d *Driver) getContainerdVersion() (containerd.Version, error) {
 	return d.client.Version(ctxWithTimeout)
 }
 
-func (d *Driver) pullImage(imageName string) (containerd.Image, error) {
-	ctxWithTimeout, cancel := context.WithTimeout(d.ctxContainerd, 90*time.Second)
+func (d *Driver) pullImage(imageName, imagePullTimeout string) (containerd.Image, error) {
+	pullTimeout, err := time.ParseDuration(imagePullTimeout)
+	if err != nil {
+		return nil, fmt.Errorf("Failed to parse image_pull_timeout: %v", err)
+	}
+
+	ctxWithTimeout, cancel := context.WithTimeout(d.ctxContainerd, pullTimeout)
 	defer cancel()
 
 	named, err := refdocker.ParseDockerRef(imageName)
diff --git a/containerd/driver.go b/containerd/driver.go
index 6258c52..b907537 100644
--- a/containerd/driver.go
+++ b/containerd/driver.go
@@ -104,6 +104,10 @@ var (
 			hclspec.NewAttr("host_dns", "bool", false),
 			hclspec.NewLiteral("true"),
 		),
+		"image_pull_timeout": hclspec.NewDefault(
+			hclspec.NewAttr("image_pull_timeout", "string", false),
+			hclspec.NewLiteral(`"5m"`),
+		),
 		"extra_hosts":     hclspec.NewAttr("extra_hosts", "list(string)", false),
 		"entrypoint":      hclspec.NewAttr("entrypoint", "list(string)", false),
 		"seccomp":         hclspec.NewAttr("seccomp", "bool", false),
@@ -153,24 +157,25 @@ type Mount struct {
 // TaskConfig contains configuration information for a task that runs with
 // this plugin
 type TaskConfig struct {
-	Image          string             `codec:"image"`
-	Command        string             `codec:"command"`
-	Args           []string           `codec:"args"`
-	CapAdd         []string           `codec:"cap_add"`
-	CapDrop        []string           `codec:"cap_drop"`
-	Cwd            string             `codec:"cwd"`
-	Devices        []string           `codec:"devices"`
-	Seccomp        bool               `codec:"seccomp"`
-	SeccompProfile string             `codec:"seccomp_profile"`
-	Sysctl         hclutils.MapStrStr `codec:"sysctl"`
-	Privileged     bool               `codec:"privileged"`
-	PidsLimit      int64              `codec:"pids_limit"`
-	HostDNS        bool               `codec:"host_dns"`
-	ExtraHosts     []string           `codec:"extra_hosts"`
-	Entrypoint     []string           `codec:"entrypoint"`
-	ReadOnlyRootfs bool               `codec:"readonly_rootfs"`
-	HostNetwork    bool               `codec:"host_network"`
-	Mounts         []Mount            `codec:"mounts"`
+	Image            string             `codec:"image"`
+	Command          string             `codec:"command"`
+	Args             []string           `codec:"args"`
+	CapAdd           []string           `codec:"cap_add"`
+	CapDrop          []string           `codec:"cap_drop"`
+	Cwd              string             `codec:"cwd"`
+	Devices          []string           `codec:"devices"`
+	Seccomp          bool               `codec:"seccomp"`
+	SeccompProfile   string             `codec:"seccomp_profile"`
+	Sysctl           hclutils.MapStrStr `codec:"sysctl"`
+	Privileged       bool               `codec:"privileged"`
+	PidsLimit        int64              `codec:"pids_limit"`
+	HostDNS          bool               `codec:"host_dns"`
+	ImagePullTimeout string             `codec:"image_pull_timeout"`
+	ExtraHosts       []string           `codec:"extra_hosts"`
+	Entrypoint       []string           `codec:"entrypoint"`
+	ReadOnlyRootfs   bool               `codec:"readonly_rootfs"`
+	HostNetwork      bool               `codec:"host_network"`
+	Mounts           []Mount            `codec:"mounts"`
 }
 
 // TaskState is the runtime state which is encoded in the handle returned to
@@ -404,7 +409,7 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *drive
 	containerConfig.ContainerName = containerName
 
 	var err error
-	containerConfig.Image, err = d.pullImage(driverConfig.Image)
+	containerConfig.Image, err = d.pullImage(driverConfig.Image, driverConfig.ImagePullTimeout)
 	if err != nil {
 		return nil, nil, fmt.Errorf("Error in pulling image %s: %v", driverConfig.Image, err)
 	}