Fix Gen4 Custom VM Cases

Author: sunnylovestiramisu

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

@@ -654,10 +654,8 @@ func (ns *GCENodeServer) NodeGetInfo(ctx context.Context, req *csi.NodeGetInfoRe
 
 	nodeID := common.CreateNodeID(ns.MetadataService.GetProject(), ns.MetadataService.GetZone(), ns.MetadataService.GetName())
 
+	// err is always nil
 	volumeLimits, err := ns.GetVolumeLimits(ctx)
-	if err != nil {
-		klog.Errorf("GetVolumeLimits failed: %v", err.Error())
-	}
 
 	resp := &csi.NodeGetInfoResponse{
 		NodeId:             nodeID,
@@ -843,13 +841,19 @@ func (ns *GCENodeServer) GetVolumeLimits(ctx context.Context) (int64, error) {
 	gen4MachineTypesPrefix := []string{"c4a-", "c4-", "n4-"}
 	for _, gen4Prefix := range gen4MachineTypesPrefix {
 		if strings.HasPrefix(machineType, gen4Prefix) {
-			cpuString := machineType[strings.LastIndex(machineType, "-")+1:]
-			cpus, err := strconv.ParseInt(cpuString, 10, 64)
-			if err != nil {
-				return volumeLimitSmall, fmt.Errorf("invalid cpuString %s for machine type: %v", cpuString, machineType)
+			machineTypeSlice := strings.Split(machineType, "-")
+			if len(machineTypeSlice) > 2 {
+				cpuString := machineTypeSlice[2]
+				cpus, err := strconv.ParseInt(cpuString, 10, 64)
+				if err != nil {
+					klog.Warningf("invalid cpuString %s for machine type: %v", cpuString, machineType)
+					return volumeLimitBig, nil
+				}
+				return common.MapNumber(cpus), nil
+			} else {
+				klog.Warningf("unconventional machine type: %v", machineType)
+				return volumeLimitBig, nil
 			}
-			return common.MapNumber(cpus), nil
-
 		}
 		if strings.HasPrefix(machineType, "x4-") {
 			return x4HyperdiskLimit, nil

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

@@ -288,11 +288,26 @@ func TestNodeGetVolumeLimits(t *testing.T) {
 			machineType:    "n4-standard-16",
 			expVolumeLimit: 31,
 		},
+		{
+			name:           "n4-micro", // This type does not exist, but testing edge cases
+			machineType:    "n4-micro",
+			expVolumeLimit: volumeLimitBig,
+		},
 		{
 			name:           "n4-highcpu-4",
 			machineType:    "n4-highcpu-4",
 			expVolumeLimit: 15,
 		},
+		{
+			name:           "n4-custom-8-12345-ext",
+			machineType:    "n4-custom-8-12345-ext",
+			expVolumeLimit: 23,
+		},
+		{
+			name:           "n4-custom-16-12345",
+			machineType:    "n4-custom-16-12345",
+			expVolumeLimit: 31,
+		},
 		{
 			name:           "invalid gen4 machine type",
 			machineType:    "n4-highcpu-4xyz",