fix: terminate goroutines gracefully

cmd/agent/run.go

@@ -2,8 +2,13 @@ package main
 
 import (
 	"context"
+	"errors"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"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"
@@ -13,15 +18,58 @@ import (
 func run(ctx context.Context, options *config.AgentOptions) error {
 	logrus.Infof("Starting VM DHCP Agent: %s", name)
 
-	agent := agent.NewAgent(ctx, options)
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+
+	eg, egctx := errgroup.WithContext(ctx)
+
+	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.Go(func() error {
+		return s.Run()
+	})
+
+	eg.Go(func() error {
+		return agent.Run(egctx)
+	})
+
+	eg.Go(func() error {
+		<-egctx.Done()
+		return s.Stop(egctx)
+	})
+
+	eg.Go(func() error {
+		for {
+			select {
+			case sig := <-sigCh:
+				logrus.Infof("received signal: %s", sig)
+				cancel()
+			case <-egctx.Done():
+				return egctx.Err()
+			}
+		}
+	})
+
+	if err := eg.Wait(); err != nil {
+		if errors.Is(err, context.Canceled) {
+			logrus.Info("context canceled")
+			return nil
+		} else {
+			return err
+		}
+	}
+
+	logrus.Info("finished clean")
+
+	return nil
 }

cmd/controller/run.go

@@ -2,11 +2,13 @@ package main
 
 import (
 	"context"
+	"errors"
 	"log"
 
 	"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,12 +65,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 {
+		<-egctx.Done()
+		return s.Stop(egctx)
+	})
+
+	eg.Go(func() error {
+		if noLeaderElection {
+			callback(egctx)
+		} else {
+			leader.RunOrDie(egctx, "kube-system", "vm-dhcp-controllers", client, callback)
+		}
+		return nil
+	})
+
+	if err := eg.Wait(); err != nil {
+		if errors.Is(err, context.Canceled) {
+			logrus.Info("context canceled")
+			return nil
+		} else {
+			logrus.Fatalf("received error: %s", err)
+		}
 	}
 
 	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,36 +46,38 @@ 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)
+	stop := make(chan struct{})
+
+	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(stop)
+		return nil
 	})
 
-	if err := eg.Wait(); err != nil {
-		logrus.Fatal(err)
-	}
+	eg.Go(func() error {
+		<-egctx.Done()
+		return a.DHCPAllocator.Stop(a.nic)
+	})
+
+	eg.Go(func() error {
+		<-egctx.Done()
+		a.ippoolEventHandler.Stop(stop)
+		return nil
+	})
 
-	return nil
+	return eg.Wait()
 }

pkg/agent/ippool/controller.go

@@ -59,29 +59,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 +96,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,18 +105,18 @@ 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{}) {
 	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")
+		logrus.Errorf("(controller.Run) timed out waiting for caches to sync")
 
 		return
 	}
@@ -126,10 +126,16 @@ func (c *Controller) Run(workers int, stopCh chan struct{}) {
 	}
 
 	<-stopCh
-	logrus.Infof("(ippool.Run) stopping IPPool controller")
+
+	logrus.Info("(controller.Run) IPPool controller terminated")
 }
 
 func (c *Controller) runWorker() {
 	for c.processNextItem() {
 	}
 }
+
+func (c *Controller) Stop(stopCh chan struct{}) {
+	logrus.Info("(controller.Stop) stopping IPPool controller")
+	close(stopCh)
+}

pkg/agent/ippool/event.go

@@ -1,8 +1,6 @@
 package ippool
 
 import (
-	"context"
-
 	"github.com/sirupsen/logrus"
 	"k8s.io/apimachinery/pkg/fields"
 	"k8s.io/apimachinery/pkg/types"
@@ -25,7 +23,6 @@ const (
 )
 
 type EventHandler struct {
-	ctx            context.Context
 	kubeConfig     string
 	kubeContext    string
 	kubeRestConfig *rest.Config
@@ -44,7 +41,6 @@ type Event struct {
 }
 
 func NewEventHandler(
-	ctx context.Context,
 	kubeConfig string,
 	kubeContext string,
 	kubeRestConfig *rest.Config,
@@ -53,7 +49,6 @@ func NewEventHandler(
 	poolCache map[string]string,
 ) *EventHandler {
 	return &EventHandler{
-		ctx:            ctx,
 		kubeConfig:     kubeConfig,
 		kubeContext:    kubeContext,
 		kubeRestConfig: kubeRestConfig,
@@ -93,8 +88,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(stopCh chan struct{}) {
+	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())
@@ -117,8 +112,16 @@ func (e *EventHandler) EventListener() (err error) {
 
 	controller := NewController(queue, indexer, informer, e.poolRef, e.dhcpAllocator, e.poolCache)
 	stop := make(chan struct{})
-	defer close(stop)
+
 	go controller.Run(1, stop)
 
-	select {}
+	<-stopCh
+	controller.Stop(stop)
+
+	logrus.Info("(eventhandler.Run) IPPool event listener terminated")
+}
+
+func (e *EventHandler) Stop(stopCh chan struct{}) {
+	logrus.Info("(eventhandler.Stop) stopping IPPool event listener")
+	close(stopCh)
 }