1 files changed, 52 insertions, 13 deletions
diff --git a/runsc/cgroup/cgroup.go b/runsc/cgroup/cgroup.go
index 2887f3d7f..65a0b6d7a 100644
--- a/runsc/cgroup/cgroup.go
+++ b/runsc/cgroup/cgroup.go
@@ -17,6 +17,7 @@
 package cgroup
 
 import (
+	"bufio"
 	"context"
 	"fmt"
 	"io/ioutil"
@@ -168,12 +169,12 @@ type Cgroup struct {
 }
 
 // New creates a new Cgroup instance if the spec includes a cgroup path.
-// Otherwise it returns nil and false.
-func New(spec *specs.Spec) (*Cgroup, bool) {
+// Returns nil otherwise.
+func New(spec *specs.Spec) *Cgroup {
 	if spec.Linux == nil || spec.Linux.CgroupsPath == "" {
-		return nil, false
+		return nil
 	}
-	return &Cgroup{Name: spec.Linux.CgroupsPath}, true
+	return &Cgroup{Name: spec.Linux.CgroupsPath}
 }
 
 // Install creates and configures cgroups according to 'res'. If cgroup path
@@ -241,19 +242,57 @@ func (c *Cgroup) Uninstall() error {
 	return nil
 }
 
-// Join adds the current process to the all controllers.
-func (c *Cgroup) Join() error {
-	return c.Add(0)
-}
+// Join adds the current process to the all controllers. Returns function that
+// restores cgroup to the original state.
+func (c *Cgroup) Join() (func(), error) {
+	// First save the current state so it can be restored.
+	undo := func() {}
+	f, err := os.Open("/proc/self/cgroup")
+	if err != nil {
+		return undo, err
+	}
+	defer f.Close()
+
+	var undoPaths []string
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		// Format: ID:controller1,controller2:path
+		// Example: 2:cpu,cpuacct:/user.slice
+		tokens := strings.Split(scanner.Text(), ":")
+		if len(tokens) != 3 {
+			return undo, fmt.Errorf("formatting cgroups file, line: %q", scanner.Text())
+		}
+		for _, ctrlr := range strings.Split(tokens[1], ",") {
+			// Skip controllers we don't handle.
+			if _, ok := controllers[ctrlr]; ok {
+				undoPaths = append(undoPaths, filepath.Join(cgroupRoot, ctrlr, tokens[2]))
+				break
+			}
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return undo, err
+	}
 
-// Add adds given process to all controllers.
-func (c *Cgroup) Add(pid int) error {
+	// Replace empty undo with the real thing before changes are made to cgroups.
+	undo = func() {
+		for _, path := range undoPaths {
+			log.Debugf("Restoring cgroup %q", path)
+			if err := setValue(path, "cgroup.procs", "0"); err != nil {
+				log.Warningf("Error restoring cgroup %q: %v", path, err)
+			}
+		}
+	}
+
+	// Now join the cgroups.
 	for key := range controllers {
-		if err := setValue(c.makePath(key), "cgroup.procs", strconv.Itoa(pid)); err != nil {
-			return err
+		path := c.makePath(key)
+		log.Debugf("Joining cgroup %q", path)
+		if err := setValue(path, "cgroup.procs", "0"); err != nil {
+			return undo, err
 		}
 	}
-	return nil
+	return undo, nil
 }
 
 // NumCPU returns the number of CPUs configured in 'cpuset/cpuset.cpus'.