[SVLS-4179] Updates serverless to use fxutil.OneShot

cmd/serverless/main.go

@@ -37,6 +37,7 @@ import (
 	"github.com/DataDog/datadog-agent/pkg/serverless/trace"
 	"github.com/DataDog/datadog-agent/pkg/serverless/trace/inferredspan"
 	"github.com/DataDog/datadog-agent/pkg/util/flavor"
+	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
 	"github.com/DataDog/datadog-agent/pkg/util/log"
 )
 
@@ -70,32 +71,28 @@ const (
 )
 
 func main() {
-	flavor.SetFlavor(flavor.ServerlessAgent)
-	config.Datadog.Set("use_v2_api.series", false, model.SourceAgentRuntime)
-	stopCh := make(chan struct{})
-
-	// Disable remote configuration for now as it just spams the debug logs
-	// and provides no value.
-	os.Setenv("DD_REMOTE_CONFIGURATION_ENABLED", "false")
-
 	// run the agent
-	serverlessDaemon, err := runAgent(stopCh)
+	err := fxutil.OneShot(runAgent)
+
 	if err != nil {
 		log.Error(err)
 		os.Exit(-1)
 	}
-
-	// handle SIGTERM signal
-	go handleSignals(serverlessDaemon, stopCh)
-
-	// block here until we receive a stop signal
-	<-stopCh
 }
 
-func runAgent(stopCh chan struct{}) (serverlessDaemon *daemon.Daemon, err error) {
+func runAgent() (err error) {
 
 	startTime := time.Now()
 
+	stopCh := make(chan struct{})
+
+	flavor.SetFlavor(flavor.ServerlessAgent)
+	config.Datadog.Set("use_v2_api.series", false, model.SourceAgentRuntime)
+
+	// Disable remote configuration for now as it just spams the debug logs
+	// and provides no value.
+	os.Setenv("DD_REMOTE_CONFIGURATION_ENABLED", "false")
+
 	// setup logger
 	// -----------
 
@@ -132,11 +129,11 @@ func runAgent(stopCh chan struct{}) (serverlessDaemon *daemon.Daemon, err error)
 		if processError != nil {
 			log.Errorf("Can't process events: %s", processError)
 		}
-		return nil, nil
+		return nil
 	}
 
 	// immediately starts the communication server
-	serverlessDaemon = daemon.StartDaemon(httpServerAddr)
+	serverlessDaemon := daemon.StartDaemon(httpServerAddr)
 	serverlessDaemon.ExecutionContext.SetInitializationTime(startTime)
 	err = serverlessDaemon.ExecutionContext.RestoreCurrentStateFromFile()
 	if err != nil {
@@ -347,21 +344,24 @@ func runAgent(stopCh chan struct{}) (serverlessDaemon *daemon.Daemon, err error)
 		}
 	}()
 
+	// setup a channel to catch SIGINT and SIGTERM signals
+	signals := []os.Signal{syscall.SIGINT, syscall.SIGTERM}
+	signalCh := make(chan os.Signal, 1)
+	signal.Notify(signalCh, signals[0], signals[1])
+	go handleTerminationSignals(serverlessDaemon, stopCh, signalCh)
+
 	// this log line is used for performance checks during CI
 	// please be careful before modifying/removing it
 	log.Debugf("serverless agent ready in %v", time.Since(startTime))
+
+	// block here until we receive a stop signal
+	<-stopCh
 	return
 }
 
