From f0ccdc47dcda328026050833ea045a988b524e3f Mon Sep 17 00:00:00 2001
From: Miao Wang <shankerwangmiao@gmail.com>
Date: Mon, 12 Jul 2021 22:26:13 +0800
Subject: [PATCH] docker: support memory limit

---
 worker/docker.go      | 2 ++
 worker/docker_test.go | 2 ++
 worker/runner.go      | 4 ++++
 3 files changed, 8 insertions(+)

diff --git a/worker/docker.go b/worker/docker.go
index b321ced..ae8cf77 100644
--- a/worker/docker.go
+++ b/worker/docker.go
@@ -13,6 +13,7 @@ type dockerHook struct {
 	image   string
 	volumes []string
 	options []string
+	memoryLimit MemBytes
 }
 
 func newDockerHook(p mirrorProvider, gCfg dockerConfig, mCfg mirrorConfig) *dockerHook {
@@ -35,6 +36,7 @@ func newDockerHook(p mirrorProvider, gCfg dockerConfig, mCfg mirrorConfig) *dock
 		image:   mCfg.DockerImage,
 		volumes: volumes,
 		options: options,
+		memoryLimit: mCfg.MemoryLimit,
 	}
 }
 
diff --git a/worker/docker_test.go b/worker/docker_test.go
index 2d1d849..0e5110b 100644
--- a/worker/docker_test.go
+++ b/worker/docker_test.go
@@ -8,6 +8,7 @@ import (
 	"path/filepath"
 	"testing"
 	"time"
+	units "github.com/docker/go-units"
 
 	"github.com/codeskyblue/go-sh"
 	. "github.com/smartystreets/goconvey/convey"
@@ -77,6 +78,7 @@ sleep 20
 			volumes: []string{
 				fmt.Sprintf("%s:%s", cmdScript, "/bin/cmd.sh"),
 			},
+			memoryLimit: 512 * units.MiB,
 		}
 		provider.AddHook(d)
 		So(provider.Docker(), ShouldNotBeNil)
diff --git a/worker/runner.go b/worker/runner.go
index 01addcd..4b9a14e 100644
--- a/worker/runner.go
+++ b/worker/runner.go
@@ -56,6 +56,10 @@ func newCmdJob(provider mirrorProvider, cmdAndArgs []string, workingDir string,
 			kv := fmt.Sprintf("%s=%s", k, v)
 			args = append(args, "-e", kv)
 		}
+		// set memlimit
+		if d.memoryLimit != 0 {
+		  args = append(args, "-m", fmt.Sprint(d.memoryLimit.Value()))
+		}
 		// apply options
 		args = append(args, d.options...)
 		// apply image and command