From f246726cb6b9b8ab4d77d8a2ef70f59c9bab25e0 Mon Sep 17 00:00:00 2001
From: Sneha Aradhey <snehaaradhey@google.com>
Date: Thu, 3 Apr 2025 04:45:27 +0000
Subject: [PATCH] update cache logic to calculate chunk size based on toatl
 cache

---
 pkg/gce-pd-csi-driver/cache.go      | 61 +++++++++++++++++++++++++++--
 pkg/gce-pd-csi-driver/cache_test.go | 60 ++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+), 4 deletions(-)

diff --git a/pkg/gce-pd-csi-driver/cache.go b/pkg/gce-pd-csi-driver/cache.go
index 32ef0faa0..633ae473e 100644
--- a/pkg/gce-pd-csi-driver/cache.go
+++ b/pkg/gce-pd-csi-driver/cache.go
@@ -175,10 +175,18 @@ func setupCaching(devicePath string, req *csi.NodeStageVolumeRequest, nodeId str
 		klog.V(4).Infof("Assuming valid data cache size and mode, resizing cache is not supported")
 	} else {
 		cacheSize := req.GetPublishContext()[common.ContextDataCacheSize]
-		chunkSize, err := fetchChunkSizeKiB(cacheSize)
+		maxChunkSizeStr := strconv.FormatInt(int64(maxChunkSize/KiB), 10)
+		var chunkSize string
+		cachePvSize, err := fetchPvSizeGiB()
 		if err != nil {
-			klog.Errorf("Errored to fetch cache size, verify the data-cache-size is valid: got %v, error: %q", cacheSize, err)
-			return mainDevicePath, err
+			klog.Errorf("Errored while fetching PV size, got %v, falling back to default chunkSize of %v", err, maxChunkSize)
+			chunkSize = maxChunkSizeStr
+		} else {
+			chunkSize, err = fetchChunkSizeKiB(cachePvSize)
+			if err != nil {
+				klog.Errorf("Errored to fetch cache size, verify the data-cache-size is valid: got %v, error: %q", chunkSize, err)
+				chunkSize = maxChunkSizeStr
+			}
 		}
 		// Check if LV exists
 		info, err = common.RunCommand("" /* pipedCmd */, nil /* pipedCmdArg */, "lvs", args...)
@@ -642,7 +650,7 @@ func watchDiskDetaches(watcher *fsnotify.Watcher, nodeName string, errorCh chan
 			// In case of an event i.e. creation or deletion of any new PV, we update the VG metadata.
 			// This might include some non-LVM changes, no harm in updating metadata multiple times.
 			reduceVolumeGroup(getVolumeGroupName(nodeName), true)
-			klog.V(2).Infof("disk attach/detach event %#v\n", event)
+			klog.V(6).Infof("disk attach/detach event %#v\n", event)
 		}
 	}
 }
@@ -674,3 +682,48 @@ func addRaidedLSSDToVg(vgName, lssdPath string) error {
 	}
 	return nil
 }
+
+func fetchPvSizeGiB() (string, error) {
+	args := []string{
+		"--select",
+		"-o",
+		"--noheadings",
+		"pv_size",
+		"--units=b",
+	}
+	// RAIDed device is always registered with its /dev/md127 equivalent in VG so cannot check it directly based on the RAIDed LSSD path which could be /dev/md/csi-driver-data-cache
+	info, err := common.RunCommand("grep" /* pipedCmd */, []string{"/dev/md"} /* pipedCmdArg */, "pvs", args...)
+	if err != nil {
+		return "", fmt.Errorf("errored while fetching PV size %v: %s", err, info)
+	}
+	infoString := strings.TrimSpace(string(info))
+	infoSlice := strings.Fields(infoString)
+	pvSize, err := fetchNumberGiB(infoSlice)
+	if err != nil {
+		return "", fmt.Errorf("Error fetching PV size for cache %v", err)
+	}
+	return pvSize, nil
+
+}
+
+func fetchNumberGiB(infoSlice []string) (string, error) {
+	re, err := regexp.Compile("^[0-9]+B$")
+	if err != nil {
+		return "", fmt.Errorf("Failed to compile regex match %v", err)
+	}
+	var pvSize string
+	for _, i := range infoSlice {
+		if re.MatchString(i) {
+			pvSize, err = strings.TrimSuffix(i, "B"), nil
+			if err != nil {
+				return "", fmt.Errorf("Failed to extract PV size %v", err)
+			}
+			break
+		}
+	}
+	pvSizeInt, err := strconv.ParseFloat(pvSize, 64)
+	if err != nil {
+		return "", fmt.Errorf("Error while fetching PV size for cache %v", err)
+	}
+	return strconv.FormatInt(int64(math.Ceil(pvSizeInt/GiB)), 10) + "GiB", nil
+}
diff --git a/pkg/gce-pd-csi-driver/cache_test.go b/pkg/gce-pd-csi-driver/cache_test.go
index 4574dde4d..340c924a1 100644
--- a/pkg/gce-pd-csi-driver/cache_test.go
+++ b/pkg/gce-pd-csi-driver/cache_test.go
@@ -55,3 +55,63 @@ func TestFetchChunkSizeKiB(t *testing.T) {
 	}
 
 }
+
+func TestFetchNumberGiB(t *testing.T) {
+	testCases := []struct {
+		name        string
+		stringInput []string
+		expOutput   string // Outputs value in GiB
+		expErr      bool
+	}{
+		{
+			name:        "valid input 1",
+			stringInput: []string{"5000000000B"},
+			expOutput:   "5GiB", //range defined in fetchChunkSizeKiB
+		},
+		{
+			name:        "valid input 2",
+			stringInput: []string{"375000000000B"}, // 1 LSSD attached
+			expOutput:   "350GiB",                  //range defined in fetchChunkSizeKiB
+		},
+		{
+			name:        "valid input 3",
+			stringInput: []string{"9000000000000B"}, // 24 LSSD attached
+			expOutput:   "8382GiB",                  //range defined in fetchChunkSizeKiB
+		},
+		{
+			name:        "valid input 4",
+			stringInput: []string{"Some text before ", "9000000000000B", "Some text after"}, // 24 LSSD attached
+			expOutput:   "8382GiB",                                                          //range defined in fetchChunkSizeKiB
+		},
+		{
+			name:        "invalid input 1",
+			stringInput: []string{"9000000000000"},
+			expErr:      true,
+		},
+		{
+			name:        "invalid input 2",
+			stringInput: []string{"A9000000000000B"},
+			expErr:      true,
+		},
+		{
+			name:        "valid input 5",
+			stringInput: []string{"900000B"}, // <1GiB gets rounded off to 0GiB
+			expOutput:   "1GiB",
+		},
+	}
+
+	for _, tc := range testCases {
+		v, err := fetchNumberGiB(tc.stringInput)
+		if err != nil {
+			if !tc.expErr {
+				t.Errorf("Errored %s", err)
+			}
+			continue
+		}
+		if v != tc.expOutput {
+			t.Errorf("Got %s want %s", v, tc.expOutput)
+		}
+
+	}
+
+}