cgroupv2: improve unit test

worker/cgroup_test.go

@@ -8,7 +8,10 @@ import (
 	"strings"
 	"testing"
 	"time"
+	"errors"
+	"syscall"
 	cgv1 "github.com/containerd/cgroups"
+	cgv2 "github.com/containerd/cgroups/v2"
 	units "github.com/docker/go-units"
 	"github.com/moby/moby/pkg/reexec"
 
@@ -20,7 +23,22 @@ func init() {
 }
 
 func TestCgroup(t *testing.T) {
-	Convey("Cgroup Should Work", t, func(ctx C) {
+	var cgcf *cgroupConfig
+	Convey("init cgroup", t, func(ctx C){
+		_, useCurrentCgroup := os.LookupEnv("USECURCGROUP")
+		cgcf = &cgroupConfig{BasePath: "/sys/fs/cgroup", Group: "tunasync", Subsystem: "cpu"}
+		if useCurrentCgroup {
+			cgcf.Group = ""
+		}
+		err := initCgroup(cgcf)
+		So(err, ShouldBeNil)
+		if cgcf.isUnified {
+			So(cgcf.cgMgrV2, ShouldNotBeNil)
+		} else {
+			So(cgcf.cgMgrV1, ShouldNotBeNil)
+		}
+
+	Convey("Cgroup Should Work", func(ctx C) {
 		tmpDir, err := ioutil.TempDir("", "tunasync")
 		defer os.RemoveAll(tmpDir)
 		So(err, ShouldBeNil)
@@ -79,15 +97,7 @@ sleep 30
 		provider, err := newCmdProvider(c)
 		So(err, ShouldBeNil)
 
-		cgcf := cgroupConfig{BasePath: "/sys/fs/cgroup", Group: "tunasync", Subsystem: "cpu"}
-		err = initCgroup(&cgcf)
-		So(err, ShouldBeNil)
-		if cgcf.isUnified {
-			So(cgcf.cgMgrV2, ShouldNotBeNil)
-		} else {
-			So(cgcf.cgMgrV1, ShouldNotBeNil)
-		}
-		cg := newCgroupHook(provider, cgcf, 0)
+		cg := newCgroupHook(provider, *cgcf, 0)
 		provider.AddHook(cg)
 
 		err = cg.preExec()
@@ -122,7 +132,7 @@ sleep 30
 
 	})
 
-	Convey("Rsync Memory Should Be Limited", t, func() {
+	Convey("Rsync Memory Should Be Limited", func() {
 		tmpDir, err := ioutil.TempDir("", "tunasync")
 		defer os.RemoveAll(tmpDir)
 		So(err, ShouldBeNil)
@@ -143,27 +153,31 @@ sleep 30
 		provider, err := newRsyncProvider(c)
 		So(err, ShouldBeNil)
 
-		cgcf := cgroupConfig{BasePath: "/sys/fs/cgroup", Group: "tunasync", Subsystem: "cpu"}
-		err = initCgroup(&cgcf)
-		So(err, ShouldBeNil)
-		if cgcf.isUnified {
-			So(cgcf.cgMgrV2, ShouldNotBeNil)
-		} else {
-			So(cgcf.cgMgrV1, ShouldNotBeNil)
-		}
-		cg := newCgroupHook(provider, cgcf, 512 * units.MiB)
+		cg := newCgroupHook(provider, *cgcf, 512 * units.MiB)
 		provider.AddHook(cg)
 
 		err = cg.preExec()
 		So(err, ShouldBeNil)
 		if cgcf.isUnified {
-			memoLimit, err := ioutil.ReadFile(filepath.Join(cgcf.BasePath, cgcf.Group, provider.Name(), "memory.max"))
+			cgpath := filepath.Join(cgcf.BasePath, cgcf.Group, provider.Name())
+			if useCurrentCgroup {
+				group, err := cgv2.NestedGroupPath(filepath.Join("..", provider.Name()))
+				So(err, ShouldBeNil)
+				cgpath = filepath.Join(cgcf.BasePath, group)
+			}
+			memoLimit, err := ioutil.ReadFile(filepath.Join(cgpath, "memory.max"))
 			So(err, ShouldBeNil)
 			So(strings.Trim(string(memoLimit), "\n"), ShouldEqual, strconv.Itoa(512*1024*1024))
 		} else {
 			for _, subsys := range(cg.cgMgrV1.Subsystems()) {
 				if subsys.Name() == cgv1.Memory {
-					memoLimit, err := ioutil.ReadFile(filepath.Join(cgcf.BasePath, "memory", cgcf.Group, provider.Name(), "memory.limit_in_bytes"))
+					cgpath := filepath.Join(cgcf.Group, provider.Name())
+					if useCurrentCgroup {
+						p, err := cgv1.NestedPath(filepath.Join("..", provider.Name()))(cgv1.Memory)
+						So(err, ShouldBeNil)
+						cgpath = p
+					}
+					memoLimit, err := ioutil.ReadFile(filepath.Join(cgcf.BasePath, "memory", cgpath, "memory.limit_in_bytes"))
 					So(err, ShouldBeNil)
 					So(strings.Trim(string(memoLimit), "\n"), ShouldEqual, strconv.Itoa(512*1024*1024))
 				}
@@ -172,4 +186,76 @@ sleep 30
 		cg.postExec()
 		So(cg.cgMgrV1, ShouldBeNil)
 	})
+	Reset(func() {
+		if cgcf.isUnified {
+			if cgcf.Group == "" {
+				wkrg, err := cgv2.NestedGroupPath("");
+				So(err, ShouldBeNil)
+				wkrMgr, err := cgv2.LoadManager("/sys/fs/cgroup", wkrg);
+				allCtrls, err := wkrMgr.Controllers()
+				So(err, ShouldBeNil)
+				err = wkrMgr.ToggleControllers(allCtrls, cgv2.Disable)
+				So(err, ShouldBeNil)
+				origMgr := cgcf.cgMgrV2
+				for {
+					logger.Debugf("Restoring pids")
+					procs, err := wkrMgr.Procs(false)
+					So(err, ShouldBeNil)
+					if len(procs) == 0 {
+						break
+					}
+					for _, p := range(procs) {
+						if err := origMgr.AddProc(p); err != nil{
+							if errors.Is(err, syscall.ESRCH) {
+								logger.Debugf("Write pid %d to sub group failed: process vanished, ignoring")
+							} else {
+								So(err, ShouldBeNil)
+							}
+						}
+					}
+				}
+				err = wkrMgr.Delete()
+				So(err, ShouldBeNil)
+			}
+		} else {
+			if cgcf.Group == "" {
+				pather := (func(p cgv1.Path) (cgv1.Path){
+					return func(subsys cgv1.Name) (string, error){
+						path, err := p(subsys);
+						if err != nil {
+							return "", err
+						}
+						if path == "/" {
+							return "", cgv1.ErrControllerNotActive
+						}
+						return path, err
+					}
+				})(cgv1.NestedPath(""))
+				wkrMgr, err := cgv1.Load(cgv1.V1, pather)
+				So(err, ShouldBeNil)
+				origMgr := cgcf.cgMgrV1
+				for _, subsys := range(wkrMgr.Subsystems()){
+					for {
+						procs, err := wkrMgr.Processes(subsys.Name(), false)
+						So(err, ShouldBeNil)
+						if len(procs) == 0 {
+							break
+						}
+						for _, proc := range(procs) {
+							if err := origMgr.Add(proc); err != nil {
+								if errors.Is(err, syscall.ESRCH) {
+									logger.Debugf("Write pid %d to sub group failed: process vanished, ignoring")
+								} else {
+									So(err, ShouldBeNil)
+								}
+							}
+						}
+					}
+				}
+				err = wkrMgr.Delete()
+				So(err, ShouldBeNil)
+			}
+		}
+	})
+	})
 }