feat: helm chart value + conditional POD_UID env var for seamless upgrades

Author: alimaazamat

cmd/dra-example-kubeletplugin/driver.go

@@ -49,15 +49,22 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
 	}
 	driver.state = state
 
-	helper, err := kubeletplugin.Start(
-		ctx,
-		driver,
+	kubeletPluginOptions := []kubeletplugin.Option{
 		kubeletplugin.KubeClient(config.coreclient),
 		kubeletplugin.NodeName(config.flags.nodeName),
 		kubeletplugin.DriverName(consts.DriverName),
 		kubeletplugin.RegistrarDirectoryPath(config.flags.kubeletRegistrarDirectoryPath),
 		kubeletplugin.PluginDataDirectoryPath(config.DriverPluginPath()),
-	)
+	}
+
+	// Enable seamless upgrades when POD_UID is available (which gets set by Helm when seamlessUpgrades.enabled=true)
+	if config.flags.podUID != "" {
+		kubeletPluginOptions = append(kubeletPluginOptions,
+			kubeletplugin.RollingUpdate(types.UID(config.flags.podUID)),
+		)
+	}
+
+	helper, err := kubeletplugin.Start(ctx, driver, kubeletPluginOptions...)
 	if err != nil {
 		return nil, err
 	}

cmd/dra-example-kubeletplugin/health.go

