From ca9f01d11e282c9949e54547982e501bb1db74a4 Mon Sep 17 00:00:00 2001
From: Sebastiaan van Stijn <github@gone.nl>
Date: Sat, 25 Jul 2026 18:24:38 +0200
Subject: [PATCH] agent/csi,manager/csi: roll back client initialization on
 failure

The CSI plugins publish their gRPC clients before initialization completes
because init() uses the client accessors during capability discovery. If
initialization fails, the plugin is left partially initialized and subsequent
calls reuse stale clients instead of retrying initialization.

Close the gRPC connection and reset the published clients when init() fails so
the plugin returns to its pre-connect state and a later attempt can retry
cleanly.

Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
---
 agent/csi/plugin/plugin.go | 29 ++++++++++++++++++-----------
 manager/csi/plugin.go      | 19 ++++++++++++++++---
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/agent/csi/plugin/plugin.go b/agent/csi/plugin/plugin.go
index a721ac25a8..4760edc3a5 100644
--- a/agent/csi/plugin/plugin.go
+++ b/agent/csi/plugin/plugin.go
@@ -2,6 +2,7 @@ package plugin
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"path/filepath"
 	"sync"
@@ -116,13 +117,17 @@ func (np *nodePlugin) connect(ctx context.Context) error {
 	}
 
 	np.cc = cc
-	// first, probe the plugin, to ensure that it exists and is ready to go
-	idc := csi.NewIdentityClient(cc)
-	np.idClient = idc
-
+	np.idClient = csi.NewIdentityClient(cc)
 	np.nodeClient = csi.NewNodeClient(cc)
 
-	return np.init(ctx)
+	if err := np.init(ctx); err != nil {
+		_ = cc.Close()
+		np.cc = nil
+		np.idClient = nil
+		np.nodeClient = nil
+		return err
+	}
+	return nil
 }
 
 func (np *nodePlugin) Client(ctx context.Context) (csi.NodeClient, error) {
@@ -135,6 +140,13 @@ func (np *nodePlugin) Client(ctx context.Context) (csi.NodeClient, error) {
 }
 
 func (np *nodePlugin) init(ctx context.Context) error {
+	if np.idClient == nil {
+		return errors.New("identity client is not initialized")
+	}
+	if np.nodeClient == nil {
+		return errors.New("node client is not initialized")
+	}
+	// first, probe the plugin, to ensure that it exists and is ready to go
 	probe, err := np.idClient.Probe(ctx, &csi.ProbeRequest{})
 	if err != nil {
 		return err
@@ -143,12 +155,7 @@ func (np *nodePlugin) init(ctx context.Context) error {
 		return status.Error(codes.FailedPrecondition, "Plugin is not Ready")
 	}
 
-	c, err := np.Client(ctx)
-	if err != nil {
-		return err
-	}
-
-	resp, err := c.NodeGetCapabilities(ctx, &csi.NodeGetCapabilitiesRequest{})
+	resp, err := np.nodeClient.NodeGetCapabilities(ctx, &csi.NodeGetCapabilitiesRequest{})
 	if err != nil {
 		// TODO(ameyag): handle
 		return err
diff --git a/manager/csi/plugin.go b/manager/csi/plugin.go
index 0aacdb33cd..510ee10203 100644
--- a/manager/csi/plugin.go
+++ b/manager/csi/plugin.go
@@ -101,20 +101,30 @@ func (p *plugin) connect(ctx context.Context) error {
 	p.cc = cc
 
 	// first, probe the plugin, to ensure that it exists and is ready to go
-	idc := csi.NewIdentityClient(cc)
-	p.idClient = idc
+	p.idClient = csi.NewIdentityClient(cc)
 
 	// controllerClient may not do anything if the plugin does not support
 	// the controller service, but it should not be an error to create it now
 	// anyway
 	p.controllerClient = csi.NewControllerClient(cc)
 
-	return p.init(ctx)
+	if err := p.init(ctx); err != nil {
+		_ = cc.Close()
+		p.cc = nil
+		p.idClient = nil
+		p.controllerClient = nil
+		return err
+	}
+
+	return nil
 }
 
 // init checks uses the identity service to check the properties of the plugin,
 // most importantly, its capabilities.
 func (p *plugin) init(ctx context.Context) error {
+	if p.idClient == nil {
+		return errors.New("identity client is not initialized")
+	}
 	probe, err := p.idClient.Probe(ctx, &csi.ProbeRequest{})
 	if err != nil {
 		return err
@@ -143,6 +153,9 @@ func (p *plugin) init(ctx context.Context) error {
 	}
 
 	if p.controller {
+		if p.controllerClient == nil {
+			return errors.New("controller client is not initialized")
+		}
 		cCapResp, err := p.controllerClient.ControllerGetCapabilities(
 			ctx, &csi.ControllerGetCapabilitiesRequest{},
 		)