fix: terminate goroutines gracefully

cmd/agent/root.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"context"
 	"fmt"
 	"os"
 
@@ -44,7 +43,6 @@ var rootCmd = &cobra.Command{
 		}
 	},
 	Run: func(cmd *cobra.Command, args []string) {
-		ctx := context.Background()
 		ipPoolNamespace, ipPoolName := kv.RSplit(ippoolRef, "/")
 		options := &config.AgentOptions{
 			DryRun:         dryRun,
@@ -57,7 +55,7 @@ var rootCmd = &cobra.Command{
 			},
 		}
 
-		if err := run(ctx, options); err != nil {
+		if err := run(options); err != nil {
 			fmt.Fprintf(os.Stderr, "%s\n", err.Error())
 			os.Exit(1)
 		}

cmd/agent/run.go

@@ -1,27 +1,54 @@
 package main
 
 import (
-	"context"
+	"errors"
+	"net/http"
 
+	"github.com/rancher/wrangler/pkg/signals"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/harvester/vm-dhcp-controller/pkg/agent"
 	"github.com/harvester/vm-dhcp-controller/pkg/config"
 	"github.com/harvester/vm-dhcp-controller/pkg/server"
 )
 
-func run(ctx context.Context, options *config.AgentOptions) error {
+func run(options *config.AgentOptions) error {
 	logrus.Infof("Starting VM DHCP Agent: %s", name)
 
-	agent := agent.NewAgent(ctx, options)
+	ctx := signals.SetupSignalContext()
+
+	agent := agent.NewAgent(options)
 
 	httpServerOptions := config.HTTPServerOptions{
 		DebugMode:     enableCacheDumpAPI,
 		DHCPAllocator: agent.DHCPAllocator,
 	}
 	s := server.NewHTTPServer(&httpServerOptions)
 	s.RegisterAgentHandlers()
-	go s.Run()
 
-	return agent.Run()
+	eg, egctx := errgroup.WithContext(ctx)
+
+	eg.Go(func() error {
+		return s.Run()
+	})
+
+	eg.Go(func() error {
+		return agent.Run(egctx)
+	})
+
+	errCh := server.Cleanup(egctx, s)
+
+	if err := eg.Wait(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+		return err
+	}
+
+	// Return cleanup error message if any
+	if err := <-errCh; err != nil {
+		return err
+	}
+
+	logrus.Info("finished clean")
+
+	return nil
 }

cmd/controller/run.go

@@ -2,11 +2,14 @@ package main
 
 import (
 	"context"
+	"errors"
 	"log"
+	"net/http"
 
 	"github.com/rancher/wrangler/pkg/leader"
 	"github.com/rancher/wrangler/pkg/signals"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/tools/clientcmd"
 
@@ -63,13 +66,34 @@ func run(options *config.ControllerOptions) error {
 	}
 	s := server.NewHTTPServer(&httpServerOptions)
 	s.RegisterControllerHandlers()
-	go s.Run()
 
-	if noLeaderElection {
-		callback(ctx)
-	} else {
-		leader.RunOrDie(ctx, "kube-system", "vm-dhcp-controllers", client, callback)
+	eg, egctx := errgroup.WithContext(ctx)
+
+	eg.Go(func() error {
+		return s.Run()
+	})
+
+	eg.Go(func() error {
+		if noLeaderElection {
+			callback(egctx)
+		} else {
+			leader.RunOrDie(egctx, "kube-system", "vm-dhcp-controllers", client, callback)
+		}
+		return nil
+	})
+
+	errCh := server.Cleanup(egctx, s)
+
+	if err := eg.Wait(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+		return err
+	}
+
+	// Return cleanup error message if any
+	if err := <-errCh; err != nil {
+		return err
 	}
 
+	logrus.Info("finished clean")
+
 	return nil
 }

cmd/webhook/run.go

@@ -81,5 +81,7 @@ func run(ctx context.Context, cfg *rest.Config, options *config.Options) error {
 
 	<-ctx.Done()
 
+	logrus.Info("Stopping webhook server")
+
 	return nil
 }

pkg/agent/agent.go

@@ -15,8 +15,6 @@ import (
 const DefaultNetworkInterface = "eth1"
 
 type Agent struct {
-	ctx context.Context
-
 	dryRun  bool
 	nic     string
 	poolRef types.NamespacedName
@@ -26,20 +24,17 @@ type Agent struct {
 	poolCache          map[string]string
 }
 
-func NewAgent(ctx context.Context, options *config.AgentOptions) *Agent {
+func NewAgent(options *config.AgentOptions) *Agent {
 	dhcpAllocator := dhcp.NewDHCPAllocator()
 	poolCache := make(map[string]string, 10)
 
 	return &Agent{
-		ctx: ctx,
-
 		dryRun:  options.DryRun,
 		nic:     options.Nic,
 		poolRef: options.IPPoolRef,
 
 		DHCPAllocator: dhcpAllocator,
 		ippoolEventHandler: ippool.NewEventHandler(
-			ctx,
 			options.KubeConfigPath,
 			options.KubeContext,
 			nil,
@@ -51,35 +46,35 @@ func NewAgent(ctx context.Context, options *config.AgentOptions) *Agent {
 	}
 }
 
-func (a *Agent) Run() error {
+func (a *Agent) Run(ctx context.Context) error {
 	logrus.Infof("monitor ippool %s", a.poolRef.String())
 
-	eg, egctx := errgroup.WithContext(a.ctx)
+	eg, egctx := errgroup.WithContext(ctx)
 
 	eg.Go(func() error {
-		select {
-		case <-egctx.Done():
-			return nil
-		default:
-			return a.DHCPAllocator.Run(a.nic, a.dryRun)
+		if a.dryRun {
+			return a.DHCPAllocator.DryRun(egctx, a.nic)
 		}
+		return a.DHCPAllocator.Run(egctx, a.nic)
 	})
 
 	eg.Go(func() error {
-		select {
-		case <-egctx.Done():
-			return nil
-		default:
-			// initialize the ippoolEventListener handler
-			if err := a.ippoolEventHandler.Init(); err != nil {
-				logrus.Fatal(err)
-			}
-			return a.ippoolEventHandler.EventListener()
+		if err := a.ippoolEventHandler.Init(); err != nil {
+			return err
 		}
+		a.ippoolEventHandler.EventListener(egctx)
+		return nil
 	})
 
+	errCh := dhcp.Cleanup(egctx, a.DHCPAllocator, a.nic)
+
 	if err := eg.Wait(); err != nil {
-		logrus.Fatal(err)
+		return err
+	}
+
+	// Return cleanup error message if any
+	if err := <-errCh; err != nil {
+		return err
 	}
 
 	return nil

pkg/agent/ippool/controller.go

@@ -15,6 +15,7 @@ import (
 )
 
 type Controller struct {
+	stopCh   chan struct{}
 	indexer  cache.Indexer
 	queue    workqueue.RateLimitingInterface
 	informer cache.Controller
@@ -33,6 +34,7 @@ func NewController(
 	poolCache map[string]string,
 ) *Controller {
 	return &Controller{
+		stopCh:        make(chan struct{}),
 		informer:      informer,
 		indexer:       indexer,
 		queue:         queue,
@@ -59,29 +61,29 @@ func (c *Controller) processNextItem() bool {
 func (c *Controller) sync(event Event) (err error) {
 	obj, exists, err := c.indexer.GetByKey(event.key)
 	if err != nil {
-		logrus.Errorf("(ippool.sync) fetching object with key %s from store failed with %v", event.key, err)
+		logrus.Errorf("(controller.sync) fetching object with key %s from store failed with %v", event.key, err)
 		return
 	}
 
 	if !exists && event.action != DELETE {
-		logrus.Infof("(ippool.sync) IPPool %s does not exist anymore", event.key)
+		logrus.Infof("(controller.sync) IPPool %s does not exist anymore", event.key)
 		return
 	}
 
 	if event.poolName != c.poolRef.Name {
-		logrus.Debugf("(ippool.sync) IPPool %s is not our target", event.key)
+		logrus.Debugf("(controller.sync) IPPool %s is not our target", event.key)
 		return
 	}
 
 	switch event.action {
 	case UPDATE:
 		ipPool, ok := obj.(*networkv1.IPPool)
 		if !ok {
-			logrus.Error("(ippool.sync) failed to assert obj during UPDATE")
+			logrus.Error("(controller.sync) failed to assert obj during UPDATE")
 		}
-		logrus.Infof("(ippool.sync) UPDATE %s/%s", ipPool.Namespace, ipPool.Name)
+		logrus.Infof("(controller.sync) UPDATE %s/%s", ipPool.Namespace, ipPool.Name)
 		if err := c.Update(ipPool); err != nil {
-			logrus.Errorf("(ippool.sync) failed to update DHCP lease store: %s", err.Error())
+			logrus.Errorf("(controller.sync) failed to update DHCP lease store: %s", err.Error())
 		}
 	}
 
@@ -96,7 +98,7 @@ func (c *Controller) handleErr(err error, key interface{}) {
 	}
 
 	if c.queue.NumRequeues(key) < 5 {
-		logrus.Errorf("(ippool.handleErr) syncing IPPool %v: %v", key, err)
+		logrus.Errorf("(controller.handleErr) syncing IPPool %v: %v", key, err)
 
 		c.queue.AddRateLimited(key)
 
@@ -105,31 +107,37 @@ func (c *Controller) handleErr(err error, key interface{}) {
 
 	c.queue.Forget(key)
 
-	logrus.Errorf("(ippool.handleErr) dropping IPPool %q out of the queue: %v", key, err)
+	logrus.Errorf("(controller.handleErr) dropping IPPool %q out of the queue: %v", key, err)
 }
 
-func (c *Controller) Run(workers int, stopCh chan struct{}) {
+func (c *Controller) Run(workers int) {
 	defer runtime.HandleCrash()
 
 	defer c.queue.ShutDown()
-	logrus.Infof("(ippool.Run) starting IPPool controller")
+	logrus.Info("(controller.Run) starting IPPool controller")
 
-	go c.informer.Run(stopCh)
-	if !cache.WaitForCacheSync(stopCh, c.informer.HasSynced) {
-		logrus.Errorf("(ippool.Run) timed out waiting for caches to sync")
+	go c.informer.Run(c.stopCh)
+	if !cache.WaitForCacheSync(c.stopCh, c.informer.HasSynced) {
+		logrus.Errorf("(controller.Run) timed out waiting for caches to sync")
 
 		return
 	}
 
 	for i := 0; i < workers; i++ {
-		go wait.Until(c.runWorker, time.Second, stopCh)
+		go wait.Until(c.runWorker, time.Second, c.stopCh)
 	}
 
-	<-stopCh
-	logrus.Infof("(ippool.Run) stopping IPPool controller")
+	<-c.stopCh
+
+	logrus.Info("(controller.Run) IPPool controller terminated")
 }
 
 func (c *Controller) runWorker() {
 	for c.processNextItem() {
 	}
 }
+
+func (c *Controller) Stop() {
+	logrus.Info("(controller.Stop) stopping IPPool controller")
+	close(c.stopCh)
+}

pkg/agent/ippool/event.go

@@ -25,7 +25,6 @@ const (
 )
 
 type EventHandler struct {
-	ctx            context.Context
 	kubeConfig     string
 	kubeContext    string
 	kubeRestConfig *rest.Config
@@ -44,7 +43,6 @@ type Event struct {
 }
 
 func NewEventHandler(
-	ctx context.Context,
 	kubeConfig string,
 	kubeContext string,
 	kubeRestConfig *rest.Config,
@@ -53,7 +51,6 @@ func NewEventHandler(
 	poolCache map[string]string,
 ) *EventHandler {
 	return &EventHandler{
-		ctx:            ctx,
 		kubeConfig:     kubeConfig,
 		kubeContext:    kubeContext,
 		kubeRestConfig: kubeRestConfig,
@@ -93,8 +90,8 @@ func (e *EventHandler) getKubeConfig() (config *rest.Config, err error) {
 	).ClientConfig()
 }
 
-func (e *EventHandler) EventListener() (err error) {
-	logrus.Infof("(ippool.EventListener) starting IPPool event listener")
+func (e *EventHandler) EventListener(ctx context.Context) {
+	logrus.Info("(eventhandler.EventListener) starting IPPool event listener")
 
 	// TODO: could be more specific on what namespaces we want to watch and what fields we need
 	watcher := cache.NewListWatchFromClient(e.k8sClientset.NetworkV1alpha1().RESTClient(), "ippools", e.poolRef.Namespace, fields.Everything())
@@ -116,9 +113,11 @@ func (e *EventHandler) EventListener() (err error) {
 	}, cache.Indexers{})
 
 	controller := NewController(queue, indexer, informer, e.poolRef, e.dhcpAllocator, e.poolCache)
-	stop := make(chan struct{})
-	defer close(stop)
-	go controller.Run(1, stop)
 
-	select {}
+	go controller.Run(1)
+
+	<-ctx.Done()
+	controller.Stop()
+
+	logrus.Info("(eventhandler.Run) IPPool event listener terminated")
 }