-// handleSignals handles OS signals, if a SIGTERM is received,
-// the serverless agent stops.
-func handleSignals(serverlessDaemon *daemon.Daemon, stopCh chan struct{}) {
-	// setup a channel to catch OS signals
-	signalCh := make(chan os.Signal, 1)
-	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
-
-	// block here until we receive the interrupt signal
-	// when received, shutdown the serverless agent.
+// handleTerminationSignals handles OS termination signals.
+// If a specified signal is received the serverless agent stops.
+func handleTerminationSignals(serverlessDaemon *daemon.Daemon, stopCh chan struct{}, signalCh chan os.Signal) {
 	for signo := range signalCh {
 		switch signo {
 		default:

cmd/serverless/main_test.go

@@ -9,83 +9,47 @@ package main
 
 import (
 	"os"
-	"path"
-	"strings"
+	"syscall"
 	"testing"
+	"time"
 
-	"github.com/stretchr/testify/assert"
-
-	"github.com/DataDog/datadog-agent/pkg/config"
+	"github.com/DataDog/datadog-agent/pkg/serverless/daemon"
 )
 
-func setupTest() {
-	config.Datadog = config.NewConfig("datadog", "DD", strings.NewReplacer(".", "_"))
-	config.InitConfig(config.Datadog)
-	os.Setenv("AWS_LAMBDA_FUNCTION_NAME", "TestFunction")
-}
-
-func TestProxyNotLoaded(t *testing.T) {
-	setupTest()
-
-	proxyHTTP := "a:1"
-	proxyHTTPS := "a:2"
-	t.Setenv("DD_PROXY_HTTP", proxyHTTP)
-	t.Setenv("DD_PROXY_HTTPS", proxyHTTPS)
-	proxyHTTPConfig := config.Datadog.GetString("proxy.http")
-	proxyHTTPSConfig := config.Datadog.GetString("proxy.https")
-	assert.Equal(t, 0, len(proxyHTTPConfig))
-	assert.Equal(t, 0, len(proxyHTTPSConfig))
-}
-
-func TestProxyLoadedFromEnvVars(t *testing.T) {
-	setupTest()
+func TestDaemonStopOnSIGINT(t *testing.T) {
+	stopCh := make(chan struct{})
+	signalCh := make(chan os.Signal, 1)
 
-	proxyHTTP := "b:1"
-	proxyHTTPS := "b:2"
-	t.Setenv("DD_PROXY_HTTP", proxyHTTP)
-	t.Setenv("DD_PROXY_HTTPS", proxyHTTPS)
+	d := daemon.StartDaemon("http://localhost:8124")
+	go handleTerminationSignals(d, stopCh, signalCh)
 
-	config.Load()
-	proxyHTTPConfig := config.Datadog.GetString("proxy.http")
-	proxyHTTPSConfig := config.Datadog.GetString("proxy.https")
+	signalCh <- syscall.SIGINT
 
-	assert.Equal(t, proxyHTTP, proxyHTTPConfig)
-	assert.Equal(t, proxyHTTPS, proxyHTTPSConfig)
+	// Use t.Run with a timeout to allow the goroutine to finish
+	t.Run("WaitForStop", func(t *testing.T) {
+		select {
+		case <-stopCh: // Expected behavior, the daemon should be stopped
+		case <-time.After(1000 * time.Millisecond):
+			t.Error("Timeout waiting for daemon to stop")
+		}
+	})
 }
 
-func TestProxyLoadedFromConfigFile(t *testing.T) {
-	setupTest()
-
-	tempDir := t.TempDir()
-	configTest := path.Join(tempDir, "datadog.yaml")
-	os.WriteFile(configTest, []byte("proxy:\n  http: \"c:1\"\n  https: \"c:2\""), 0644)
-
-	config.Datadog.AddConfigPath(tempDir)
-	config.Load()
-	proxyHTTPConfig := config.Datadog.GetString("proxy.http")
-	proxyHTTPSConfig := config.Datadog.GetString("proxy.https")
-
-	assert.Equal(t, "c:1", proxyHTTPConfig)
-	assert.Equal(t, "c:2", proxyHTTPSConfig)
-}
-
-func TestProxyLoadedFromConfigFileAndEnvVars(t *testing.T) {
-	setupTest()
-
-	proxyHTTPEnvVar := "d:1"
-	proxyHTTPSEnvVar := "d:2"
-	t.Setenv("DD_PROXY_HTTP", proxyHTTPEnvVar)
-	t.Setenv("DD_PROXY_HTTPS", proxyHTTPSEnvVar)
+func TestDaemonStopOnSIGTERM(t *testing.T) {
+	stopCh := make(chan struct{})
+	signalCh := make(chan os.Signal, 1)
 
-	tempDir := t.TempDir()
-	configTest := path.Join(tempDir, "datadog.yaml")
-	os.WriteFile(configTest, []byte("proxy:\n  http: \"e:1\"\n  https: \"e:2\""), 0644)
+	d := daemon.StartDaemon("http://localhost:8124")
+	go handleTerminationSignals(d, stopCh, signalCh)
 
-	config.Datadog.AddConfigPath(tempDir)
-	config.Load()
-	proxyHTTPConfig := config.Datadog.GetString("proxy.http")
-	proxyHTTPSConfig := config.Datadog.GetString("proxy.https")
+	signalCh <- syscall.SIGTERM
 
-	assert.Equal(t, proxyHTTPEnvVar, proxyHTTPConfig)
-	assert.Equal(t, proxyHTTPSEnvVar, proxyHTTPSConfig)
+	// Use t.Run with a timeout to allow the goroutine to finish
+	t.Run("WaitForStop", func(t *testing.T) {
+		select {
+		case <-stopCh: // Expected behavior, the daemon should be stopped
+		case <-time.After(1000 * time.Millisecond):
+			t.Error("Timeout waiting for daemon to stop")
+		}
+	})
 }