From 8b31fcf3a238c78f8ca4b1b86633fe69450ea2d6 Mon Sep 17 00:00:00 2001
From: Jing Xu <jinxu@google.com>
Date: Mon, 4 Jun 2018 17:19:03 -0700
Subject: [PATCH] Set GCE PD attachable volume limit based on machineType

This PR implements the function to return attachable volume limit based
on machineType for GCE PD. This is part of the design in kubernetes/community#2051/
---
 pkg/volume/gce_pd/gce_pd.go | 43 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

diff --git a/pkg/volume/gce_pd/gce_pd.go b/pkg/volume/gce_pd/gce_pd.go
index d4f58784799..75f3d65f9f6 100644
--- a/pkg/volume/gce_pd/gce_pd.go
+++ b/pkg/volume/gce_pd/gce_pd.go
@@ -17,10 +17,12 @@ limitations under the License.
 package gce_pd
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path"
 	"strconv"
+	"strings"
 
 	"github.com/golang/glog"
 	"k8s.io/api/core/v1"
@@ -56,6 +58,18 @@ const (
 	gcePersistentDiskPluginName = "kubernetes.io/gce-pd"
 )
 
+// The constants are used to map from the machine type (number of CPUs) to the limit of
+// persistent disks that can be attached to an instance. Please refer to gcloud doc
+// https://cloud.google.com/compute/docs/disks/#increased_persistent_disk_limits
+const (
+	OneCPU         = 1
+	EightCPUs      = 8
+	VolumeLimit16  = 16
+	VolumeLimit32  = 32
+	VolumeLimit64  = 64
+	VolumeLimit128 = 128
+)
+
 func getPath(uid types.UID, volName string, host volume.VolumeHost) string {
 	return host.GetPodVolumeDir(uid, kstrings.EscapeQualifiedNameForDisk(gcePersistentDiskPluginName), volName)
 }
@@ -104,7 +118,7 @@ func (plugin *gcePersistentDiskPlugin) GetAccessModes() []v1.PersistentVolumeAcc
 
 func (plugin *gcePersistentDiskPlugin) GetVolumeLimits() (map[string]int64, error) {
 	volumeLimits := map[string]int64{
-		util.GCEVolumeLimitKey: 16,
+		util.GCEVolumeLimitKey: VolumeLimit16,
 	}
 	cloud := plugin.host.GetCloudProvider()
 
@@ -120,6 +134,33 @@ func (plugin *gcePersistentDiskPlugin) GetVolumeLimits() (map[string]int64, erro
 		return nil, fmt.Errorf("Expected gce cloud got %s", cloud.ProviderName())
 	}
 
+	instances, ok := cloud.Instances()
+	if !ok {
+		glog.Warning("Failed to get instances from cloud provider")
+		return volumeLimits, nil
+	}
+
+	instanceType, err := instances.InstanceType(context.TODO(), plugin.host.GetNodeName())
+	if err != nil {
+		glog.Errorf("Failed to get instance type from GCE cloud provider")
+		return volumeLimits, nil
+	}
+	if strings.HasPrefix(instanceType, "n1-") {
+		splits := strings.Split(instanceType, "-")
+		if len(splits) < 3 {
+			return volumeLimits, nil
+		}
+		last := splits[2]
+		if num, err := strconv.Atoi(last); err == nil {
+			if num == OneCPU {
+				volumeLimits[util.GCEVolumeLimitKey] = VolumeLimit32
+			} else if num < EightCPUs {
+				volumeLimits[util.GCEVolumeLimitKey] = VolumeLimit64
+			} else {
+				volumeLimits[util.GCEVolumeLimitKey] = VolumeLimit128
+			}
+		}
+	}
 	return volumeLimits, nil
 }