From daa23ce6732e0d56ddd46c0975beca613e8dbdbe Mon Sep 17 00:00:00 2001
From: Jacob McCann <jmccann.git@gmail.com>
Date: Sat, 24 Jun 2017 07:16:16 -0500
Subject: [PATCH] Allow specifying limits for pipeline steps from server config
 (#2085)

---
 drone/server/server.go  | 38 ++++++++++++++++++++++++++++++++++++++
 model/resource_limit.go | 11 +++++++++++
 server/hook.go          |  1 +
 server/rpc.go           |  1 +
 4 files changed, 51 insertions(+)
 create mode 100644 model/resource_limit.go

diff --git a/drone/server/server.go b/drone/server/server.go
index 8d96921bf..8caef5988 100644
--- a/drone/server/server.go
+++ b/drone/server/server.go
@@ -132,6 +132,36 @@ var Command = cli.Command{
 			Usage:  "database driver configuration string",
 			Value:  "drone.sqlite",
 		},
+		//
+		// resource limit parameters
+		//
+		cli.Int64Flag{
+			Name:   "limit-mem-swap",
+			EnvVar: "DRONE_LIMIT_MEM_SWAP",
+		},
+		cli.Int64Flag{
+			Name:   "limit-mem",
+			EnvVar: "DRONE_LIMIT_MEM",
+		},
+		cli.Int64Flag{
+			Name:   "limit-shm-size",
+			EnvVar: "DRONE_LIMIT_SHM_SIZE",
+		},
+		cli.Int64Flag{
+			Name:   "limit-cpu-quota",
+			EnvVar: "DRONE_LIMIT_CPU_QUOTA",
+		},
+		cli.Int64Flag{
+			Name:   "limit-cpu-shares",
+			EnvVar: "DRONE_LIMIT_CPU_SHARES",
+		},
+		cli.StringFlag{
+			Name:   "limit-cpu-set",
+			EnvVar: "DRONE_LIMIT_CPU_SET",
+		},
+		//
+		// remote parameters
+		//
 		cli.BoolFlag{
 			EnvVar: "DRONE_GITHUB",
 			Name:   "github",
@@ -443,6 +473,14 @@ func setupEvilGlobals(c *cli.Context, v store.Store) {
 		droneserver.Config.Services.Senders = sender.NewRemote(endpoint)
 	}
 
+	// limits
+	droneserver.Config.Pipeline.Limits.MemSwapLimit = c.Int64("limit-mem-swap")
+	droneserver.Config.Pipeline.Limits.MemLimit = c.Int64("limit-mem")
+	droneserver.Config.Pipeline.Limits.ShmSize = c.Int64("limit-shm-size")
+	droneserver.Config.Pipeline.Limits.CPUQuota = c.Int64("limit-cpu-quota")
+	droneserver.Config.Pipeline.Limits.CPUShares = c.Int64("limit-cpu-shares")
+	droneserver.Config.Pipeline.Limits.CPUSet = c.String("limit-cpu-set")
+
 	// server configuration
 	droneserver.Config.Server.Cert = c.String("server-cert")
 	droneserver.Config.Server.Key = c.String("server-key")
diff --git a/model/resource_limit.go b/model/resource_limit.go
new file mode 100644
index 000000000..af2b68c7c
--- /dev/null
+++ b/model/resource_limit.go
@@ -0,0 +1,11 @@
+package model
+
+// ResourceLimit is the resource limit to set on pipeline steps
+type ResourceLimit struct {
+	MemSwapLimit int64
+	MemLimit     int64
+	ShmSize      int64
+	CPUQuota     int64
+	CPUShares    int64
+	CPUSet       string
+}
diff --git a/server/hook.go b/server/hook.go
index 668d9b496..e8b21efcb 100644
--- a/server/hook.go
+++ b/server/hook.go
@@ -484,6 +484,7 @@ func (b *builder) Build() ([]*buildItem, error) {
 			compiler.WithEnviron(environ),
 			compiler.WithEnviron(b.Envs),
 			compiler.WithEscalated(Config.Pipeline.Privileged...),
+			compiler.WithResourceLimit(Config.Pipeline.Limits.MemSwapLimit, Config.Pipeline.Limits.MemLimit, Config.Pipeline.Limits.ShmSize, Config.Pipeline.Limits.CPUQuota, Config.Pipeline.Limits.CPUShares, Config.Pipeline.Limits.CPUSet),
 			compiler.WithVolumes(Config.Pipeline.Volumes...),
 			compiler.WithNetworks(Config.Pipeline.Networks...),
 			compiler.WithLocal(false),
diff --git a/server/rpc.go b/server/rpc.go
index 7a8247863..099bcfccb 100644
--- a/server/rpc.go
+++ b/server/rpc.go
@@ -58,6 +58,7 @@ var Config = struct {
 		// Admins map[string]struct{}
 	}
 	Pipeline struct {
+		Limits     model.ResourceLimit
 		Volumes    []string
 		Networks   []string
 		Privileged []string