From ebd71ffa037a35f0c6a539cbea342159af678b77 Mon Sep 17 00:00:00 2001 From: balanza Date: Mon, 13 Jan 2025 23:08:53 +0100 Subject: [PATCH] test plugin cancellation --- Makefile | 3 +- plugin_examples/{ => dummy}/dummy.go | 2 +- plugin_examples/sleep/sleep.go | 65 ++++++++++++++ test/cli/plugins.bats | 127 +++++++++++++++++++++++++++ 4 files changed, 195 insertions(+), 2 deletions(-) rename plugin_examples/{ => dummy}/dummy.go (94%) create mode 100644 plugin_examples/sleep/sleep.go diff --git a/Makefile b/Makefile index 8f878a15..a8fd6d4e 100644 --- a/Makefile +++ b/Makefile @@ -25,7 +25,8 @@ agent: .PHONY: build-plugin-examples build-plugin-examples: - $(GO_BUILD) -o $(BUILD_DIR)/$(CURRENT_ARCH)/plugin_examples/dummy ./plugin_examples/dummy.go + $(GO_BUILD) -o $(BUILD_DIR)/$(CURRENT_ARCH)/plugin_examples/dummy ./plugin_examples/dummy/dummy.go + $(GO_BUILD) -o $(BUILD_DIR)/$(CURRENT_ARCH)/plugin_examples/sleep ./plugin_examples/sleep/sleep.go .PHONY: cross-compiled $(ARCHS) cross-compiled: $(ARCHS) diff --git a/plugin_examples/dummy.go b/plugin_examples/dummy/dummy.go similarity index 94% rename from plugin_examples/dummy.go rename to plugin_examples/dummy/dummy.go index dc5920e5..fe4f6241 100644 --- a/plugin_examples/dummy.go +++ b/plugin_examples/dummy/dummy.go @@ -1,6 +1,6 @@ package main -// go build -o /usr/etc/trento/dummy ./plugin_examples/dummy.go +// go build -o /usr/etc/trento/dummy ./plugin_examples/dummy/dummy.go import ( "context" diff --git a/plugin_examples/sleep/sleep.go b/plugin_examples/sleep/sleep.go new file mode 100644 index 00000000..2db142ed --- /dev/null +++ b/plugin_examples/sleep/sleep.go @@ -0,0 +1,65 @@ +package main + +// go build -o /usr/etc/trento/sleep ./plugin_examples/sleep/sleep.go + +import ( + "context" + "fmt" + "os/exec" + "sync" + + "github.com/hashicorp/go-plugin" + log "github.com/sirupsen/logrus" + "github.com/trento-project/agent/pkg/factsengine/entities" + "github.com/trento-project/agent/pkg/factsengine/plugininterface" +) + +type sleepGatherer struct { +} + +func (s sleepGatherer) Gather(ctx context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) { + facts := []entities.Fact{} + + log.Infof("Starting sleep plugin facts gathering process") + + wg := sync.WaitGroup{} + + for _, factReq := range factsRequests { + log.Infof("Sleeping for %s", factReq.Argument) + fact := entities.NewFactGatheredWithRequest(factReq, &entities.FactValueString{Value: fmt.Sprint(factReq.Argument)}) + facts = append(facts, fact) + + time := fmt.Sprint(factReq.Argument) + wg.Add(1) + go func(time string) { + defer wg.Done() + cmd := exec.CommandContext(ctx, "sleep", time) + cmd.Run() + }(time) + + } + + wg.Wait() + + log.Infof("Requested sleep plugin facts gathered") + return facts, nil +} + +func main() { + d := &sleepGatherer{} + + handshakeConfig := plugin.HandshakeConfig{ + ProtocolVersion: 1, + MagicCookieKey: "TRENTO_PLUGIN", + MagicCookieValue: "gatherer", + } + + var pluginMap = map[string]plugin.Plugin{ + "gatherer": &plugininterface.GathererPlugin{Impl: d}, + } + + plugin.Serve(&plugin.ServeConfig{ // nolint + HandshakeConfig: handshakeConfig, + Plugins: pluginMap, + }) +} diff --git a/test/cli/plugins.bats b/test/cli/plugins.bats index c899f60b..b9c29451 100644 --- a/test/cli/plugins.bats +++ b/test/cli/plugins.bats @@ -7,6 +7,10 @@ setup() { PATH="$BUILD_DIR:$PATH" } +teardown() { + # kill all the processes started by the test + pkill -P $$ || true +} @test "it should include the dummy plugin into list" { run trento-agent facts list --plugins-folder $BUILD_DIR/plugin_examples @@ -32,3 +36,126 @@ setup() { [ "$status" -eq 0 ] echo $output | grep -q "Name: 2" } + +@test "it should remove all the processes on complete" { + # declare the expected processes + cmd_agent="trento-agent facts gather --plugins-folder $BUILD_DIR/plugin_examples --gatherer sleep --argument 2s" + cmd_plugin="$BUILD_DIR/plugin_examples/sleep" + cmd_sleep="sleep 2s" + + # start the agent in background + eval "$cmd_agent &" + pid=$! + + # retrieve the pid of the exepcted process + pid_agent=$(pgrep -f "$cmd_agent") + pid_plugin=$(pgrep -f "$cmd_plugin") + pid_sleep=$(pgrep -f "$cmd_sleep") + + # double check the test is correct + [ $pid -eq $pid_agent ] + + # ensure no duplicated processes are running + assert_one "$pid_agent" + assert_one "$pid_plugin" + assert_one "$pid_sleep" + + # ensure the process tree is correct + assert_parent "$pid_plugin" "$pid_agent" + assert_parent "$pid_sleep" "$pid_plugin" + + # wait for the process to finish + while kill -0 $pid 2>/dev/null; do + sleep 1 + done + + # test processes are killed + assert_no_pid "$pid_agent" + assert_no_pid "$pid_plugin" + assert_no_pid "$pid_sleep" +} + +@test "it should remove all the processes on agent process stopped (SIGINT)" { + # declare the expected processes + cmd_agent="trento-agent facts gather --plugins-folder $BUILD_DIR/plugin_examples --gatherer sleep --argument 2s" + cmd_plugin="$BUILD_DIR/plugin_examples/sleep" + cmd_sleep="sleep 2s" + + # start the agent in background + eval "$cmd_agent &" + pid=$! + + # retrieve the pid of the exepcted process + pid_agent=$(pgrep -f "$cmd_agent") + pid_plugin=$(pgrep -f "$cmd_plugin") + pid_sleep=$(pgrep -f "$cmd_sleep") + + # double check the test is correct + [ $pid -eq $pid_agent ] + + # ensure no duplicated processes are running + assert_one "$pid_agent" + assert_one "$pid_plugin" + assert_one "$pid_sleep" + + # ensure the process tree is correct + assert_parent "$pid_plugin" "$pid_agent" + assert_parent "$pid_sleep" "$pid_plugin" + + # kill the agent + kill -INT $pid_agent + + # test processes are killed + assert_no_pid "$pid_agent" + assert_no_pid "$pid_plugin" + assert_no_pid "$pid_sleep" +} + +@test "it should remove all the processes on agent process stopped (SIGTERM)" { + # declare the expected processes + cmd_agent="trento-agent facts gather --plugins-folder $BUILD_DIR/plugin_examples --gatherer sleep --argument 2s" + cmd_plugin="$BUILD_DIR/plugin_examples/sleep" + cmd_sleep="sleep 2s" + + # start the agent in background + eval "$cmd_agent &" + pid=$! + + # retrieve the pid of the exepcted process + pid_agent=$(pgrep -f "$cmd_agent") + pid_plugin=$(pgrep -f "$cmd_plugin") + pid_sleep=$(pgrep -f "$cmd_sleep") + + # double check the test is correct + [ $pid -eq $pid_agent ] + + # ensure no duplicated processes are running + assert_one "$pid_agent" + assert_one "$pid_plugin" + assert_one "$pid_sleep" + + # ensure the process tree is correct + assert_parent "$pid_plugin" "$pid_agent" + assert_parent "$pid_sleep" "$pid_plugin" + + # kill the agent + kill -TERM $pid_agent + + # test processes are killed + assert_no_pid "$pid_agent" + assert_no_pid "$pid_plugin" + assert_no_pid "$pid_sleep" +} + + +function assert_one { + [ $(echo "$1" | wc -l) == 1 ] +} + +function assert_parent { + [ $(ps -o ppid= -p "$1") == "$2" ] +} + +function assert_no_pid { + [ $(ps -p "$1" | wc -l) == 1 ] +} \ No newline at end of file