Restrict data caching setup to labelled node pool

Author: Sneha-at

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

@@ -27,8 +27,11 @@ import (
 	"strings"
 	"time"
 
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
 	"k8s.io/klog/v2"
 
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/gcp-compute-persistent-disk-csi-driver/pkg/common"
 	"sigs.k8s.io/gcp-compute-persistent-disk-csi-driver/pkg/deviceutils"
 	gce "sigs.k8s.io/gcp-compute-persistent-disk-csi-driver/pkg/gce-cloud-provider/compute"
@@ -73,6 +76,7 @@ var (
 	enableStoragePoolsFlag        = flag.Bool("enable-storage-pools", false, "If set to true, the CSI Driver will allow volumes to be provisioned in Storage Pools")
 	enableControllerDataCacheFlag = flag.Bool("enable-controller-data-cache", false, "If set to true, the CSI Driver will allow volumes to be provisioned with data cache configuration")
 	enableNodeDataCacheFlag       = flag.Bool("enable-node-data-cache", false, "If set to true, the CSI Driver will allow volumes to be provisioned with data cache configuration")
+	nodeName                      = flag.String("node-name", "", "The node this driver is running on")
 
 	multiZoneVolumeHandleDiskTypesFlag = flag.String("multi-zone-volume-handle-disk-types", "", "Comma separated list of allowed disk types that can use the multi-zone volumeHandle. Used only if --multi-zone-volume-handle-enable")
 	multiZoneVolumeHandleEnableFlag    = flag.Bool("multi-zone-volume-handle-enable", false, "If set to true, the multi-zone volumeHandle feature will be enabled")
@@ -91,7 +95,9 @@ var (
 )
 
 const (
-	driverName = "pd.csi.storage.gke.io"
+	driverName          = "pd.csi.storage.gke.io"
+	dataCacheLabel      = "datacache-storage-gke-io"
+	dataCacheLabelValue = "enabled"
 )
 
 func init() {
@@ -235,12 +241,14 @@ func handle() {
 	}
 
 	if *enableNodeDataCacheFlag {
-		klog.V(2).Info("Raiding local ssds to setup data cache")
-		err := driver.RaidLocalSsds()
-		if err != nil {
-			klog.Fatalf("Failed to Raid local SSDs, unable to setup data caching, got error %v", err)
+		if nodeName == nil || *nodeName == "" {
+			klog.Fatalf("Data cache enabled, but --node-name not passed")
+		}
+		if err := setupDataCache(ctx, *nodeName); err != nil {
+			klog.Fatalf("DataCache setup failed: %v", err)
 		}
 	}
+
 	err = gceDriver.SetupGCEDriver(driverName, version, extraVolumeLabels, extraTags, identityServer, controllerServer, nodeServer)
 	if err != nil {
 		klog.Fatalf("Failed to initialize GCE CSI Driver: %v", err.Error())
@@ -289,3 +297,31 @@ func urlFlag(target **url.URL, name string, usage string) {
 		return err
 	})
 }
+
+func setupDataCache(ctx context.Context, nodeName string) error {
+	cfg, err := rest.InClusterConfig()
+	if err != nil {
+		return err
+	}
+	kubeClient, err := kubernetes.NewForConfig(cfg)
+	if err != nil {
+		return err
+	}
+	node, err := kubeClient.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+	if err != nil {
+		// We could retry, but this error will also crashloop the driver which may be as good a way to retry as any.
+		return err
+	}
+	if val, found := node.GetLabels()[dataCacheLabel]; !found || val != dataCacheLabelValue {
+		klog.V(2).Infof("Datacache not enabled for node %s; node label %s=%s and not %s", nodeName, dataCacheLabel, val, dataCacheLabelValue)
+		return nil
+	}
+	// Setup data cache only if enabled fro nodes
+	klog.V(2).Info("Raiding local ssds to setup data cache")
+	if err := driver.RaidLocalSsds(); err != nil {
+		return fmt.Errorf("Failed to Raid local SSDs, unable to setup data caching, got error %v", err)
+	}
+
+	klog.V(2).Infof("Datacache enabled for node %s", nodeName)
+	return nil
+}

deploy/kubernetes/base/controller/cluster_setup.yaml

@@ -199,6 +199,9 @@ rules:
     verbs:     ['use']
     resourceNames:
     - csi-gce-pd-node-psp
+  - apiGroups: [""]
+    resources: ["nodes"]
+    verbs: ["get", "list"]
 ---
 
 kind: ClusterRole

deploy/kubernetes/base/node_linux/node.yaml

@@ -46,8 +46,15 @@ spec:
             - "--v=5"
             - "--endpoint=unix:/csi/csi.sock"
             - "--run-controller-service=false"
+            - "--enable-node-data-cache"
+            - "--node-name=$(KUBE_NODE_NAME)"
           securityContext:
             privileged: true
+          env:
+            - name: KUBE_NODE_NAME
+              valueFrom:
+                fieldRef:
+                  fieldPath: spec.nodeName
           volumeMounts:
             - name: kubelet-dir
               mountPath: /var/lib/kubelet