Restore to original cgroup after sandbox and gofer processes are created

The original code assumed that it was safe to join and not restore cgroup,
but Container.Run will not exit after calling start, making cgroup cleanup
fail because there were still processes inside the cgroup.

PiperOrigin-RevId: 228529199
Change-Id: I12a48d9adab4bbb02f20d71ec99598c336cbfe51

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'.