Increase robustness for the Managed DHCP add-on under disastrous events

.github/workflows/codeql.yml

@@ -13,9 +13,9 @@ name: "CodeQL"
 
 on:
   push:
-    branches: [ "main" ]
+    branches: [ "main", "v*" ]
   pull_request:
-    branches: [ "main" ]
+    branches: [ "main", "v*" ]
   schedule:
     - cron: '25 5 * * 6'
 

.github/workflows/main.yml

@@ -4,6 +4,7 @@ on:
   push:
     branches:
     - main
+    - v*
 
 jobs:
   build-main:

.github/workflows/scan.yml

@@ -3,12 +3,13 @@ on:
   push:
     branches:
     - main
+    - v*
   pull_request:
 
 jobs:
   build:
     name: Build
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-24.04
     steps:
     - name: Checkout code
       uses: actions/checkout@v4

cmd/agent/root.go

@@ -22,6 +22,7 @@ var (
 	dryRun             bool
 	nic                string
 	enableCacheDumpAPI bool
+	noLeaderElection   bool
 	kubeConfigPath     string
 	kubeContext        string
 	ippoolRef          string
@@ -74,6 +75,7 @@ func init() {
 	rootCmd.Flags().StringVar(&kubeContext, "kubecontext", os.Getenv("KUBECONTEXT"), "Context name")
 	rootCmd.Flags().BoolVar(&dryRun, "dry-run", false, "Run vm-dhcp-agent without starting the DHCP server")
 	rootCmd.Flags().BoolVar(&enableCacheDumpAPI, "enable-cache-dump-api", false, "Enable cache dump APIs")
+	rootCmd.Flags().BoolVar(&noLeaderElection, "no-leader-election", false, "Run vm-dhcp-agent with leader election disabled")
 	rootCmd.Flags().StringVar(&ippoolRef, "ippool-ref", os.Getenv("IPPOOL_REF"), "The IPPool object the agent should sync with")
 	rootCmd.Flags().StringVar(&nic, "nic", agent.DefaultNetworkInterface, "The network interface the embedded DHCP server listens on")
 }

cmd/agent/run.go

@@ -1,12 +1,17 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"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/rest"
+	"k8s.io/client-go/tools/clientcmd"
 
 	"github.com/harvester/vm-dhcp-controller/pkg/agent"
 	"github.com/harvester/vm-dhcp-controller/pkg/config"
@@ -18,37 +23,71 @@ func run(options *config.AgentOptions) error {
 
 	ctx := signals.SetupSignalContext()
 
-	agent := agent.NewAgent(options)
+	cfg, err := buildRestConfig(options.KubeConfigPath, options.KubeContext)
+	if err != nil {
+		return err
+	}
 
-	httpServerOptions := config.HTTPServerOptions{
-		DebugMode:     enableCacheDumpAPI,
-		DHCPAllocator: agent.DHCPAllocator,
+	client, err := kubernetes.NewForConfig(cfg)
+	if err != nil {
+		return err
 	}
-	s := server.NewHTTPServer(&httpServerOptions)
-	s.RegisterAgentHandlers()
 
-	eg, egctx := errgroup.WithContext(ctx)
+	callback := func(ctx context.Context) {
+		agent := agent.NewAgent(options)
 
-	eg.Go(func() error {
-		return s.Run()
-	})
+		httpServerOptions := config.HTTPServerOptions{
+			DebugMode:     enableCacheDumpAPI,
+			DHCPAllocator: agent.DHCPAllocator,
+		}
+		s := server.NewHTTPServer(&httpServerOptions)
+		s.RegisterAgentHandlers()
 
-	eg.Go(func() error {
-		return agent.Run(egctx)
-	})
+		eg, egctx := errgroup.WithContext(ctx)
 
-	errCh := server.Cleanup(egctx, s)
+		eg.Go(func() error {
+			return s.Run()
+		})
 
-	if err := eg.Wait(); err != nil && !errors.Is(err, http.ErrServerClosed) {
-		return err
+		eg.Go(func() error {
+			return agent.Run(egctx)
+		})
+
+		errCh := server.Cleanup(egctx, s)
+
+		if err := eg.Wait(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			logrus.Errorf("agent runtime error: %v", err)
+		}
+
+		if err := <-errCh; err != nil {
+			logrus.Errorf("cleanup error: %v", err)
+		}
 	}
 
-	// Return cleanup error message if any
-	if err := <-errCh; err != nil {
-		return err
+	if noLeaderElection {
+		callback(ctx)
+		<-ctx.Done()
+	} else {
+		leader.RunOrDie(ctx, "kube-system", "vm-dhcp-agent-"+options.IPPoolRef.Name, client, callback)
 	}
 
 	logrus.Info("finished clean")
 
 	return nil
 }
+
+func buildRestConfig(kubeconfig, kubecontext string) (*rest.Config, error) {
+	if kubeconfig == "" {
+		if cfg, err := rest.InClusterConfig(); err == nil {
+			return cfg, nil
+		}
+	}
+
+	loadingRules := &clientcmd.ClientConfigLoadingRules{ExplicitPath: kubeconfig}
+	overrides := &clientcmd.ConfigOverrides{}
+	if kubecontext != "" {
+		overrides.CurrentContext = kubecontext
+	}
+
+	return clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, overrides).ClientConfig()
+}

cmd/controller/root.go

@@ -47,16 +47,10 @@ var rootCmd = &cobra.Command{
 		}
 	},
 	Run: func(cmd *cobra.Command, args []string) {
-		var image *config.Image
-		imageTokens := strings.Split(agentImage, ":")
-		if len(imageTokens) == 2 {
-			image = config.NewImage(imageTokens[0], imageTokens[1])
-		} else {
-			fmt.Fprintf(os.Stderr, "Error parse agent image name\n")
-			if err := cmd.Help(); err != nil {
-				os.Exit(1)
-			}
-			os.Exit(0)
+		image, err := parseImageNameAndTag(agentImage)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+			os.Exit(1)
 		}
 
 		options := &config.ControllerOptions{
@@ -96,3 +90,26 @@ func init() {
 func execute() {
 	cobra.CheckErr(rootCmd.Execute())
 }
+
+func parseImageNameAndTag(image string) (*config.Image, error) {
+	idx := strings.LastIndex(image, ":")
+
+	if idx == -1 {
+		return config.NewImage(image, "latest"), nil
+	}
+
+	// If the last colon is immediately followed by the end of the string, it's invalid (no tag).
+	if idx == len(image)-1 {
+		return nil, fmt.Errorf("invalid image name: colon without tag")
+	}
+
+	if strings.Count(image, ":") > 2 {
+		return nil, fmt.Errorf("invalid image name: multiple colons found")
+	}
+
+	if idx <= strings.LastIndex(image, "/") {
+		return config.NewImage(image, "latest"), nil
+	}
+
+	return config.NewImage(image[:idx], image[idx+1:]), nil
+}

cmd/controller/root_test.go

@@ -0,0 +1,78 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/harvester/vm-dhcp-controller/pkg/config"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestParseImageNameAndTag(t *testing.T) {
+	tests := []struct {
+		name     string
+		image    string
+		expected *config.Image
+		err      bool
+	}{
+		{
+			name:  "valid image with registry and tag",
+			image: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller:v0.3.3",
+			expected: &config.Image{
+				Repository: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller",
+				Tag:        "v0.3.3",
+			},
+			err: false,
+		},
+		{
+			name:  "valid image with only image and tag",
+			image: "rancher/harvester-vm-dhcp-controller:v0.3.3",
+			expected: &config.Image{
+				Repository: "rancher/harvester-vm-dhcp-controller",
+				Tag:        "v0.3.3",
+			},
+			err: false,
+		},
+		{
+			name:  "valid image without tag",
+			image: "rancher/harvester-vm-dhcp-controller",
+			expected: &config.Image{
+				Repository: "rancher/harvester-vm-dhcp-controller",
+				Tag:        "latest",
+			},
+			err: false,
+		},
+		{
+			name:  "valid image with port but no tag",
+			image: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller",
+			expected: &config.Image{
+				Repository: "myregistry.local:5000/rancher/harvester-vm-dhcp-controller",
+				Tag:        "latest",
+			},
+			err: false,
+		},
+		{
+			name:     "invalid image with colon but no tag",
+			image:    "myregistry.local:5000/rancher/harvester-vm-dhcp-controller:",
+			expected: nil,
+			err:      true,
+		},
+		{
+			name:     "invalid image with multiple colons",
+			image:    "myregistry.local:5000/rancher/harvester-vm-dhcp-controller:v0.3.3:latest",
+			expected: nil,
+			err:      true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := parseImageNameAndTag(tt.image)
+			if tt.err {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expected, result)
+			}
+		})
+	}
+}

cmd/controller/run.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"log"
 	"net/http"
+	"time"
 
 	"github.com/rancher/wrangler/pkg/leader"
 	"github.com/rancher/wrangler/pkg/signals"
@@ -16,6 +17,7 @@ import (
 	"github.com/harvester/vm-dhcp-controller/pkg/config"
 	"github.com/harvester/vm-dhcp-controller/pkg/controller"
 	"github.com/harvester/vm-dhcp-controller/pkg/server"
+	"github.com/harvester/vm-dhcp-controller/pkg/util"
 )
 
 var (
@@ -46,6 +48,8 @@ func run(options *config.ControllerOptions) error {
 		logrus.Fatalf("Error building controllers: %s", err.Error())
 	}
 
+	go util.CleanupTerminatingPods(ctx, client, options.AgentNamespace, "controller", time.Minute)
+
 	callback := func(ctx context.Context) {
 		if err := management.Register(ctx, cfg, controller.RegisterFuncList); err != nil {
 			panic(err)

cmd/webhook/run.go

@@ -2,11 +2,14 @@ package main
 
 import (
 	"context"
+	"time"
 
+	"github.com/harvester/vm-dhcp-controller/pkg/util"
 	"github.com/harvester/webhook/pkg/config"
 	"github.com/harvester/webhook/pkg/server"
 	"github.com/rancher/wrangler/pkg/start"
 	"github.com/sirupsen/logrus"
+	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
 
 	ctlcore "github.com/harvester/vm-dhcp-controller/pkg/generated/controllers/core"
@@ -70,6 +73,13 @@ func run(ctx context.Context, cfg *rest.Config, options *config.Options) error {
 		return err
 	}
 
+	client, err := kubernetes.NewForConfig(cfg)
+	if err != nil {
+		return err
+	}
+
+	go util.CleanupTerminatingPods(ctx, client, options.Namespace, "webhook", time.Minute)
+
 	webhookServer := server.NewWebhookServer(ctx, cfg, name, options)
 
 	if err := webhookServer.RegisterValidators(

pkg/agent/agent.go

@@ -41,6 +41,7 @@ func NewAgent(options *config.AgentOptions) *Agent {
 			options.IPPoolRef,
 			dhcpAllocator,
 			poolCache,
+			options.Nic,
 		),
 		poolCache: poolCache,
 	}

pkg/agent/ippool/controller.go

@@ -23,6 +23,7 @@ type Controller struct {
 	poolRef       types.NamespacedName
 	dhcpAllocator *dhcp.DHCPAllocator
 	poolCache     map[string]string
+	nic           string
 }
 
 func NewController(
@@ -32,6 +33,7 @@ func NewController(
 	poolRef types.NamespacedName,
 	dhcpAllocator *dhcp.DHCPAllocator,
 	poolCache map[string]string,
+	nic string,
 ) *Controller {
 	return &Controller{
 		stopCh:        make(chan struct{}),
@@ -41,6 +43,7 @@ func NewController(
 		poolRef:       poolRef,
 		dhcpAllocator: dhcpAllocator,
 		poolCache:     poolCache,
+		nic:           nic,
 	}
 }