Refactor:Remove grpc wrapper

go.mod

@@ -25,6 +25,7 @@ require (
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.21.0
 	go.opentelemetry.io/otel/sdk v1.21.0
 	golang.org/x/net v0.19.0
+	golang.org/x/sync v0.5.0
 	google.golang.org/grpc v1.60.1
 	google.golang.org/protobuf v1.31.0
 	k8s.io/api v0.29.0
@@ -131,7 +132,6 @@ require (
 	golang.org/x/exp v0.0.0-20231214170342-aacd6d4b4611 // indirect
 	golang.org/x/mod v0.14.0 // indirect
 	golang.org/x/oauth2 v0.15.0 // indirect
-	golang.org/x/sync v0.5.0 // indirect
 	golang.org/x/sys v0.15.0 // indirect
 	golang.org/x/term v0.15.0 // indirect
 	golang.org/x/text v0.14.0 // indirect

pkg/azuredisk/azuredisk.go

@@ -18,6 +18,7 @@ package azuredisk
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"reflect"
 	"strconv"
@@ -26,7 +27,9 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2022-08-01/compute"
 	"github.com/container-storage-interface/spec/lib/go/csi"
+	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -78,6 +81,9 @@ type DriverOptions struct {
 	EnableOtelTracing            bool
 	WaitForSnapshotReady         bool
 	CheckDiskLUNCollision        bool
+	Kubeconfig                   string
+	Endpoint                     string
+	DisableAVSetNodes            bool
 }
 
 // CSIDriver defines the interface for a CSI driver.
@@ -86,7 +92,7 @@ type CSIDriver interface {
 	csi.NodeServer
 	csi.IdentityServer
 
-	Run(endpoint, kubeconfig string, disableAVSetNodes, testMode bool)
+	Run(ctx context.Context) error
 }
 
 type hostUtil interface {
@@ -127,6 +133,8 @@ type DriverCore struct {
 	enableOtelTracing            bool
 	shouldWaitForSnapshotReady   bool
 	checkDiskLUNCollision        bool
+	endpoint                     string
+	disableAVSetNodes            bool
 }
 
 // Driver is the v1 implementation of the Azure Disk CSI Driver.
@@ -169,6 +177,9 @@ func newDriverV1(options *DriverOptions) *Driver {
 	driver.enableOtelTracing = options.EnableOtelTracing
 	driver.shouldWaitForSnapshotReady = options.WaitForSnapshotReady
 	driver.checkDiskLUNCollision = options.CheckDiskLUNCollision
+	driver.endpoint = options.Endpoint
+	driver.disableAVSetNodes = options.DisableAVSetNodes
+	driver.kubeconfig = options.Kubeconfig
 	driver.volumeLocks = volumehelper.NewVolumeLocks()
 	driver.ioHandler = azureutils.NewOSIOHandler()
 	driver.hostUtil = hostutil.NewHostUtil()
@@ -186,7 +197,7 @@ func newDriverV1(options *DriverOptions) *Driver {
 }
 
 // Run driver initialization
-func (d *Driver) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMock bool) {
+func (d *Driver) Run(ctx context.Context) error {
 	versionMeta, err := GetVersionYAML(d.Name)
 	if err != nil {
 		klog.Fatalf("%v", err)
@@ -196,13 +207,12 @@ func (d *Driver) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMock
 	userAgent := GetUserAgent(d.Name, d.customUserAgent, d.userAgentSuffix)
 	klog.V(2).Infof("driver userAgent: %s", userAgent)
 
-	cloud, err := azureutils.GetCloudProvider(context.Background(), kubeconfig, d.cloudConfigSecretName, d.cloudConfigSecretNamespace,
+	cloud, err := azureutils.GetCloudProvider(context.Background(), d.kubeconfig, d.cloudConfigSecretName, d.cloudConfigSecretNamespace,
 		userAgent, d.allowEmptyCloudConfig, d.enableTrafficManager, d.trafficManagerPort)
 	if err != nil {
 		klog.Fatalf("failed to get Azure Cloud Provider, error: %v", err)
 	}
 	d.cloud = cloud
-	d.kubeconfig = kubeconfig
 
 	if d.cloud != nil {
 		if d.vmType != "" {
@@ -221,7 +231,7 @@ func (d *Driver) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMock
 				d.cloud.DisableAvailabilitySetNodes = false
 			}
 
-			if d.cloud.VMType == azurecloudconsts.VMTypeVMSS && !d.cloud.DisableAvailabilitySetNodes && disableAVSetNodes {
+			if d.cloud.VMType == azurecloudconsts.VMTypeVMSS && !d.cloud.DisableAvailabilitySetNodes && d.disableAVSetNodes {
 				klog.V(2).Infof("DisableAvailabilitySetNodes for controller since current VMType is vmss")
 				d.cloud.DisableAvailabilitySetNodes = true
 			}
@@ -283,11 +293,39 @@ func (d *Driver) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMock
 		csi.NodeServiceCapability_RPC_GET_VOLUME_STATS,
 		csi.NodeServiceCapability_RPC_SINGLE_NODE_MULTI_WRITER,
 	})
+	grpcInterceptor := grpc.UnaryInterceptor(csicommon.LogGRPC)
+	if d.enableOtelTracing {
+		grpcInterceptor = grpc.ChainUnaryInterceptor(csicommon.LogGRPC, otelgrpc.UnaryServerInterceptor())
+	}
+	opts := []grpc.ServerOption{
+		grpcInterceptor,
+	}
+
+	if d.enableOtelTracing {
+		opts = append(opts, grpc.StatsHandler(otelgrpc.NewServerHandler()))
+	}
+
+	s := grpc.NewServer(opts...)
+	csi.RegisterIdentityServer(s, d)
+	csi.RegisterControllerServer(s, d)
+	csi.RegisterNodeServer(s, d)
 
-	s := csicommon.NewNonBlockingGRPCServer()
+	go func() {
+		//graceful shutdown
+		<-ctx.Done()
+		s.GracefulStop()
+	}()
 	// Driver d act as IdentityServer, ControllerServer and NodeServer
-	s.Start(endpoint, d, d, d, testingMock, d.enableOtelTracing)
-	s.Wait()
+	listener, err := csicommon.Listen(ctx, d.endpoint)
+	if err != nil {
+		klog.Fatalf("failed to listen to endpoint, error: %v", err)
+	}
+	err = s.Serve(listener)
+	if errors.Is(err, grpc.ErrServerStopped) {
+		klog.Infof("gRPC server stopped serving")
+		return nil
+	}
+	return err
 }
 
 func (d *Driver) isGetDiskThrottled() bool {

pkg/azuredisk/azuredisk_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/Azure/go-autorest/autorest/date"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc/status"
 	"k8s.io/apimachinery/pkg/types"
 	clientset "k8s.io/client-go/kubernetes"
@@ -107,7 +108,14 @@ func TestRun(t *testing.T) {
 				cntl := gomock.NewController(t)
 				defer cntl.Finish()
 				d, _ := NewFakeDriver(cntl)
-				d.Run("tcp://127.0.0.1:0", "", true, true)
+				ctx, cancelFn := context.WithCancel(context.Background())
+				var routines errgroup.Group
+				routines.Go(func() error { return d.Run(ctx) })
+				time.Sleep(time.Millisecond * 500)
+				cancelFn()
+				time.Sleep(time.Millisecond * 500)
+				err := routines.Wait()
+				assert.Nil(t, err)
 			},
 		},
 		{
@@ -116,7 +124,14 @@ func TestRun(t *testing.T) {
 				cntl := gomock.NewController(t)
 				defer cntl.Finish()
 				d, _ := NewFakeDriver(cntl)
-				d.Run("tcp://127.0.0.1:0", "", true, true)
+				ctx, cancelFn := context.WithCancel(context.Background())
+				var routines errgroup.Group
+				routines.Go(func() error { return d.Run(ctx) })
+				time.Sleep(time.Millisecond * 500)
+				cancelFn()
+				time.Sleep(time.Millisecond * 500)
+				err := routines.Wait()
+				assert.Nil(t, err)
 			},
 		},
 		{
@@ -139,7 +154,14 @@ func TestRun(t *testing.T) {
 				d, _ := NewFakeDriver(cntl)
 				d.setCloud(&azure.Cloud{})
 				d.setNodeID("")
-				d.Run("tcp://127.0.0.1:0", "", true, true)
+				ctx, cancelFn := context.WithCancel(context.Background())
+				var routines errgroup.Group
+				routines.Go(func() error { return d.Run(ctx) })
+				time.Sleep(time.Millisecond * 500)
+				cancelFn()
+				time.Sleep(time.Millisecond * 500)
+				err := routines.Wait()
+				assert.Nil(t, err)
 			},
 		},
 		{
@@ -165,8 +187,16 @@ func TestRun(t *testing.T) {
 					EnablePerfOptimization: true,
 					VMSSCacheTTLInSeconds:  10,
 					VMType:                 "vmss",
+					Endpoint:               "tcp://127.0.0.1:0",
 				})
-				d.Run("tcp://127.0.0.1:0", "", true, true)
+				ctx, cancelFn := context.WithCancel(context.Background())
+				var routines errgroup.Group
+				routines.Go(func() error { return d.Run(ctx) })
+				time.Sleep(time.Millisecond * 500)
+				cancelFn()
+				time.Sleep(time.Millisecond * 500)
+				err := routines.Wait()
+				assert.Nil(t, err)
 			},
 		},
 	}

pkg/azuredisk/azuredisk_v2.go

@@ -21,13 +21,16 @@ package azuredisk
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"reflect"
 
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2022-08-01/compute"
 	"github.com/container-storage-interface/spec/lib/go/csi"
+	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -78,13 +81,16 @@ func newDriverV2(options *DriverOptions) *DriverV2 {
 	driver.enableOtelTracing = options.EnableOtelTracing
 	driver.ioHandler = azureutils.NewOSIOHandler()
 	driver.hostUtil = hostutil.NewHostUtil()
+	driver.kubeconfig = options.Kubeconfig
+	driver.disableAVSetNodes = options.DisableAVSetNodes
+	driver.endpoint = options.Endpoint
 
 	topologyKey = fmt.Sprintf("topology.%s/zone", driver.Name)
 	return &driver
 }
 
 // Run driver initialization
-func (d *DriverV2) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMock bool) {
+func (d *DriverV2) Run(ctx context.Context) error {
 	versionMeta, err := GetVersionYAML(d.Name)
 	if err != nil {
 		klog.Fatalf("%v", err)
@@ -94,7 +100,7 @@ func (d *DriverV2) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMo
 	userAgent := GetUserAgent(d.Name, d.customUserAgent, d.userAgentSuffix)
 	klog.V(2).Infof("driver userAgent: %s", userAgent)
 
-	cloud, err := azureutils.GetCloudProvider(context.Background(), kubeconfig, d.cloudConfigSecretName, d.cloudConfigSecretNamespace,
+	cloud, err := azureutils.GetCloudProvider(context.Background(), d.kubeconfig, d.cloudConfigSecretName, d.cloudConfigSecretNamespace,
 		userAgent, d.allowEmptyCloudConfig, d.enableTrafficManager, d.trafficManagerPort)
 	if err != nil {
 		klog.Fatalf("failed to get Azure Cloud Provider, error: %v", err)
@@ -118,7 +124,7 @@ func (d *DriverV2) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMo
 				d.cloud.DisableAvailabilitySetNodes = false
 			}
 
-			if d.cloud.VMType == consts.VMTypeVMSS && !d.cloud.DisableAvailabilitySetNodes && disableAVSetNodes {
+			if d.cloud.VMType == consts.VMTypeVMSS && !d.cloud.DisableAvailabilitySetNodes && d.disableAVSetNodes {
 				klog.V(2).Infof("DisableAvailabilitySetNodes for controller since current VMType is vmss")
 				d.cloud.DisableAvailabilitySetNodes = true
 			}
@@ -165,11 +171,34 @@ func (d *DriverV2) Run(endpoint, kubeconfig string, disableAVSetNodes, testingMo
 		csi.NodeServiceCapability_RPC_GET_VOLUME_STATS,
 		csi.NodeServiceCapability_RPC_SINGLE_NODE_MULTI_WRITER,
 	})
-
-	s := csicommon.NewNonBlockingGRPCServer()
+	grpcInterceptor := grpc.UnaryInterceptor(csicommon.LogGRPC)
+	if d.enableOtelTracing {
+		grpcInterceptor = grpc.ChainUnaryInterceptor(csicommon.LogGRPC, otelgrpc.UnaryServerInterceptor())
+	}
+	opts := []grpc.ServerOption{
+		grpcInterceptor,
+	}
+	s := grpc.NewServer(opts...)
+	csi.RegisterIdentityServer(s, d)
+	csi.RegisterControllerServer(s, d)
+	csi.RegisterNodeServer(s, d)
+
+	go func() {
+		//graceful shutdown
+		<-ctx.Done()
+		s.GracefulStop()
+	}()
 	// Driver d act as IdentityServer, ControllerServer and NodeServer
-	s.Start(endpoint, d, d, d, testingMock, d.enableOtelTracing)
-	s.Wait()
+	listener, err := csicommon.Listen(ctx, d.endpoint)
+	if err != nil {
+		klog.Fatalf("failed to listen to endpoint, error: %v", err)
+	}
+	err = s.Serve(listener)
+	if errors.Is(err, grpc.ErrServerStopped) {
+		klog.Infof("gRPC server stopped serving")
+		return nil
+	}
+	return err
 }
 
 func (d *DriverV2) checkDiskExists(ctx context.Context, diskURI string) (*compute.Disk, error) {

pkg/azuredisk/fake_azuredisk.go

@@ -114,6 +114,9 @@ func newFakeDriverV1(ctrl *gomock.Controller) (*fakeDriverV1, error) {
 	driver.useCSIProxyGAInterface = true
 	driver.allowEmptyCloudConfig = true
 	driver.shouldWaitForSnapshotReady = true
+	driver.endpoint = "tcp://127.0.0.1:0"
+	driver.disableAVSetNodes = true
+	driver.kubeconfig = ""
 
 	driver.cloud = azure.GetTestCloud(ctrl)
 	mounter, err := mounter.NewSafeMounter(driver.enableWindowsHostProcess, driver.useCSIProxyGAInterface)

pkg/azuredisk/fake_azuredisk_v2.go

@@ -64,6 +64,9 @@ func newFakeDriverV2(ctrl *gomock.Controller) (*fakeDriverV2, error) {
 	driver.hostUtil = azureutils.NewFakeHostUtil()
 	driver.useCSIProxyGAInterface = true
 	driver.allowEmptyCloudConfig = true
+	driver.endpoint = "tcp://127.0.0.1:0"
+	driver.disableAVSetNodes = true
+	driver.kubeconfig = ""
 
 	driver.cloud = azure.GetTestCloud(ctrl)
 	mounter, err := mounter.NewSafeMounter(driver.enableWindowsHostProcess, driver.useCSIProxyGAInterface)

pkg/azurediskplugin/main.go

@@ -135,13 +135,17 @@ func handle() {
 		EnableOtelTracing:            *enableOtelTracing,
 		WaitForSnapshotReady:         *waitForSnapshotReady,
 		CheckDiskLUNCollision:        *checkDiskLUNCollision,
+		Endpoint:                     *endpoint,
+		Kubeconfig:                   *kubeconfig,
+		DisableAVSetNodes:            *disableAVSetNodes,
 	}
 	driver := azuredisk.NewDriver(&driverOptions)
 	if driver == nil {
 		klog.Fatalln("Failed to initialize azuredisk CSI Driver")
 	}
-	testingMock := false
-	driver.Run(*endpoint, *kubeconfig, *disableAVSetNodes, testingMock)
+	if err := driver.Run(context.Background()); err != nil {
+		klog.Fatalf("Failed to run azuredisk CSI Driver: %v", err)
+	}
 }
 
 func exportMetrics() {