130 lines
2.9 KiB
Go
130 lines
2.9 KiB
Go
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
package cgroups
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
|
|
"golang.org/x/sys/unix"
|
|
)
|
|
|
|
type Cgroup2Manager struct {
|
|
cgroupFDs map[ProcessType]int
|
|
}
|
|
|
|
var _ Manager = (*Cgroup2Manager)(nil)
|
|
|
|
type cgroup2Config struct {
|
|
rootPath string
|
|
processTypes map[ProcessType]Cgroup2Config
|
|
}
|
|
|
|
type Cgroup2ManagerOption func(*cgroup2Config)
|
|
|
|
func WithCgroup2RootSysFSPath(path string) Cgroup2ManagerOption {
|
|
return func(config *cgroup2Config) {
|
|
config.rootPath = path
|
|
}
|
|
}
|
|
|
|
func WithCgroup2ProcessType(processType ProcessType, path string, properties map[string]string) Cgroup2ManagerOption {
|
|
return func(config *cgroup2Config) {
|
|
if config.processTypes == nil {
|
|
config.processTypes = make(map[ProcessType]Cgroup2Config)
|
|
}
|
|
config.processTypes[processType] = Cgroup2Config{Path: path, Properties: properties}
|
|
}
|
|
}
|
|
|
|
type Cgroup2Config struct {
|
|
Path string
|
|
Properties map[string]string
|
|
}
|
|
|
|
func NewCgroup2Manager(opts ...Cgroup2ManagerOption) (*Cgroup2Manager, error) {
|
|
config := cgroup2Config{
|
|
rootPath: "/sys/fs/cgroup",
|
|
}
|
|
|
|
for _, opt := range opts {
|
|
opt(&config)
|
|
}
|
|
|
|
cgroupFDs, err := createCgroups(config)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create cgroups: %w", err)
|
|
}
|
|
|
|
return &Cgroup2Manager{cgroupFDs: cgroupFDs}, nil
|
|
}
|
|
|
|
func createCgroups(configs cgroup2Config) (map[ProcessType]int, error) {
|
|
var (
|
|
results = make(map[ProcessType]int)
|
|
errs []error
|
|
)
|
|
|
|
for procType, config := range configs.processTypes {
|
|
fullPath := filepath.Join(configs.rootPath, config.Path)
|
|
fd, err := createCgroup(fullPath, config.Properties)
|
|
if err != nil {
|
|
errs = append(errs, fmt.Errorf("failed to create %s cgroup: %w", procType, err))
|
|
|
|
continue
|
|
}
|
|
results[procType] = fd
|
|
}
|
|
|
|
if len(errs) > 0 {
|
|
for procType, fd := range results {
|
|
err := unix.Close(fd)
|
|
if err != nil {
|
|
errs = append(errs, fmt.Errorf("failed to close cgroup fd for %s: %w", procType, err))
|
|
}
|
|
}
|
|
|
|
return nil, errors.Join(errs...)
|
|
}
|
|
|
|
return results, nil
|
|
}
|
|
|
|
func createCgroup(fullPath string, properties map[string]string) (int, error) {
|
|
if err := os.MkdirAll(fullPath, 0o755); err != nil {
|
|
return -1, fmt.Errorf("failed to create cgroup root: %w", err)
|
|
}
|
|
|
|
var errs []error
|
|
for name, value := range properties {
|
|
if err := os.WriteFile(filepath.Join(fullPath, name), []byte(value), 0o644); err != nil {
|
|
errs = append(errs, fmt.Errorf("failed to write cgroup property: %w", err))
|
|
}
|
|
}
|
|
if len(errs) > 0 {
|
|
return -1, errors.Join(errs...)
|
|
}
|
|
|
|
return unix.Open(fullPath, unix.O_RDONLY, 0)
|
|
}
|
|
|
|
func (c Cgroup2Manager) GetFileDescriptor(procType ProcessType) (int, bool) {
|
|
fd, ok := c.cgroupFDs[procType]
|
|
|
|
return fd, ok
|
|
}
|
|
|
|
func (c Cgroup2Manager) Close() error {
|
|
var errs []error
|
|
for procType, fd := range c.cgroupFDs {
|
|
if err := unix.Close(fd); err != nil {
|
|
errs = append(errs, fmt.Errorf("failed to close cgroup fd for %s: %w", procType, err))
|
|
}
|
|
delete(c.cgroupFDs, procType)
|
|
}
|
|
|
|
return errors.Join(errs...)
|
|
}
|