@@ -21,11 +21,8 @@ import (
 	"fmt"
 	"net"
 	"net/url"
-	"os"
 	"path"
-	"path/filepath"
 	"strconv"
-	"strings"
 	"sync"
 
 	"google.golang.org/grpc"
@@ -66,18 +63,11 @@ func startHealthcheck(ctx context.Context, config *Config) (*healthcheck, error)
 
 	regSockPath := (&url.URL{
 		Scheme: "unix",
-		// Support both legacy and seamless upgrade socket naming conventions
 		Path: func() string {
-			socketPath, err := findSocketPath(
-				config.flags.kubeletRegistrarDirectoryPath,
-				consts.DriverName,
-				"-reg.sock",
-			)
-			if err != nil {
-				// Fallback to legacy path name
-				return path.Join(config.flags.kubeletRegistrarDirectoryPath, consts.DriverName+"-reg.sock")
+			if config.flags.podUID != "" {
+				return path.Join(config.flags.kubeletRegistrarDirectoryPath, consts.DriverName+"-"+config.flags.podUID+"-reg.sock")
 			}
-			return socketPath
+			return path.Join(config.flags.kubeletRegistrarDirectoryPath, consts.DriverName+"-reg.sock")
 		}(),
 	}).String()
 	log.Info("connecting to registration socket", "path", regSockPath)
@@ -92,16 +82,10 @@ func startHealthcheck(ctx context.Context, config *Config) (*healthcheck, error)
 	draSockPath := (&url.URL{
 		Scheme: "unix",
 		Path: func() string {
-			socketPath, err := findSocketPath(
-				config.DriverPluginPath(),
-				"dra",
-				".sock",
-			)
-			if err != nil {
-				// Fallback to legacy path name
-				return path.Join(config.DriverPluginPath(), "dra.sock")
+			if config.flags.podUID != "" {
+				return path.Join(config.DriverPluginPath(), "dra-"+config.flags.podUID+".sock")
 			}
-			return socketPath
+			return path.Join(config.DriverPluginPath(), "dra.sock")
 		}(),
 	}).String()
 	log.Info("connecting to DRA socket", "path", draSockPath)
@@ -171,39 +155,3 @@ func (h *healthcheck) Check(ctx context.Context, req *grpc_health_v1.HealthCheck
 	status.Status = grpc_health_v1.HealthCheckResponse_SERVING
 	return status, nil
 }
-
-// Finds driver's socket paths whether its legacy (fixed filename) or seamless upgrade (UID-based filename) formats.
-func findSocketPath(dir, baseName, suffix string) (string, error) {
-	// First try the legacy path name: {baseName}{suffix}
-	legacyPath := filepath.Join(dir, baseName+suffix)
-	if _, err := os.Stat(legacyPath); err == nil {
-		return legacyPath, nil
-	}
-
-	// Then try the seamless upgrade format: {baseName}-{uid}{suffix}
-	entries, err := os.ReadDir(dir)
-	if err != nil {
-		return "", fmt.Errorf("failed to read directory %s: %w", dir, err)
-	}
-
-	for _, entry := range entries {
-		if entry.IsDir() {
-			continue
-		}
-		name := entry.Name()
-
-		// Look for files matching pattern: {baseName}-{uid}{suffix}
-		if strings.HasPrefix(name, baseName+"-") && strings.HasSuffix(name, suffix) {
-			// Verify it's not the legacy format (which would be {baseName}{suffix})
-			if name != baseName+suffix {
-				socketPath := filepath.Join(dir, name)
-				klog.Info("Found seamless upgrade socket", "path", socketPath, "uid", strings.TrimPrefix(strings.TrimSuffix(name, suffix), baseName+"-"))
-				return socketPath, nil
-			}
-		}
-	}
-
-	// If neither path name types are found, return the legacy path for error reporting
-	return legacyPath, fmt.Errorf("socket file not found in dir %s (tried legacy path name %s and seamless upgrade with uid path name %s-*%s)",
-		dir, baseName+suffix, baseName, suffix)
-}

cmd/dra-example-kubeletplugin/main.go

@@ -49,6 +49,7 @@ type Flags struct {
 	kubeletRegistrarDirectoryPath string
 	kubeletPluginsDirectoryPath   string
 	healthcheckPort               int
+	podUID                        string
 }
 
 type Config struct {
@@ -114,6 +115,12 @@ func newApp() *cli.App {
 			Destination: &flags.healthcheckPort,
 			EnvVars:     []string{"HEALTHCHECK_PORT"},
 		},
+		&cli.StringFlag{
+			Name:        "pod-uid",
+			Usage:       "UID of the pod (used for seamless upgrades to create unique socket names).",
+			Destination: &flags.podUID,
+			EnvVars:     []string{"POD_UID"},
+		},
 	}
 	cliFlags = append(cliFlags, flags.kubeClientConfig.Flags()...)
 	cliFlags = append(cliFlags, flags.loggingConfig.Flags()...)

deployments/helm/dra-example-driver/templates/kubeletplugin.yaml

@@ -12,10 +12,18 @@ spec:
     matchLabels:
       {{- include "dra-example-driver.selectorLabels" . | nindent 6 }}
       app.kubernetes.io/component: kubeletplugin
+  {{- if .Values.kubeletPlugin.seamlessUpgrades.enabled }}
+  updateStrategy:
+    type: RollingUpdate
+    rollingUpdate:
+      maxSurge: 1
+      maxUnavailable: 0
+  {{- else }}
   {{- with .Values.kubeletPlugin.updateStrategy }}
   updateStrategy:
     {{- toYaml . | nindent 4 }}
   {{- end }}
+  {{- end }}
   template:
     metadata:
       {{- with .Values.kubeletPlugin.podAnnotations }}
@@ -80,6 +88,12 @@ spec:
         - name: HEALTHCHECK_PORT
           value: {{ .Values.kubeletPlugin.containers.plugin.healthcheckPort | quote }}
         {{- end }}
+        {{- if .Values.kubeletPlugin.seamlessUpgrades.enabled }}
+        - name: POD_UID
+          valueFrom:
+            fieldRef:
+              fieldPath: metadata.uid
+        {{- end }}
         volumeMounts:
         - name: plugins-registry
           mountPath: {{ .Values.kubeletPlugin.kubeletRegistrarDirectoryPath | quote }}

deployments/helm/dra-example-driver/values.yaml

@@ -46,6 +46,8 @@ controller:
 
 kubeletPlugin:
   priorityClassName: "system-node-critical"
+  seamlessUpgrades:
+    enabled: false
   updateStrategy:
     type: RollingUpdate
   podAnnotations: {}