Implementation with unit tests

Author: julianKatz

cmd/gce-pd-csi-driver/main.go

@@ -27,6 +27,8 @@ import (
 	"strings"
 	"time"
 
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
 	"k8s.io/klog/v2"
 	"k8s.io/utils/strings/slices"
 	"sigs.k8s.io/gcp-compute-persistent-disk-csi-driver/pkg/common"
@@ -94,6 +96,8 @@ var (
 
 	extraTagsStr = flag.String("extra-tags", "", "Extra tags to attach to each Compute Disk, Image, Snapshot created. It is a comma separated list of parent id, key and value like '<parent_id1>/<tag_key1>/<tag_value1>,...,<parent_idN>/<tag_keyN>/<tag_valueN>'. parent_id is the Organization or the Project ID or Project name where the tag key and the tag value resources exist. A maximum of 50 tags bindings is allowed for a resource. See https://cloud.google.com/resource-manager/docs/tags/tags-overview, https://cloud.google.com/resource-manager/docs/tags/tags-creating-and-managing for details")
 
+	diskTopology = flag.Bool("disk-topology", false, "If set to true, the driver will add a topology.gke.io/[disk-type] topology label to the Topologies returned in CreateVolumeResponse")
+
 	version string
 )
 
@@ -227,7 +231,10 @@ func handle() {
 		}
 		initialBackoffDuration := time.Duration(*errorBackoffInitialDurationMs) * time.Millisecond
 		maxBackoffDuration := time.Duration(*errorBackoffMaxDurationMs) * time.Millisecond
-		controllerServer = driver.NewControllerServer(gceDriver, cloudProvider, initialBackoffDuration, maxBackoffDuration, fallbackRequisiteZones, *enableStoragePoolsFlag, *enableDataCacheFlag, multiZoneVolumeHandleConfig, listVolumesConfig, provisionableDisksConfig, *enableHdHAFlag)
+		args := &driver.GCEControllerServerArgs{
+			EnableDiskTopology: *diskTopology,
+		}
+		controllerServer = driver.NewControllerServer(gceDriver, cloudProvider, initialBackoffDuration, maxBackoffDuration, fallbackRequisiteZones, *enableStoragePoolsFlag, *enableDataCacheFlag, multiZoneVolumeHandleConfig, listVolumesConfig, provisionableDisksConfig, *enableHdHAFlag, args)
 	} else if *cloudConfigFilePath != "" {
 		klog.Warningf("controller service is disabled but cloud config given - it has no effect")
 	}
