From c1bc91d3e9acebde52a575a6bac2b04338dc2388 Mon Sep 17 00:00:00 2001
From: Tim Megela <megela@google.com>
Date: Fri, 18 Oct 2024 04:03:20 -0700
Subject: [PATCH] Refactors the start daemon code

PiperOrigin-RevId: 687246053
---
 cmd/main.go                                   |  8 +++----
 .../{startdaemon/startdaemon.go => daemon.go} | 24 ++++++++++++++-----
 2 files changed, 22 insertions(+), 10 deletions(-)
 rename internal/daemon/{startdaemon/startdaemon.go => daemon.go} (91%)

diff --git a/cmd/main.go b/cmd/main.go
index 06b06bf..354af92 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -28,7 +28,7 @@ import (
 	"go.uber.org/zap/zapcore"
 	"github.com/GoogleCloudPlatform/sapagent/shared/gce/metadataserver"
 	"github.com/GoogleCloudPlatform/sapagent/shared/log"
-	"github.com/GoogleCloudPlatform/workloadagent/internal/daemon/startdaemon"
+	"github.com/GoogleCloudPlatform/workloadagent/internal/daemon"
 	"github.com/GoogleCloudPlatform/workloadagent/internal/onetime/logusage"
 	"github.com/GoogleCloudPlatform/workloadagent/internal/onetime"
 	"github.com/GoogleCloudPlatform/workloadagent/internal/onetime/version"
@@ -65,10 +65,10 @@ func main() {
 	}
 	rootCmd.AddCommand(version.NewCommand())
 	rootCmd.AddCommand(logusage.NewCommand(lp, cloudProps))
-	daemon := startdaemon.NewDaemon(lp, cloudProps)
-	rootCmd.AddCommand(daemon)
+	d := daemon.NewDaemon(lp, cloudProps)
+	rootCmd.AddCommand(d)
 	// Add any additional windows or linux specific subcommands.
-	rootCmd.AddCommand(additionalSubcommands(ctx, daemon, lp, cloudProps)...)
+	rootCmd.AddCommand(additionalSubcommands(ctx, d, lp, cloudProps)...)
 
 	for _, cmd := range rootCmd.Commands() {
 		if cmd.Name() != "startdaemon" {
diff --git a/internal/daemon/startdaemon/startdaemon.go b/internal/daemon/daemon.go
similarity index 91%
rename from internal/daemon/startdaemon/startdaemon.go
rename to internal/daemon/daemon.go
index 499430a..03fe65e 100644
--- a/internal/daemon/startdaemon/startdaemon.go
+++ b/internal/daemon/daemon.go
@@ -14,8 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-// Package startdaemon implements startdaemon mode execution subcommand in Workload Agent.
-package startdaemon
+// Package daemon implements daemon mode execution subcommand in Workload Agent.
+package daemon
 
 import (
 	"context"
@@ -67,11 +67,25 @@ func NewDaemon(lp log.Parameters, cloudProps *cpb.CloudProperties) *cobra.Comman
 		Short: "Start daemon mode of the agent",
 		Long:  "startdaemon [--config <path-to-config-file>]",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return d.startdaemonHandler(cmd.Context())
+			return d.Execute(cmd.Context())
 		},
 	}
 	cmd.Flags().StringVar(&d.configFilePath, "config", "", "configuration path for startdaemon mode")
 	cmd.Flags().StringVar(&d.configFilePath, "c", "", "configuration path for startdaemon mode")
+	d.setupLogging()
+	return cmd
+}
+
+// Execute runs the daemon command.
+func (d *Daemon) Execute(ctx context.Context) error {
+	d.setupLogging()
+	// Run the daemon handler that will start any services
+	ctx, cancel := context.WithCancel(ctx)
+	d.startdaemonHandler(ctx, cancel)
+	return nil
+}
+
+func (d *Daemon) setupLogging() {
 	// Configure daemon logging with default values until the config file is loaded.
 	d.lp.CloudLogName = `google-cloud-workload-agent`
 	d.lp.LogFileName = `/var/log/google-cloud-workload-agent.log`
@@ -82,10 +96,9 @@ func NewDaemon(lp log.Parameters, cloudProps *cpb.CloudProperties) *cobra.Comman
 		os.Chmod(logDir, 0777)
 	}
 	log.SetupLogging(d.lp)
-	return cmd
 }
 
-func (d *Daemon) startdaemonHandler(ctx context.Context) error {
+func (d *Daemon) startdaemonHandler(ctx context.Context, cancel context.CancelFunc) error {
 	var err error
 	d.config, err = configuration.Load(d.configFilePath, os.ReadFile, d.cloudProps)
 	if err != nil {
@@ -118,7 +131,6 @@ func (d *Daemon) startdaemonHandler(ctx context.Context) error {
 
 	shutdownch := make(chan os.Signal, 1)
 	signal.Notify(shutdownch, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
-	ctx, cancel := context.WithCancel(ctx)
 
 	// Add any additional services here.
 	d.services = []Service{