From 20d3fae3dbc9aa07b1a18c7159f1a8df41347d75 Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Tue, 17 Sep 2019 11:14:28 -0400
Subject: [PATCH] Add Opt for modifying shm size

Closes #3654

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 oci/spec_opts.go           | 25 ++++++++++++++++++
 oci/spec_opts_test.go      | 54 ++++++++++++++++++++++++++++++++++++++
 oci/spec_opts_unix_test.go |  2 +-
 3 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/oci/spec_opts.go b/oci/spec_opts.go
index 33bbcf23c..ad6b52a9f 100644
--- a/oci/spec_opts.go
+++ b/oci/spec_opts.go
@@ -1222,3 +1222,28 @@ func WithEnvFile(path string) SpecOpts {
 		return WithEnv(vars)(nil, nil, nil, s)
 	}
 }
+
+// ErrNoShmMount is returned when there is no /dev/shm mount specified in the config
+// and an Opts was trying to set a configuration value on the mount.
+var ErrNoShmMount = errors.New("no /dev/shm mount specified")
+
+// WithDevShmSize sets the size of the /dev/shm mount for the container.
+//
+// The size value is specified in kb, kilobytes.
+func WithDevShmSize(kb int64) SpecOpts {
+	return func(ctx context.Context, _ Client, c *containers.Container, s *Spec) error {
+		for _, m := range s.Mounts {
+			if m.Source == "shm" && m.Type == "tmpfs" {
+				for i, o := range m.Options {
+					if strings.HasPrefix(o, "size=") {
+						m.Options[i] = fmt.Sprintf("size=%dk", kb)
+						return nil
+					}
+				}
+				m.Options = append(m.Options, fmt.Sprintf("size=%dk", kb))
+				return nil
+			}
+		}
+		return ErrNoShmMount
+	}
+}
diff --git a/oci/spec_opts_test.go b/oci/spec_opts_test.go
index 77cbe2261..9ac0683be 100644
--- a/oci/spec_opts_test.go
+++ b/oci/spec_opts_test.go
@@ -27,6 +27,7 @@ import (
 	"os"
 	"reflect"
 	"runtime"
+	"strings"
 	"testing"
 
 	"github.com/containerd/containerd/content"
@@ -506,3 +507,56 @@ func TestDropCaps(t *testing.T) {
 		}
 	}
 }
+
+func TestDevShmSize(t *testing.T) {
+	t.Parallel()
+	var (
+		s   Spec
+		c   = containers.Container{ID: t.Name()}
+		ctx = namespaces.WithNamespace(context.Background(), "test")
+	)
+
+	err := populateDefaultUnixSpec(ctx, &s, c.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected := "1024k"
+	if err := WithDevShmSize(1024)(nil, nil, nil, &s); err != nil {
+		t.Fatal(err)
+	}
+	m := getShmMount(&s)
+	if m == nil {
+		t.Fatal("no shm mount found")
+	}
+	o := getShmSize(m.Options)
+	if o == "" {
+		t.Fatal("shm size not specified")
+	}
+	parts := strings.Split(o, "=")
+	if len(parts) != 2 {
+		t.Fatal("invalid size format")
+	}
+	size := parts[1]
+	if size != expected {
+		t.Fatalf("size %s not equal %s", size, expected)
+	}
+}
+
+func getShmMount(s *Spec) *specs.Mount {
+	for _, m := range s.Mounts {
+		if m.Source == "shm" && m.Type == "tmpfs" {
+			return &m
+		}
+	}
+	return nil
+}
+
+func getShmSize(opts []string) string {
+	for _, o := range opts {
+		if strings.HasPrefix(o, "size=") {
+			return o
+		}
+	}
+	return ""
+}
diff --git a/oci/spec_opts_unix_test.go b/oci/spec_opts_unix_test.go
index 22373fb67..4dd9a48ab 100644
--- a/oci/spec_opts_unix_test.go
+++ b/oci/spec_opts_unix_test.go
@@ -32,7 +32,7 @@ func TestWithImageConfigNoEnv(t *testing.T) {
 	t.Parallel()
 	var (
 		s   Spec
-		c   = containers.Container{ID: "TestWithImageConfigNoEnv"}
+		c   = containers.Container{ID: t.Name()}
 		ctx = namespaces.WithNamespace(context.Background(), "test")
 	)