@@ -239,6 +246,7 @@ func handle() {
 		if err != nil {
 			klog.Fatalf("Failed to get safe mounter: %v", err.Error())
 		}
+
 		deviceUtils := deviceutils.NewDeviceUtils()
 		statter := mountmanager.NewStatter(mounter)
 		meta, err := metadataservice.NewMetadataService()
@@ -249,13 +257,26 @@ func handle() {
 		if err != nil {
 			klog.Fatalf("Failed to get node info from API server: %v", err.Error())
 		}
-		nsArgs := driver.NodeServerArgs{
+
+		nsArgs := &driver.NodeServerArgs{
 			EnableDeviceInUseCheck:   *enableDeviceInUseCheck,
 			DeviceInUseTimeout:       *deviceInUseTimeout,
 			EnableDataCache:          *enableDataCacheFlag,
 			DataCacheEnabledNodePool: isDataCacheEnabledNodePool,
+			EnableDiskTopology:       *diskTopology,
 		}
+
+		if *diskTopology {
+			klog.V(2).Infof("Setting up kubeClient")
+			kubeClient, err := instantiateKubeClient()
+			if err != nil {
+				klog.Fatalf("Failed to instantiate Kubernetes client: %v", err)
+			}
+			nsArgs.KubeClient = kubeClient
+		}
+
 		nodeServer = driver.NewNodeServer(gceDriver, mounter, deviceUtils, meta, statter, nsArgs)
+
 		if *maxConcurrentFormatAndMount > 0 {
 			nodeServer = nodeServer.WithSerializedFormatAndMount(*formatAndMountTimeout, *maxConcurrentFormatAndMount)
 		}
@@ -292,6 +313,18 @@ func handle() {
 	gceDriver.Run(*endpoint, *grpcLogCharCap, *enableOtelTracing, metricsManager)
 }
 
+func instantiateKubeClient() (*kubernetes.Clientset, error) {
+	cfg, err := rest.InClusterConfig()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create REST Config for k8s client: %w", err)
+	}
+	kubeClient, err := kubernetes.NewForConfig(cfg)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create k8s client: %w", err)
+	}
+	return kubeClient, nil
+}
+
 func notEmpty(v string) bool {
 	return v != ""
 }

deploy/kubernetes/base/controller/cluster_setup.yaml

@@ -205,7 +205,7 @@ rules:
     verbs:     ['use']
     resourceNames:
     - csi-gce-pd-node-psp
-  - apiGroups: [""]
+  - apiGroups: [""]  # The core API group
     resources: ["nodes"]
     verbs: ["get", "list"]
 ---
@@ -220,6 +220,9 @@ rules:
     verbs:     ['use']
     resourceNames:
     - csi-gce-pd-node-psp-win
+  - apiGroups: [""]  # The core API group
+    resources: ["nodes"]
+    verbs: ["get", "list"]
 ---
 
 apiVersion: rbac.authorization.k8s.io/v1

pkg/common/constants.go

@@ -18,7 +18,8 @@ package common
 
 const (
 	// Keys for Topology. This key will be shared amongst drivers from GCP
-	TopologyKeyZone = "topology.gke.io/zone"
+	TopologyKeyPrefix = "topology.gke.io"
+	TopologyKeyZone   = TopologyKeyPrefix + "/zone"
 
 	// VolumeAttributes for Partition
 	VolumeAttributePartition = "partition"

pkg/common/utils.go

@@ -764,3 +764,11 @@ func MapNumber(num int64) int64 {
 	}
 	return 0
 }
+
+func IsGKETopologyLabel(key string) bool {
+	return strings.HasPrefix(key, TopologyKeyPrefix)
+}
+
+func TopologyLabelKey(diskType string) string {
+	return fmt.Sprintf("%s/%s", TopologyKeyPrefix, diskType)
+}

pkg/gce-cloud-provider/metadata/fake.go

@@ -20,13 +20,13 @@ type fakeServiceManager struct{}
 
 var _ MetadataService = &fakeServiceManager{}
 
-const (
-	FakeZone    = "country-region-zone"
-	FakeProject = "test-project"
+var (
+	FakeMachineType = "n1-standard-1"
+	FakeZone        = "country-region-zone"
+	FakeProject     = "test-project"
+	FakeName        = "test-name"
 )
 
-var FakeMachineType = "n1-standard-1"
-
 func NewFakeService() MetadataService {
 	return &fakeServiceManager{}
 }
@@ -40,7 +40,7 @@ func (manager *fakeServiceManager) GetProject() string {
 }
 
 func (manager *fakeServiceManager) GetName() string {
-	return "test-name"
+	return FakeName
 }
 
 func (manager *fakeServiceManager) GetMachineType() string {
@@ -50,3 +50,11 @@ func (manager *fakeServiceManager) GetMachineType() string {
 func SetMachineType(s string) {
 	FakeMachineType = s
 }
+
+func SetZone(s string) {
+	FakeZone = s
+}
+
+func SetName(s string) {
+	FakeName = s
+}

pkg/gce-pd-csi-driver/controller.go

@@ -120,6 +120,12 @@ type GCEControllerServer struct {
 	// Embed UnimplementedControllerServer to ensure the driver returns Unimplemented for any
 	// new RPC methods that might be introduced in future versions of the spec.
 	csi.UnimplementedControllerServer
+
+	EnableDiskTopology bool
+}
+
+type GCEControllerServerArgs struct {
+	EnableDiskTopology bool
 }
 
 type MultiZoneVolumeHandleConfig struct {
@@ -320,7 +326,7 @@ func (gceCS *GCEControllerServer) createVolumeInternal(ctx context.Context, req
 	if len(req.GetName()) == 0 {
 		return nil, status.Error(codes.InvalidArgument, "CreateVolume Name must be provided")
 	}
-	if volumeCapabilities == nil || len(volumeCapabilities) == 0 {
+	if len(volumeCapabilities) == 0 {
 		return nil, status.Error(codes.InvalidArgument, "CreateVolume Volume capabilities must be provided")
 	}
 
@@ -517,7 +523,7 @@ func (gceCS *GCEControllerServer) createMultiZoneDisk(ctx context.Context, req *
 	// Use the first response as a template
 	volumeId := fmt.Sprintf("projects/%s/zones/%s/disks/%s", gceCS.CloudProvider.GetDefaultProject(), common.MultiZoneValue, req.GetName())
 	klog.V(4).Infof("CreateVolume succeeded for multi-zone disks in zones %s: %v", zones, multiZoneVolKey)
-	return generateCreateVolumeResponseWithVolumeId(createdDisks[0], zones, params, dataCacheParams, enableDataCache, volumeId), nil
+	return gceCS.generateCreateVolumeResponseWithVolumeId(createdDisks[0], zones, params, dataCacheParams, enableDataCache, volumeId), nil
 }
 
 func (gceCS *GCEControllerServer) getZonesWithDiskNameAndType(ctx context.Context, name string, diskType string) ([]string, error) {
@@ -623,7 +629,7 @@ func (gceCS *GCEControllerServer) createSingleDeviceDisk(ctx context.Context, re
 		return nil, common.LoggedError("CreateVolume failed: %v", err)
 	}
 
-	return generateCreateVolumeResponseWithVolumeId(disk, zones, params, dataCacheParams, enableDataCache, volumeID), err
+	return gceCS.generateCreateVolumeResponseWithVolumeId(disk, zones, params, dataCacheParams, enableDataCache, volumeID), err
 }
 
 func getAccessMode(req *csi.CreateVolumeRequest, params common.DiskParameters) (string, error) {
@@ -2304,9 +2310,11 @@ func getZonesFromTopology(topList []*csi.Topology) ([]string, error) {
 func getZoneFromSegment(seg map[string]string) (string, error) {
 	var zone string
 	for k, v := range seg {
-		switch k {
-		case common.TopologyKeyZone:
+		switch {
+		case k == common.TopologyKeyZone:
 			zone = v
+		case common.IsGKETopologyLabel(k):
+			continue
 		default:
 			return "", fmt.Errorf("topology segment has unknown key %v", k)
 		}
@@ -2396,21 +2404,30 @@ func extractVolumeContext(context map[string]string) (*PDCSIContext, error) {
 		case contextForceAttach:
 			b, err := common.ConvertStringToBool(val)
 			if err != nil {
-				return nil, fmt.Errorf("Bad volume context force attach: %v", err)
+				return nil, fmt.Errorf("bad volume context force attach: %w", err)
 			}
 			info.ForceAttach = b
 		}
 	}
 	return info, nil
 }
 
-func generateCreateVolumeResponseWithVolumeId(disk *gce.CloudDisk, zones []string, params common.DiskParameters, dataCacheParams common.DataCacheParameters, enableDataCache bool, volumeId string) *csi.CreateVolumeResponse {
+func (gceCS *GCEControllerServer) generateCreateVolumeResponseWithVolumeId(disk *gce.CloudDisk, zones []string, params common.DiskParameters, dataCacheParams common.DataCacheParameters, enableDataCache bool, volumeId string) *csi.CreateVolumeResponse {
 	tops := []*csi.Topology{}
 	for _, zone := range zones {
-		tops = append(tops, &csi.Topology{
-			Segments: map[string]string{common.TopologyKeyZone: zone},
-		})
+		top := &csi.Topology{
+			Segments: map[string]string{
+				common.TopologyKeyZone: zone,
+			},
+		}
+
+		if gceCS.EnableDiskTopology {
+			top.Segments[common.TopologyLabelKey(params.DiskType)] = "true"
+		}
+
+		tops = append(tops, top)
 	}
+
 	realDiskSizeBytes := common.GbToBytes(disk.GetSizeGb())
 	createResp := &csi.CreateVolumeResponse{
 		Volume: &csi.Volume{