Skip to content

Commit

Permalink
Implement memory monitoring
Browse files Browse the repository at this point in the history
  • Loading branch information
yawangwang committed Nov 16, 2023
1 parent 1fd3df5 commit ed6b88e
Show file tree
Hide file tree
Showing 13 changed files with 475 additions and 27 deletions.
3 changes: 1 addition & 2 deletions go.work.sum
Original file line number Diff line number Diff line change
Expand Up @@ -73,11 +73,10 @@ cloud.google.com/go/lifesciences v0.8.0/go.mod h1:lFxiEOMqII6XggGbOnKiyZ7IBwoIqA
cloud.google.com/go/logging v1.4.2 h1:Mu2Q75VBDQlW1HlBMjTX4X84UFR73G1TiLlRYc/b7tA=
cloud.google.com/go/logging v1.4.2/go.mod h1:jco9QZSx8HiVVqLJReq7z7bVdj0P1Jb9PDFs63T+axo=
github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
github.com/godbus/dbus v0.0.0-20190422162347-ade71ed3457e h1:BWhy2j3IXJhjCbC68FptL43tDKIq8FladmaTs3Xs7Z8=
github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE=
github.com/googleapis/gax-go v2.0.2+incompatible h1:silFMLAnr330+NRuag/VjIGF7TLp/LBrV2CJKFLWEww=
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU=
golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
26 changes: 26 additions & 0 deletions launcher/container_runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,9 @@ import (
"github.com/google/go-tpm-tools/cel"
"github.com/google/go-tpm-tools/client"
"github.com/google/go-tpm-tools/launcher/agent"
npd "github.com/google/go-tpm-tools/launcher/internal/healthmonitoring/nodeproblemdetector"
"github.com/google/go-tpm-tools/launcher/internal/signaturediscovery"
"github.com/google/go-tpm-tools/launcher/internal/systemctl"
"github.com/google/go-tpm-tools/launcher/launcherfile"
"github.com/google/go-tpm-tools/launcher/spec"
"github.com/google/go-tpm-tools/launcher/teeserver"
Expand All @@ -55,6 +57,8 @@ const tokenFileTmp = ".token.tmp"

const teeServerSocket = "teeserver.sock"

const systemStatsConfigFilePath = "/etc/node_problem_detector/system-stats-monitor.json"

// Since we only allow one container on a VM, using a deterministic id is probably fine
const (
containerID = "tee-container"
Expand Down Expand Up @@ -507,6 +511,28 @@ func (r *ContainerRunner) Run(ctx context.Context) error {
defer teeServer.Shutdown(ctx)
}

// customize node-problem-detector.service and start it.
if r.launchSpec.MemoryMonitoringEnabled {
r.logger.Println("MemoryMonitoring is enabled")
config := npd.NewSystemStatsConfig()
// collects "memory/bytes_used" metrics only when memory monitoring enabled.
config.EnableMemoryBytesUsed()
// override the default config file.
if err := config.WriteFile(systemStatsConfigFilePath); err != nil {
return fmt.Errorf("failed to override the default config file [%s] for node-problem-detector: %v", systemStatsConfigFilePath, err)
}
s, err := systemctl.New()
if err != nil {
return fmt.Errorf("failed to create systemctl client: %v", err)
}
defer s.Close()

if err := s.Start("node-problem-detector.service"); err != nil {
return fmt.Errorf("failed to start node-problem-detector.service: %v", err)
}
r.logger.Println("node-problem-detector.service successfully started.")
}

var streamOpt cio.Opt
switch r.launchSpec.LogRedirect {
case spec.Nowhere:
Expand Down
3 changes: 3 additions & 0 deletions launcher/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,10 @@ require (
github.com/containerd/fifo v1.0.0 // indirect
github.com/containerd/ttrpc v1.1.0 // indirect
github.com/containerd/typeurl v1.0.2 // indirect
github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf // indirect
github.com/coreos/go-systemd/v22 v22.5.0 // indirect
github.com/docker/go-events v0.0.0-20190806004212-e31b211e4f1c // indirect
github.com/godbus/dbus v0.0.0-20190422162347-ade71ed3457e // indirect
github.com/gogo/googleapis v1.4.1 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
Expand Down
11 changes: 11 additions & 0 deletions launcher/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -336,9 +336,13 @@ github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3Ee
github.com/coreos/go-systemd v0.0.0-20161114122254-48702e0da86b/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
github.com/coreos/go-systemd v0.0.0-20180511133405-39ca1b05acc7/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf h1:iW4rZ826su+pqaw19uhpSCzhj44qo35pNgKFGqzDKkU=
github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
github.com/coreos/go-systemd/v22 v22.0.0/go.mod h1:xO0FLkIi5MaZafQlIrOotqXZ90ih+1atmu1JpKERPPk=
github.com/coreos/go-systemd/v22 v22.1.0/go.mod h1:xO0FLkIi5MaZafQlIrOotqXZ90ih+1atmu1JpKERPPk=
github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
github.com/coreos/go-systemd/v22 v22.5.0 h1:RrqgGjYQKalulkV8NGVIfkXQf6YYmOyiJKk8iXXhfZs=
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
github.com/coreos/pkg v0.0.0-20160727233714-3ac0863d7acf/go.mod h1:E3G3o1h8I7cfcXa63jLwjI0eiQQMgzzUDFVpN/nH/eA=
github.com/coreos/pkg v0.0.0-20180928190104-399ea9e2e55f/go.mod h1:E3G3o1h8I7cfcXa63jLwjI0eiQQMgzzUDFVpN/nH/eA=
github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE=
Expand Down Expand Up @@ -459,6 +463,7 @@ github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LB
github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
github.com/godbus/dbus v0.0.0-20151105175453-c7fdd8b5cd55/go.mod h1:/YcGZj5zSblfDWMMoOzV4fas9FZnQYTkDnsGvmh2Grw=
github.com/godbus/dbus v0.0.0-20180201030542-885f9cc04c9c/go.mod h1:/YcGZj5zSblfDWMMoOzV4fas9FZnQYTkDnsGvmh2Grw=
github.com/godbus/dbus v0.0.0-20190422162347-ade71ed3457e h1:BWhy2j3IXJhjCbC68FptL43tDKIq8FladmaTs3Xs7Z8=
github.com/godbus/dbus v0.0.0-20190422162347-ade71ed3457e/go.mod h1:bBOAhwG1umN6/6ZUMtDFBMQR8jRg9O75tm9K00oMsK4=
github.com/godbus/dbus/v5 v5.0.3/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
Expand Down Expand Up @@ -547,6 +552,7 @@ github.com/google/go-querystring v1.0.0/go.mod h1:odCYkC5MyYFN7vkCjXpyrEuKhc/BUO
github.com/google/go-replayers/grpcreplay v0.1.0/go.mod h1:8Ig2Idjpr6gifRd6pNVggX6TC1Zw6Jx74AKp7QNH2QE=
github.com/google/go-replayers/httpreplay v0.1.0/go.mod h1:YKZViNhiGgqdBlUbI2MwGpq4pXxNmhJLPHQ7cv2b5no=
github.com/google/go-sev-guest v0.9.3 h1:GOJ+EipURdeWFl/YYdgcCxyPeMgQUWlI056iFkBD8UU=
github.com/google/go-sev-guest v0.9.3/go.mod h1:hc1R4R6f8+NcJwITs0L90fYWTsBpd1Ix+Gur15sqHDs=
github.com/google/go-tpm v0.9.0 h1:sQF6YqWMi+SCXpsmS3fd21oPy/vSddwZry4JnmltHVk=
github.com/google/go-tpm v0.9.0/go.mod h1:FkNVkc6C+IsvDI9Jw1OveJmxGZUUaKxtrpOS47QWKfU=
github.com/google/go-tpm-tools v0.3.13-0.20230620182252-4639ecce2aba h1:qJEJcuLzH5KDR0gKc0zcktin6KSAwL7+jWKBYceddTc=
Expand Down Expand Up @@ -592,6 +598,7 @@ github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.2.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4=
github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/wire v0.3.0/go.mod h1:i1DMg/Lu8Sz5yYl25iOdmc5CT5qusaa+zmRWs16741s=
github.com/googleapis/enterprise-certificate-proxy v0.2.4 h1:uGy6JWR/uMIILU8wbf+OkstIrNiMjGpEIyhx8f6W7s4=
github.com/googleapis/enterprise-certificate-proxy v0.2.4/go.mod h1:AwSRAtLfXpU5Nm3pW+v7rGDHp09LsPtGY9MduiEsR9k=
Expand Down Expand Up @@ -881,6 +888,7 @@ github.com/pact-foundation/pact-go v1.0.4/go.mod h1:uExwJY4kCzNPcHRj+hCR/HBbOOIw
github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
github.com/pborman/uuid v1.2.0/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
github.com/pborman/uuid v1.2.1 h1:+ZZIw58t/ozdjRaXh/3awHfmWRbzYxJoAdNJxe/3pvw=
github.com/pborman/uuid v1.2.1/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
github.com/pelletier/go-buffruneio v0.2.0/go.mod h1:JkE26KsDizTr40EUHkXVtNPvgGtbSNq5BcowyYOWdKo=
github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
github.com/pelletier/go-toml v1.8.1/go.mod h1:T2/BmBdy8dvIRq1a/8aqjN41wvWlN4lrapLU/GW4pbc=
Expand Down Expand Up @@ -1171,6 +1179,7 @@ golang.org/x/crypto v0.0.0-20210322153248-0c34fe9e7dc2/go.mod h1:T9bdIzuCu7OtxOm
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
golang.org/x/crypto v0.0.0-20220314234659-1baeb1ce4c0b/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
golang.org/x/crypto v0.13.0 h1:mvySKfSWJ+UKUii46M40LOvyWfN0s2U+46/jDd0e6Ck=
golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/exp v0.0.0-20190510132918-efd6b22b2522/go.mod h1:ZjyILWgesfNpC6sMxTJOJm9Kp84zZh5NQWvqDGG3Qr8=
Expand Down Expand Up @@ -1414,6 +1423,7 @@ golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o=
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
Expand All @@ -1427,6 +1437,7 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
Expand Down
1 change: 0 additions & 1 deletion launcher/image/preload.sh
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,6 @@ configure_systemd_units_for_hardened() {
disable_unit "konlet-startup.service"
disable_unit "crash-reporter.service"
disable_unit "device_policy_manager.service"
disable_unit "node-problem-detector.service"
disable_unit "docker-events-collector-fluent-bit.service"
disable_unit "sshd.service"
disable_unit "var-lib-toolbox.mount"
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
// Package nodeproblemdetector provides configurations for node-problem-detector.service.
package nodeproblemdetector

import (
"encoding/json"
"fmt"
"os"
"time"
)

var defaultInvokeIntervalString = (60 * time.Second).String()

type metricConfig struct {
DisplayName string `json:"displayName"`
}

type memoryStatsConfig struct {
MetricsConfigs map[string]metricConfig `json:"metricsConfigs"`
}

// SystemStatsConfig contains configurations for `System Stats Monitor`,
// a problem daemon in node-problem-detector that collects pre-defined health-related metrics from different system components.
// For now we only consider collecting memory related metrics.
// View the comprehensive configuration details on https://github.com/kubernetes/node-problem-detector/tree/master/pkg/systemstatsmonitor#detailed-configuration-options
type SystemStatsConfig struct {
MemoryStatsConfig memoryStatsConfig `json:"memory"`
InvokeInterval string `json:"invokeInterval"`
}

// NewSystemStatsConfig returns a new SystemStatsConfig struct with default configurations.
func NewSystemStatsConfig() SystemStatsConfig {
return SystemStatsConfig{
MemoryStatsConfig: memoryStatsConfig{MetricsConfigs: map[string]metricConfig{}},
InvokeInterval: defaultInvokeIntervalString,
}
}

// EnableMemoryBytesUsed enables "memory/bytes_used" for memory monitoring.
func (ssc *SystemStatsConfig) EnableMemoryBytesUsed() {
ssc.MemoryStatsConfig.MetricsConfigs["memory/bytes_used"] = metricConfig{DisplayName: "memory/bytes_used"}
}

// WithInvokeInterval overrides the default invokeInterval.
func (ssc *SystemStatsConfig) WithInvokeInterval(interval time.Duration) {
ssc.InvokeInterval = interval.String()
}

// WriteFile writes systemStatsConfig data to the named file, creating it if necessary.
func (ssc *SystemStatsConfig) WriteFile(path string) error {
bytes, err := json.Marshal(ssc)
if err != nil {
return fmt.Errorf("failed to marshal struct [%v]: %w", ssc, err)
}
return os.WriteFile(path, bytes, 0644)
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
package nodeproblemdetector

import (
"bytes"
"io"
"os"
"path"
"testing"
"time"

"github.com/google/go-cmp/cmp"
)

func TestEnableMemoryBytesUsed(t *testing.T) {
got := NewSystemStatsConfig()
got.EnableMemoryBytesUsed()

want := SystemStatsConfig{
MemoryStatsConfig: memoryStatsConfig{
MetricsConfigs: map[string]metricConfig{
"memory/bytes_used": {DisplayName: "memory/bytes_used"},
},
},
InvokeInterval: defaultInvokeIntervalString,
}
if !cmp.Equal(got, want) {
t.Errorf("EnableMemoryBytesUsed() failed, got: %v, want: %v", got, want)
}
}

func TestWithInvokeInterval(t *testing.T) {
got := SystemStatsConfig{}
got.WithInvokeInterval(2 * time.Second)

want := SystemStatsConfig{InvokeInterval: (2 * time.Second).String()}
if !cmp.Equal(got, want) {
t.Errorf("WithInvokeInterval() failed, got: %v, want: %v", got, want)
}
}

func TestWriteFile(t *testing.T) {
tmpDir := t.TempDir()
tmpConfigFile := path.Join(tmpDir, "system-stats-monitor.json")

config := NewSystemStatsConfig()
config.EnableMemoryBytesUsed()
if err := config.WriteFile(tmpConfigFile); err != nil {
t.Fatalf("WriteFile() failed: %v", err)
}

file, err := os.OpenFile(tmpConfigFile, os.O_RDONLY, 0)
if err != nil {
t.Fatalf("failed to open file %s: %v", tmpConfigFile, err)
}

gotBytes, err := io.ReadAll(file)
if err != nil {
t.Fatalf("failed to read from file %s: %v", tmpConfigFile, err)
}

wantBytes := []byte(`{"memory":{"metricsConfigs":{"memory/bytes_used":{"displayName":"memory/bytes_used"}}},"invokeInterval":"1m0s"}`)
if !bytes.Equal(gotBytes, wantBytes) {
t.Errorf("WriteFile() did not write expected contents, got %s, want %s", gotBytes, wantBytes)
}
}
62 changes: 62 additions & 0 deletions launcher/internal/systemctl/systemctl.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
// Package systemctl implements a subset of systemctl operations.
package systemctl

import (
"fmt"
"log"

"github.com/coreos/go-systemd/dbus"
)

// Systemd is an interface to connect to host systemd with selected functions.
type Systemd interface {
Start(string) error
Stop(string) error
Close()
}

// Systemctl is a wrap around of dbus.Conn and implements the Systemd interface.
type Systemctl struct {
dbus *dbus.Conn
}

var _ Systemd = (*Systemctl)(nil)

// New connects to systemd over dbus.
func New() (*Systemctl, error) {
conn, err := dbus.New()
if err != nil {
return nil, err
}
return &Systemctl{dbus: conn}, nil
}

// Start is the equivalent of `systemctl start $unit`.
func (s *Systemctl) Start(unit string) error {
return runSystemdCmd(s.dbus.StartUnit, "start", unit)
}

// Stop is the equivalent of `systemctl stop $unit`.
func (s *Systemctl) Stop(unit string) error {
return runSystemdCmd(s.dbus.StopUnit, "stop", unit)
}

// Close disconnects from dbus.
func (s *Systemctl) Close() { s.dbus.Close() }

func runSystemdCmd(cmdFunc func(string, string, chan<- string) (int, error), cmd string, unit string) error {
progress := make(chan string, 1)

// Run systemd command in "replace" mode to start the unit and its dependencies,
// possibly replacing already queued jobs that conflict w∏ith this.
if _, err := cmdFunc(unit, "replace", progress); err != nil {
return fmt.Errorf("failed to run systemctl [%s] for unit [%s]: %v", cmd, unit, err)
}

if result := <-progress; result != "done" {
return fmt.Errorf("systemctl [%s] result was [%s], want done", cmd, result)
}

log.Printf("Finished up systemctl [%s] for unit [%s]", cmd, unit)
return nil
}
50 changes: 50 additions & 0 deletions launcher/internal/systemctl/systemctl_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
package systemctl

import (
"errors"
"testing"
)

func TestRunSystmedCmd(t *testing.T) {
doneUnitFunc := func(unit string, mode string, progress chan<- string) (int, error) {
progress <- "done"
return 1, nil
}
failedCallUnitFunc := func(unit string, mode string, progress chan<- string) (int, error) {
return 1, errors.New("something went wrong")
}
failedUnitFunc := func(unit string, mode string, progress chan<- string) (int, error) {
progress <- "failed"
return 1, nil
}

testCases := []struct {
name string
sytemdCmdFunc func(unit string, flag string, progress chan<- string) (int, error)
wantErr bool
}{
{
name: "success",
sytemdCmdFunc: doneUnitFunc,
wantErr: false,
},
{
name: "failed call",
sytemdCmdFunc: failedCallUnitFunc,
wantErr: true,
},
{
name: "failed unit run",
sytemdCmdFunc: failedUnitFunc,
wantErr: true,
},
}

for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
if err := runSystemdCmd(tc.sytemdCmdFunc, "test", "test_unit"); (err != nil) != tc.wantErr {
t.Errorf("runSystemdCmd() did not return expected error, got error: %v, but wantErr %v", err, tc.wantErr)
}
})
}
}
Loading

0 comments on commit ed6b88e

Please sign in to comment.