Merge pull request #50 from jingxiaolu/master

estesp · web-flow · commit 5e610833b720 · 2018-05-15T13:50:38.000-04:00
Add interface AddTask to control groups.
diff --git a/cgroup.go b/cgroup.go
@@ -137,6 +137,36 @@ func (c *cgroup) add(process Process) error {
 	return nil
 }
 
+// AddTask moves the provided tasks (threads) into the new cgroup
+func (c *cgroup) AddTask(process Process) error {
+	if process.Pid <= 0 {
+		return ErrInvalidPid
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.err != nil {
+		return c.err
+	}
+	return c.addTask(process)
+}
+
+func (c *cgroup) addTask(process Process) error {
+	for _, s := range pathers(c.subsystems) {
+		p, err := c.path(s.Name())
+		if err != nil {
+			return err
+		}
+		if err := ioutil.WriteFile(
+			filepath.Join(s.Path(p), cgroupTasks),
+			[]byte(strconv.Itoa(process.Pid)),
+			defaultFilePerm,
+		); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // Delete will remove the control group from each of the subsystems registered
 func (c *cgroup) Delete() error {
 	c.mu.Lock()
diff --git a/cgroup_test.go b/cgroup_test.go
@@ -101,6 +101,29 @@ func TestAdd(t *testing.T) {
 	}
 }
 
+func TestAddTask(t *testing.T) {
+	mock, err := newMock()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer mock.delete()
+	control, err := New(mock.hierarchy, StaticPath("test"), &specs.LinuxResources{})
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if err := control.AddTask(Process{Pid: 1234}); err != nil {
+		t.Error(err)
+		return
+	}
+	for _, s := range Subsystems() {
+		if err := checkTaskid(mock, filepath.Join(string(s), "test"), 1234); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+}
+
 func TestListPids(t *testing.T) {
 	mock, err := newMock()
 	if err != nil {
@@ -159,6 +182,21 @@ func checkPid(mock *mockCgroup, path string, expected int) error {
 	return nil
 }
 
+func checkTaskid(mock *mockCgroup, path string, expected int) error {
+	data, err := readValue(mock, filepath.Join(path, cgroupTasks))
+	if err != nil {
+		return err
+	}
+	v, err := strconv.Atoi(string(data))
+	if err != nil {
+		return err
+	}
+	if v != expected {
+		return fmt.Errorf("expectd task id %d but received %d", expected, v)
+	}
+	return nil
+}
+
 func TestLoad(t *testing.T) {
 	mock, err := newMock()
 	if err != nil {
@@ -222,6 +260,16 @@ func TestCreateSubCgroup(t *testing.T) {
 			return
 		}
 	}
+	if err := sub.AddTask(Process{Pid: 5678}); err != nil {
+		t.Error(err)
+		return
+	}
+	for _, s := range Subsystems() {
+		if err := checkTaskid(mock, filepath.Join(string(s), "test", "child"), 5678); err != nil {
+			t.Error(err)
+			return
+		}
+	}
 }
 
 func TestFreezeThaw(t *testing.T) {
diff --git a/control.go b/control.go
@@ -24,6 +24,7 @@ import (
 
 const (
 	cgroupProcs    = "cgroup.procs"
+	cgroupTasks    = "tasks"
 	defaultDirPerm = 0755
 )
 
@@ -48,8 +49,10 @@ type Process struct {
 type Cgroup interface {
 	// New creates a new cgroup under the calling cgroup
 	New(string, *specs.LinuxResources) (Cgroup, error)
-	// Add adds a process to the cgroup
+	// Add adds a process to the cgroup (cgroup.procs)
 	Add(Process) error
+	// AddTask adds a process to the cgroup (tasks)
+	AddTask(Process) error
 	// Delete removes the cgroup as a whole
 	Delete() error
 	// MoveTo moves all the processes under the calling cgroup to the provided one
