switch to implementation based on cgroup-toolscgroups-poc

2 files changed, 114 insertions, 44 deletions

diff --git a/internal/git/cgroups.go b/internal/git/cgroups.go
new file mode 100644
index 000000000..e362042a2
--- /dev/null
+++ b/internal/git/cgroups.go
@@ -0,0 +1,93 @@
+package git
+
+import (
+	"bufio"
+	"encoding/binary"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"os"
+	"os/exec"
+	"strconv"
+	"strings"
+	"sync"
+)
+
+var (
+	cgroupOnce sync.Once
+	cgroupSet  map[string]string
+	cgroupErr  error
+)
+
+// TODO: instead of discovery here, we could provide it as a
+// config flag GITALY_CGROUP_PREFIX, which users can then set
+// to something like "/system.slice/gitlab-runsvdir.service".
+// that wouold allos us to drop all of this code. but would also
+// requires to pre-create cgroups ahead of time
+func getCgroupSet() (map[string]string, error) {
+	cgroupOnce.Do(func() {
+		cgroupSet, cgroupErr = parseCgroupFile("/proc/self/cgroup")
+		if cgroupErr != nil {
+			return
+		}
+
+		cgroupCount := uint64(4096)
+		if os.Getenv("GITALY_CGROUP_COUNT") != "" {
+			cgroupCount, cgroupErr = strconv.ParseUint(os.Getenv("GITALY_CGROUP_COUNT"), 64, 10)
+			if cgroupErr != nil {
+				return
+			}
+		}
+
+		for i := uint64(0); i < cgroupCount; i++ {
+			src := make([]byte, 8)
+			binary.BigEndian.PutUint64(src, i)
+
+			dst := make([]byte, hex.EncodedLen(len(src)))
+			hex.Encode(dst, src)
+			dst = dst[len(dst)-3:]
+
+			// TODO: batch these up
+			cmd := exec.Command("sudo", "cgcreate", "-g", "cpu:"+cgroupSet["cpu"]+"/"+string(dst), "-g", "memory:"+cgroupSet["memory"]+"/"+string(dst))
+			cgroupErr := cmd.Run()
+			if cgroupErr != nil {
+				return
+			}
+		}
+	})
+	return cgroupSet, cgroupErr
+}
+
+func parseCgroupFile(path string) (map[string]string, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	return parseCgroupFromReader(f)
+}
+
+func parseCgroupFromReader(r io.Reader) (map[string]string, error) {
+	var (
+		cgroups = make(map[string]string)
+		s       = bufio.NewScanner(r)
+	)
+	for s.Scan() {
+		var (
+			text  = s.Text()
+			parts = strings.SplitN(text, ":", 3)
+		)
+		if len(parts) < 3 {
+			return nil, fmt.Errorf("invalid cgroup entry: %q", text)
+		}
+		for _, subs := range strings.Split(parts[1], ",") {
+			if subs != "" {
+				cgroups[subs] = parts[2]
+			}
+		}
+	}
+	if err := s.Err(); err != nil {
+		return nil, err
+	}
+	return cgroups, nil
+}
diff --git a/internal/git/command.go b/internal/git/command.go
index 19922687b..561a345f5 100644
--- a/internal/git/command.go
+++ b/internal/git/command.go
@@ -3,15 +3,14 @@ package git
 import (
 	"context"
 	"crypto/sha1"
+	"encoding/hex"
+	"os"
 	"os/exec"
 
 	"gitlab.com/gitlab-org/gitaly/internal/command"
 	"gitlab.com/gitlab-org/gitaly/internal/git/alternates"
 	"gitlab.com/gitlab-org/gitaly/internal/git/repository"
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
-
-	"github.com/containerd/cgroups"
-	specs "github.com/opencontainers/runtime-spec/specs-go"
 )
 
 // unsafeCmdWithEnv creates a git.unsafeCmd with the given args, environment, and Repository
@@ -51,62 +50,40 @@ func argsAndEnv(repo repository.GitRepo, args ...string) ([]string, []string, er
 // unsafeBareCmd creates a git.Command with the given args, stdin/stdout/stderr, and env
 func unsafeBareCmd(ctx context.Context, repo repository.GitRepo, stream CmdStream, env []string, args ...string) (*command.Command, error) {
 	if repo == nil {
+		env = append(env, command.GitEnv...)
 		return command.New(ctx, exec.Command(command.GitPath(), args...), stream.In, stream.Out, stream.Err, env...)
 	}
 
-	// TODO: extract this out into a platform-specific package
-
-	// 16^3 = 4k
-	repoPath, err := helper.GetRepoPath(repo)
-	if err != nil {
-		return nil, err
-	}
-	hash := sha1.Sum([]byte(repoPath))
-	groupName := string(hash[0:3])
-
-	subCgroup, err := cgroups.Load(cgroups.V1, cgroups.NestedPath(groupName))
-	if err != nil && err != cgroups.ErrCgroupDeleted {
-		return nil, err
-	}
+	args = append([]string{command.GitPath()}, args...)
 
-	if err == cgroups.ErrCgroupDeleted {
-		// cgroup does not yet exist, let's create it
-		control, err := cgroups.Load(cgroups.V1, cgroups.NestedPath(""))
+	if os.Getenv("GITALY_CGROUP_ENABLED") == "true" {
+		// 16^3 = 4k
+		repoPath, err := helper.GetRepoPath(repo)
 		if err != nil {
 			return nil, err
 		}
+		hash := sha1.Sum([]byte(repoPath))
 
-		// TODO: adjust limits dynamically, possibly based on percentage
-		// 100/1024 and 1GiB for now
-		cpuShares := uint64(100)
-		memoryLimit := int64(1024 * 1024 * 1024)
-		subCgroup, err = control.New(groupName, &specs.LinuxResources{
-			CPU: &specs.LinuxCPU{
-				Shares: &cpuShares,
-			},
-			Memory: &specs.LinuxMemory{
-				Limit: &memoryLimit,
-			},
-		})
+		dst := make([]byte, hex.EncodedLen(len(hash)))
+		hex.Encode(dst, hash[:])
+		groupName := string(dst[0:3])
+
+		cgroupSet, err = getCgroupSet()
 		if err != nil {
 			return nil, err
 		}
-	}
 
-	env = append(env, command.GitEnv...)
-
-	cmd, err := command.New(ctx, exec.Command(command.GitPath(), args...), stream.In, stream.Out, stream.Err, env...)
-	if err != nil {
-		return nil, err
-	}
+		cgexecArgs := []string{"sudo", "cgexec"}
+		for _, type_ := range []string{"cpu", "memory"} {
+			cgroupName := cgroupSet[type_] + "/" + groupName
+			cgexecArgs = append(cgexecArgs, "-g", type_+":"+cgroupName)
+		}
 
-	if err := subCgroup.Add(cgroups.Process{Pid: cmd.Pid()}); err != nil {
-		return nil, err
+		args = append(cgexecArgs, args...)
 	}
 
-	// TODO: check if we need to add children
-
-	return cmd, err
+	env = append(env, command.GitEnv...)
+	return command.New(ctx, exec.Command(args[0], args[1:]...), stream.In, stream.Out, stream.Err, env...)
 }
 
 // unsafeCmdWithoutRepo works like Command but without a git repository