cgroups for archivejc-cgroups-archive

7 files changed, 45 insertions, 24 deletions

diff --git a/cmd/gitaly/main.go b/cmd/gitaly/main.go
index db7c73a06..b62a90cfe 100644
--- a/cmd/gitaly/main.go
+++ b/cmd/gitaly/main.go
@@ -240,6 +240,7 @@ func run(cfg config.Cfg) error {
 			CatfileCache:       catfileCache,
 			DiskCache:          diskCache,
 			PackObjectsCache:   streamcache.New(cfg.PackObjectsCache, glog.Default()),
+			CgroupsManager:     cgroups.NewManager(cfg.Cgroups),
 		})
 		b.RegisterStarter(starter.New(c, srv))
 	}
diff --git a/internal/cgroups/cgroups.go b/internal/cgroups/cgroups.go
index 733f89d74..6801a5e9e 100644
--- a/internal/cgroups/cgroups.go
+++ b/internal/cgroups/cgroups.go
@@ -2,6 +2,7 @@ package cgroups
 
 import (
 	"gitlab.com/gitlab-org/gitaly/v14/internal/command"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/git/repository"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/config/cgroups"
 )
 
@@ -12,7 +13,7 @@ type Manager interface {
 	// instance of the Manager.
 	Setup() error
 	// AddCommand adds a Command to a cgroup
-	AddCommand(*command.Command) error
+	AddCommand(repository.GitRepo, *command.Command) error
 	// Cleanup cleans up cgroups created in Setup.
 	// It is expected to be called once at Gitaly shutdown from any
 	// instance of the Manager.
diff --git a/internal/cgroups/noop.go b/internal/cgroups/noop.go
index caac82e2d..ccd87a476 100644
--- a/internal/cgroups/noop.go
+++ b/internal/cgroups/noop.go
@@ -2,6 +2,7 @@ package cgroups
 
 import (
 	"gitlab.com/gitlab-org/gitaly/v14/internal/command"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/git/repository"
 )
 
 // NoopManager is a cgroups manager that does nothing
@@ -13,7 +14,7 @@ func (cg *NoopManager) Setup() error {
 }
 
 //nolint: revive,stylecheck // This is unintentionally missing documentation.
-func (cg *NoopManager) AddCommand(cmd *command.Command) error {
+func (cg *NoopManager) AddCommand(repo repository.GitRepo, cmd *command.Command) error {
 	return nil
 }
 
diff --git a/internal/cgroups/v1_linux.go b/internal/cgroups/v1_linux.go
index 42a797aa1..5dbbef788 100644
--- a/internal/cgroups/v1_linux.go
+++ b/internal/cgroups/v1_linux.go
@@ -54,8 +54,8 @@ func (cg *CGroupV1Manager) Setup() error {
 }
 
 //nolint: revive,stylecheck // This is unintentionally missing documentation.
-func (cg *CGroupV1Manager) AddCommand(cmd *command.Command) error {
-	checksum := crc32.ChecksumIEEE([]byte(strings.Join(cmd.Args(), "")))
+func (cg *CGroupV1Manager) AddCommand(repo repsitory.GitRepo, cmd *command.Command) error {
+	checksum := crc32.ChecksumIEEE([]byte(repo.GetRelativePath()))
 	groupID := uint(checksum) % cg.cfg.Count
 	cgroupPath := cg.cgroupPath(int(groupID))
 
diff --git a/internal/gitaly/service/dependencies.go b/internal/gitaly/service/dependencies.go
index dcdec1783..f3b8cb063 100644
--- a/internal/gitaly/service/dependencies.go
+++ b/internal/gitaly/service/dependencies.go
@@ -4,6 +4,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/client"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/backchannel"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/cache"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/cgroups"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git/catfile"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/config"
@@ -31,6 +32,7 @@ type Dependencies struct {
 	CatfileCache        catfile.Cache
 	DiskCache           cache.Cache
 	PackObjectsCache    streamcache.Cache
+	CgroupsManager      cgroups.Manager
 }
 
 // GetCfg returns service configuration.
@@ -97,3 +99,7 @@ func (dc *Dependencies) GetDiskCache() cache.Cache {
 func (dc *Dependencies) GetPackObjectsCache() streamcache.Cache {
 	return dc.PackObjectsCache
 }
+
+func (dc *Dependencies) GetCgroupsManager() cgroups.Manager {
+	return dc.CgroupsManager
+}
diff --git a/internal/gitaly/service/repository/archive.go b/internal/gitaly/service/repository/archive.go
index 844e00367..2beaa1912 100644
--- a/internal/gitaly/service/repository/archive.go
+++ b/internal/gitaly/service/repository/archive.go
@@ -10,6 +10,8 @@ import (
 	"path/filepath"
 	"strings"
 
+	"gitlab.com/gitlab-org/gitaly/v14/internal/git/repository"
+
 	"gitlab.com/gitlab-org/gitaly/v14/internal/command"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/commit"
@@ -33,6 +35,7 @@ type archiveParams struct {
 	tlsCfg      []byte
 	binDir      string
 	loggingDir  string
+	repo        repository.GitRepo
 }
 
 func (s *server) GetArchive(in *gitalypb.GetArchiveRequest, stream gitalypb.RepositoryService_GetArchiveServer) error {
@@ -104,6 +107,7 @@ func (s *server) GetArchive(in *gitalypb.GetArchiveRequest, stream gitalypb.Repo
 		tlsCfg:      tlsCfg,
 		binDir:      s.binDir,
 		loggingDir:  s.loggingCfg.Dir,
+		repo:        in.GetRepository(),
 	})
 }
 
@@ -228,6 +232,10 @@ func (s *server) handleArchive(p archiveParams) error {
 		return err
 	}
 
+	if err := s.cgroupsManager.AddCommand(p.repo, archiveCommand); err != nil {
+		return err
+	}
+
 	if p.compressCmd != nil {
 		command, err := command.New(p.ctx, p.compressCmd, archiveCommand, p.writer, nil)
 		if err != nil {
diff --git a/internal/gitaly/service/repository/server.go b/internal/gitaly/service/repository/server.go
index d211325e5..dbd9fbbd9 100644
--- a/internal/gitaly/service/repository/server.go
+++ b/internal/gitaly/service/repository/server.go
@@ -2,6 +2,7 @@ package repository
 
 import (
 	"gitlab.com/gitlab-org/gitaly/v14/client"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/cgroups"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git/catfile"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git/localrepo"
@@ -16,16 +17,17 @@ import (
 
 type server struct {
 	gitalypb.UnimplementedRepositoryServiceServer
-	ruby          *rubyserver.Server
-	conns         *client.Pool
-	locator       storage.Locator
-	txManager     transaction.Manager
-	gitCmdFactory git.CommandFactory
-	cfg           config.Cfg
-	binDir        string
-	loggingCfg    config.Logging
-	catfileCache  catfile.Cache
-	git2go        git2go.Executor
+	ruby           *rubyserver.Server
+	conns          *client.Pool
+	locator        storage.Locator
+	txManager      transaction.Manager
+	gitCmdFactory  git.CommandFactory
+	cfg            config.Cfg
+	binDir         string
+	loggingCfg     config.Logging
+	catfileCache   catfile.Cache
+	git2go         git2go.Executor
+	cgroupsManager cgroups.Manager
 }
 
 // NewServer creates a new instance of a gRPC repo server
@@ -37,18 +39,20 @@ func NewServer(
 	gitCmdFactory git.CommandFactory,
 	catfileCache catfile.Cache,
 	connsPool *client.Pool,
+	cgroupsManager cgroups.Manager,
 ) gitalypb.RepositoryServiceServer {
 	return &server{
-		ruby:          rs,
-		locator:       locator,
-		txManager:     txManager,
-		gitCmdFactory: gitCmdFactory,
-		conns:         connsPool,
-		cfg:           cfg,
-		binDir:        cfg.BinDir,
-		loggingCfg:    cfg.Logging,
-		catfileCache:  catfileCache,
-		git2go:        git2go.NewExecutor(cfg, locator),
+		ruby:           rs,
+		locator:        locator,
+		txManager:      txManager,
+		gitCmdFactory:  gitCmdFactory,
+		conns:          connsPool,
+		cfg:            cfg,
+		binDir:         cfg.BinDir,
+		loggingCfg:     cfg.Logging,
+		catfileCache:   catfileCache,
+		git2go:         git2go.NewExecutor(cfg, locator),
+		cgroupsManager: cgroupsManager,
 	}
 }