fix: handle normal teardown procedure

Signed-off-by: Zespre Chang <zespre.chang@suse.com>

cmd/agent/run.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sync/errgroup"
@@ -33,10 +34,36 @@ func run(ctx context.Context, options *config.AgentOptions) error {
 		return agent.Run(egctx)
 	})
 
-	eg.Go(func() error {
-		<-egctx.Done()
-		return s.Stop(egctx)
-	})
+	errCh := cleanup(egctx, s)
+
+	if err := eg.Wait(); err != nil {
+		return err
+	}
+
+	// Return cleanup error message if any
+	if err := <-errCh; err != nil {
+		return err
+	}
+
+	logrus.Info("finished clean")
+
+	return nil
+}
+
+func cleanup(ctx context.Context, srv *server.HTTPServer) <-chan error {
+	errCh := make(chan error)
+
+	go func() {
+		defer close(errCh)
+		<-ctx.Done()
+
+		if errors.Is(ctx.Err(), context.Canceled) {
+			errCh <- srv.Stop(ctx)
+			return
+		}
+
+		errCh <- nil
+	}()
 
-	return eg.Wait()
+	return errCh
 }

cmd/controller/run.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"log"
 
 	"github.com/rancher/wrangler/pkg/leader"
@@ -71,11 +72,6 @@ func run(options *config.ControllerOptions) error {
 		return s.Run()
 	})
 
-	eg.Go(func() error {
-		<-egctx.Done()
-		return s.Stop(egctx)
-	})
-
 	eg.Go(func() error {
 		if noLeaderElection {
 			callback(egctx)
@@ -85,5 +81,36 @@ func run(options *config.ControllerOptions) error {
 		return nil
 	})
 
-	return eg.Wait()
+	errCh := cleanup(egctx, s)
+
+	if err := eg.Wait(); err != nil {
+		return err
+	}
+
+	// Return cleanup error message if any
+	if err := <-errCh; err != nil {
+		return err
+	}
+
+	logrus.Info("finished clean")
+
+	return nil
+}
+
+func cleanup(ctx context.Context, srv *server.HTTPServer) <-chan error {
+	errCh := make(chan error)
+
+	go func() {
+		defer close(errCh)
+		<-ctx.Done()
+
+		if errors.Is(ctx.Err(), context.Canceled) {
+			errCh <- srv.Stop(ctx)
+			return
+		}
+
+		errCh <- nil
+	}()
+
+	return errCh
 }