From dc4509a0685324a8e268409297c084e21c24ee20 Mon Sep 17 00:00:00 2001 From: Jason Lynch Date: Sun, 5 Apr 2026 09:32:57 -0400 Subject: [PATCH] feat: cpu and memory limits for systemd Adds CPU and memory limits for databases deployed with systemd using these service properties: - `CPUQuota`: Expressed as a percentage where 100% = 1 CPU core. - `MemoryLimit`: Expressed as an integer with optional units, such as `1G`. Defaults to bytes when units are unspecified. These are populated from existing properties on our instance spec. This commit includes a small refactor to the way we produce and test unit file options. These changes are aimed at making it easier to add new types of unit files when we implement supporting services for systemd. PLAT-417 --- .../PatroniUnitOptions/cpu_limit.service | 17 ++ .../fractional_cpu_limit.service | 17 ++ .../PatroniUnitOptions/memory_max.service | 17 ++ .../PatroniUnitOptions/minimal.service | 16 ++ .../orchestrator/systemd/main_test.go | 16 ++ .../orchestrator/systemd/orchestrator.go | 2 +- .../orchestrator/systemd/patroni_unit.go | 92 +++---- .../orchestrator/systemd/unit_options.go | 245 ++++++++++++++++++ .../orchestrator/systemd/unit_options_test.go | 95 +++++++ 9 files changed, 456 insertions(+), 61 deletions(-) create mode 100644 server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/cpu_limit.service create mode 100644 server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/fractional_cpu_limit.service create mode 100644 server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/memory_max.service create mode 100644 server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/minimal.service create mode 100644 server/internal/orchestrator/systemd/main_test.go create mode 100644 server/internal/orchestrator/systemd/unit_options.go create mode 100644 server/internal/orchestrator/systemd/unit_options_test.go diff --git a/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/cpu_limit.service b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/cpu_limit.service new file mode 100644 index 00000000..aa30b0b5 --- /dev/null +++ b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/cpu_limit.service @@ -0,0 +1,17 @@ +[Unit] +After=syslog.target network.target + +[Service] +Type=simple +User=postgres +ExecStart=/usr/local/bin/patroni /var/lib/pgsql/18/storefront-n1-689qacsi/configs/patroni.yaml +ExecReload=/bin/kill -s HUP $MAINPID +KillMode=process +TimeoutSec=30 +CPUQuota=1400% +Restart=on-failure +Environment="PATH=/usr/pgsql-18/bin:/root/.local/bin:/root/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/usr/local/go/bin" +Environment="PGSERVICEFILE=/var/lib/pgsql/18/storefront-n1-689qacsi/configs/pg_service.conf" + +[Install] +WantedBy=multi-user.target diff --git a/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/fractional_cpu_limit.service b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/fractional_cpu_limit.service new file mode 100644 index 00000000..8d3e2ee9 --- /dev/null +++ b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/fractional_cpu_limit.service @@ -0,0 +1,17 @@ +[Unit] +After=syslog.target network.target + +[Service] +Type=simple +User=postgres +ExecStart=/usr/local/bin/patroni /var/lib/pgsql/18/storefront-n1-689qacsi/configs/patroni.yaml +ExecReload=/bin/kill -s HUP $MAINPID +KillMode=process +TimeoutSec=30 +CPUQuota=50% +Restart=on-failure +Environment="PATH=/usr/pgsql-18/bin:/root/.local/bin:/root/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/usr/local/go/bin" +Environment="PGSERVICEFILE=/var/lib/pgsql/18/storefront-n1-689qacsi/configs/pg_service.conf" + +[Install] +WantedBy=multi-user.target diff --git a/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/memory_max.service b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/memory_max.service new file mode 100644 index 00000000..b9e78ed2 --- /dev/null +++ b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/memory_max.service @@ -0,0 +1,17 @@ +[Unit] +After=syslog.target network.target + +[Service] +Type=simple +User=postgres +ExecStart=/usr/local/bin/patroni /var/lib/pgsql/18/storefront-n1-689qacsi/configs/patroni.yaml +ExecReload=/bin/kill -s HUP $MAINPID +KillMode=process +TimeoutSec=30 +MemoryMax=8589934592 +Restart=on-failure +Environment="PATH=/usr/pgsql-18/bin:/root/.local/bin:/root/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/usr/local/go/bin" +Environment="PGSERVICEFILE=/var/lib/pgsql/18/storefront-n1-689qacsi/configs/pg_service.conf" + +[Install] +WantedBy=multi-user.target diff --git a/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/minimal.service b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/minimal.service new file mode 100644 index 00000000..859019f8 --- /dev/null +++ b/server/internal/orchestrator/systemd/golden_test/TestUnitOptions/PatroniUnitOptions/minimal.service @@ -0,0 +1,16 @@ +[Unit] +After=syslog.target network.target + +[Service] +Type=simple +User=postgres +ExecStart=/usr/local/bin/patroni /var/lib/pgsql/18/storefront-n1-689qacsi/configs/patroni.yaml +ExecReload=/bin/kill -s HUP $MAINPID +KillMode=process +TimeoutSec=30 +Restart=on-failure +Environment="PATH=/usr/pgsql-18/bin:/root/.local/bin:/root/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/usr/local/go/bin" +Environment="PGSERVICEFILE=/var/lib/pgsql/18/storefront-n1-689qacsi/configs/pg_service.conf" + +[Install] +WantedBy=multi-user.target diff --git a/server/internal/orchestrator/systemd/main_test.go b/server/internal/orchestrator/systemd/main_test.go new file mode 100644 index 00000000..968063fa --- /dev/null +++ b/server/internal/orchestrator/systemd/main_test.go @@ -0,0 +1,16 @@ +package systemd_test + +import ( + "flag" + "os" + "testing" +) + +var update bool + +func TestMain(m *testing.M) { + flag.BoolVar(&update, "update", false, "update golden test outputs") + flag.Parse() + + os.Exit(m.Run()) +} diff --git a/server/internal/orchestrator/systemd/orchestrator.go b/server/internal/orchestrator/systemd/orchestrator.go index dbde56a9..89007553 100644 --- a/server/internal/orchestrator/systemd/orchestrator.go +++ b/server/internal/orchestrator/systemd/orchestrator.go @@ -259,7 +259,7 @@ func (o *Orchestrator) GenerateInstanceResources(spec *database.InstanceSpec) (* DatabaseID: spec.DatabaseID, HostID: spec.HostID, Name: patroniServiceName(spec.InstanceID), - Options: patroniUnitOptions(paths, o.packageManager.BinDir(pgMajor)), + Options: PatroniUnitOptions(paths, o.packageManager.BinDir(pgMajor), spec.CPUs, spec.MemoryBytes), ExtraDependencies: []resource.Identifier{ patroniConfig.Identifier(), instanceDir.Identifier(), diff --git a/server/internal/orchestrator/systemd/patroni_unit.go b/server/internal/orchestrator/systemd/patroni_unit.go index 878cea45..05994ed2 100644 --- a/server/internal/orchestrator/systemd/patroni_unit.go +++ b/server/internal/orchestrator/systemd/patroni_unit.go @@ -6,71 +6,43 @@ import ( "path/filepath" "github.com/coreos/go-systemd/v22/unit" - "github.com/pgEdge/control-plane/server/internal/orchestrator/common" ) -func patroniUnitOptions(paths common.InstancePaths, pgBinPath string) []*unit.UnitOption { - pathEnv := "PATH=" + pgBinPath +func PatroniUnitOptions( + paths common.InstancePaths, + pgBinPath string, + cpus float64, + memoryBytes uint64, +) []*unit.UnitOption { + pathEnv := pgBinPath if p := os.Getenv("PATH"); p != "" { pathEnv += ":" + p } + patroniCmd := fmt.Sprintf("%s %s", paths.PatroniPath, paths.Instance.PatroniConfig()) + pgServiceFileEnv := filepath.Join(paths.Instance.Configs(), "pg_service.conf") - return []*unit.UnitOption{ - { - Section: "Unit", - Name: "After", - Value: "syslog.target network.target", - }, - { - Section: "Service", - Name: "Type", - Value: "simple", - }, - { - Section: "Service", - Name: "User", - Value: "postgres", - }, - { - Section: "Service", - Name: "ExecStart", - Value: fmt.Sprintf("%s %s", paths.PatroniPath, paths.Instance.PatroniConfig()), - }, - { - Section: "Service", - Name: "ExecReload", - Value: "/bin/kill -s HUP $MAINPID", - }, - { - Section: "Service", - Name: "KillMode", - Value: "process", - }, - { - Section: "Service", - Name: "TimeoutSec", - Value: "30", - }, - { - Section: "Service", - Name: "Restart", - Value: "on-failure", - }, - { - Section: "Service", - Name: "Environment", - Value: pathEnv, - }, - { - Section: "Service", - Name: "Environment", - Value: "PGSERVICEFILE=" + filepath.Join(paths.Instance.Configs(), "pg_service.conf"), - }, - { - Section: "Install", - Name: "WantedBy", - Value: "multi-user.target", - }, - } + return UnitFile{ + Unit: UnitSection{ + After: []string{"syslog.target", "network.target"}, + }, + Service: ServiceSection{ + Type: ServiceTypeSimple, + User: "postgres", + ExecStart: patroniCmd, + ExecReload: "/bin/kill -s HUP $MAINPID", + KillMode: ServiceKillModeProcess, + TimeoutSec: 30, + CPUs: cpus, + MemoryBytes: memoryBytes, + Restart: ServiceRestartOnFailure, + Environment: map[string]string{ + "PATH": pathEnv, + "PGSERVICEFILE": pgServiceFileEnv, + }, + }, + Install: InstallSection{ + WantedBy: []string{"multi-user.target"}, + }, + }.Options() } diff --git a/server/internal/orchestrator/systemd/unit_options.go b/server/internal/orchestrator/systemd/unit_options.go new file mode 100644 index 00000000..aacd8cac --- /dev/null +++ b/server/internal/orchestrator/systemd/unit_options.go @@ -0,0 +1,245 @@ +package systemd + +import ( + "fmt" + "maps" + "slices" + "strconv" + "strings" + + "github.com/coreos/go-systemd/v22/unit" +) + +type ServiceType string + +func (s ServiceType) String() string { + return string(s) +} + +const ( + ServiceTypeSimple ServiceType = "simple" + ServiceTypeExec ServiceType = "exec" + ServiceTypeForking ServiceType = "forking" + ServiceTypeOneShot ServiceType = "oneshot" + ServiceTypeNotify ServiceType = "notify" + ServiceTypeDBus ServiceType = "dbus" + ServiceTypeIdle ServiceType = "idle" + ServiceTypeNotifyReload ServiceType = "notify-reload" +) + +type ServiceKillMode string + +func (s ServiceKillMode) String() string { + return string(s) +} + +const ( + ServiceKillModeControlGroup ServiceKillMode = "control-group" + ServiceKillModeMixed ServiceKillMode = "mixed" + ServiceKillModeProcess ServiceKillMode = "process" + ServiceKillModeNone ServiceKillMode = "none" +) + +type ServiceRestart string + +func (s ServiceRestart) String() string { + return string(s) +} + +const ( + ServiceRestartNo ServiceRestart = "no" + ServiceRestartAlways ServiceRestart = "always" + ServiceRestartOnSuccess ServiceRestart = "on-success" + ServiceRestartOnFailure ServiceRestart = "on-failure" + ServiceRestartOnAbnormal ServiceRestart = "on-abnormal" + ServiceRestartOnAbort ServiceRestart = "on-abort" + ServiceRestartOnWatchdog ServiceRestart = "on-watchdog" +) + +type UnitFile struct { + Unit UnitSection + Service ServiceSection + Install InstallSection +} + +func (f UnitFile) Options() []*unit.UnitOption { + return slices.Concat( + f.Unit.Options(), + f.Service.Options(), + f.Install.Options(), + ) +} + +type UnitSection struct { + After []string +} + +func (u UnitSection) Options() []*unit.UnitOption { + var opts []*unit.UnitOption + if len(u.After) != 0 { + opts = append(opts, UnitAfterOption(u.After...)) + } + return opts +} + +type ServiceSection struct { + Type ServiceType + User string + ExecStart string + ExecReload string + KillMode ServiceKillMode + TimeoutSec int + CPUs float64 + MemoryBytes uint64 + Restart ServiceRestart + Environment map[string]string +} + +func (s ServiceSection) Options() []*unit.UnitOption { + var opts []*unit.UnitOption + if s.Type != "" { + opts = append(opts, ServiceTypeOption(s.Type)) + } + if s.User != "" { + opts = append(opts, ServiceUserOption(s.User)) + } + if s.ExecStart != "" { + opts = append(opts, ServiceExecStartOption(s.ExecStart)) + } + if s.ExecReload != "" { + opts = append(opts, ServiceExecReloadOption(s.ExecReload)) + } + if s.KillMode != "" { + opts = append(opts, ServiceKillModeOption(s.KillMode)) + } + if s.TimeoutSec != 0 { + opts = append(opts, ServiceTimeoutSecOption(s.TimeoutSec)) + } + if s.CPUs > 0 { + opts = append(opts, ServiceCPUQuotaOption(s.CPUs)) + } + if s.MemoryBytes != 0 { + opts = append(opts, ServiceMemoryMaxOption(s.MemoryBytes)) + } + if s.Restart != "" { + opts = append(opts, ServiceRestartOption(s.Restart)) + } + for _, name := range slices.Sorted(maps.Keys(s.Environment)) { + opts = append(opts, ServiceEnvironmentOption(name, s.Environment[name])) + } + return opts +} + +type InstallSection struct { + WantedBy []string +} + +func (s InstallSection) Options() []*unit.UnitOption { + var opts []*unit.UnitOption + if len(s.WantedBy) != 0 { + opts = append(opts, InstallWantedByOption(s.WantedBy...)) + } + return opts +} + +const ( + sectionNameUnit = "Unit" + sectionNameService = "Service" + sectionNameInstall = "Install" +) + +func UnitAfterOption(values ...string) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameUnit, + Name: "After", + Value: strings.Join(values, " "), + } +} + +func ServiceTypeOption(value ServiceType) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "Type", + Value: value.String(), + } +} + +func ServiceUserOption(value string) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "User", + Value: value, + } +} + +func ServiceExecStartOption(value string) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "ExecStart", + Value: value, + } +} + +func ServiceExecReloadOption(value string) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "ExecReload", + Value: value, + } +} + +func ServiceKillModeOption(value ServiceKillMode) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "KillMode", + Value: value.String(), + } +} + +func ServiceTimeoutSecOption(value int) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "TimeoutSec", + Value: strconv.Itoa(value), + } +} + +func ServiceRestartOption(value ServiceRestart) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "Restart", + Value: value.String(), + } +} + +func ServiceEnvironmentOption(name, value string) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "Environment", + Value: fmt.Sprintf("%q", name+"="+value), + } +} + +func ServiceCPUQuotaOption(cpus float64) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "CPUQuota", + Value: fmt.Sprintf("%.f%%", cpus*100), + } +} + +func ServiceMemoryMaxOption(memoryBytes uint64) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameService, + Name: "MemoryMax", + Value: strconv.FormatUint(memoryBytes, 10), + } +} + +func InstallWantedByOption(value ...string) *unit.UnitOption { + return &unit.UnitOption{ + Section: sectionNameInstall, + Name: "WantedBy", + Value: strings.Join(value, " "), + } +} diff --git a/server/internal/orchestrator/systemd/unit_options_test.go b/server/internal/orchestrator/systemd/unit_options_test.go new file mode 100644 index 00000000..9201deb9 --- /dev/null +++ b/server/internal/orchestrator/systemd/unit_options_test.go @@ -0,0 +1,95 @@ +package systemd_test + +import ( + "bytes" + "fmt" + "io" + "testing" + + "github.com/coreos/go-systemd/v22/unit" + "github.com/pgEdge/control-plane/server/internal/orchestrator/common" + "github.com/pgEdge/control-plane/server/internal/orchestrator/systemd" + "github.com/pgEdge/control-plane/server/internal/testutils" +) + +func TestUnitOptions(t *testing.T) { + golden := &testutils.GoldenTest[[]*unit.UnitOption]{ + FileExtension: ".service", + Marshal: func(v any) ([]byte, error) { + opts, ok := v.([]*unit.UnitOption) + if !ok { + return nil, fmt.Errorf("expected []*unit.UnitOption, but got %T", v) + } + + var buf bytes.Buffer + _, err := io.Copy(&buf, unit.Serialize(opts)) + if err != nil { + return nil, fmt.Errorf("failed to serialize unit options: %w", err) + } + + return buf.Bytes(), nil + }, + Unmarshal: func(data []byte, v any) error { + opts, ok := v.(*[]*unit.UnitOption) + if !ok { + return fmt.Errorf("expected *[]*unit.UnitOption, but got %T", v) + } + out, err := unit.Deserialize(bytes.NewReader(data)) + if err != nil { + return fmt.Errorf("failed to deserialize unit options: %w", err) + } + *opts = out + + return nil + }, + } + + t.Run("PatroniUnitOptions", func(t *testing.T) { + paths := common.InstancePaths{ + Instance: common.Paths{BaseDir: "/var/lib/pgsql/18/storefront-n1-689qacsi"}, + Host: common.Paths{BaseDir: "/var/lib/pgsql/18/storefront-n1-689qacsi"}, + PgBackRestPath: "/usr/bin/pgbackrest", + PatroniPath: "/usr/local/bin/patroni", + } + pgBinPath := "/usr/pgsql-18/bin" + + for _, tc := range []struct { + name string + paths common.InstancePaths + pgBinPath string + cpus float64 + memoryBytes uint64 + }{ + { + name: "minimal", + paths: paths, + pgBinPath: pgBinPath, + }, + { + name: "cpu limit", + paths: paths, + pgBinPath: pgBinPath, + cpus: 14, + }, + { + name: "fractional cpu limit", + paths: paths, + pgBinPath: pgBinPath, + cpus: 0.5, + }, + { + name: "memory max", + paths: paths, + pgBinPath: pgBinPath, + memoryBytes: 8_589_934_592, // 8GiB in bytes + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("PATH", "/root/.local/bin:/root/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/usr/local/go/bin") + + actual := systemd.PatroniUnitOptions(tc.paths, tc.pgBinPath, tc.cpus, tc.memoryBytes) + golden.Run(t, actual, update) + }) + } + }) +}