From 87c58833373b7e85175a25caee12d9bb4a35eaa7 Mon Sep 17 00:00:00 2001
From: Pengfei Ni <peni@microsoft.com>
Date: Mon, 23 Jul 2018 14:56:17 +0800
Subject: [PATCH] Implement GetLabelsForVolume for AzureDisk

---
 pkg/cloudprovider/providers/azure/azure.go    |  3 +
 .../azure/azure_managedDiskController.go      | 60 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/pkg/cloudprovider/providers/azure/azure.go b/pkg/cloudprovider/providers/azure/azure.go
index f06faf7d323..80a54f406de 100644
--- a/pkg/cloudprovider/providers/azure/azure.go
+++ b/pkg/cloudprovider/providers/azure/azure.go
@@ -59,6 +59,9 @@ var (
 	defaultExcludeMasterFromStandardLB = true
 )
 
+// Azure implements PVLabeler.
+var _ cloudprovider.PVLabeler = (*Cloud)(nil)
+
 // Config holds the configuration parsed from the --cloud-config flag
 // All fields are required unless otherwise specified
 type Config struct {
diff --git a/pkg/cloudprovider/providers/azure/azure_managedDiskController.go b/pkg/cloudprovider/providers/azure/azure_managedDiskController.go
index 61c6515ef91..22d3d3bff30 100644
--- a/pkg/cloudprovider/providers/azure/azure_managedDiskController.go
+++ b/pkg/cloudprovider/providers/azure/azure_managedDiskController.go
@@ -17,16 +17,21 @@ limitations under the License.
 package azure
 
 import (
+	"context"
 	"fmt"
 	"path"
+	"strconv"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
 	"github.com/Azure/azure-sdk-for-go/services/storage/mgmt/2017-10-01/storage"
 	"github.com/golang/glog"
 
+	"k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	kwait "k8s.io/apimachinery/pkg/util/wait"
+	kubeletapis "k8s.io/kubernetes/pkg/kubelet/apis"
+	"k8s.io/kubernetes/pkg/volume"
 	"k8s.io/kubernetes/pkg/volume/util"
 )
 
@@ -201,3 +206,58 @@ func getResourceGroupFromDiskURI(diskURI string) (string, error) {
 	}
 	return fields[4], nil
 }
+
+// GetLabelsForVolume implements PVLabeler.GetLabelsForVolume
+func (c *Cloud) GetLabelsForVolume(ctx context.Context, pv *v1.PersistentVolume) (map[string]string, error) {
+	// Ignore if not AzureDisk.
+	if pv.Spec.AzureDisk == nil {
+		return nil, nil
+	}
+
+	// Ignore any volumes that are being provisioned
+	if pv.Spec.AzureDisk.DiskName == volume.ProvisionedVolumeName {
+		return nil, nil
+	}
+
+	return c.GetAzureDiskLabels(pv.Spec.AzureDisk.DataDiskURI)
+}
+
+// GetAzureDiskLabels gets availability zone labels for Azuredisk.
+func (c *Cloud) GetAzureDiskLabels(diskURI string) (map[string]string, error) {
+	// Get disk's resource group.
+	diskName := path.Base(diskURI)
+	resourceGroup, err := getResourceGroupFromDiskURI(diskURI)
+	if err != nil {
+		glog.Errorf("Failed to get resource group for AzureDisk %q: %v", diskName, err)
+		return nil, err
+	}
+
+	// Get information of the disk.
+	ctx, cancel := getContextWithCancel()
+	defer cancel()
+	disk, err := c.DisksClient.Get(ctx, resourceGroup, diskName)
+	if err != nil {
+		glog.Errorf("Failed to get information for AzureDisk %q: %v", diskName, err)
+		return nil, err
+	}
+
+	// Check whether availability zone is specified.
+	if disk.Zones == nil || len(*disk.Zones) == 0 {
+		glog.V(4).Infof("Azure disk %q is not zoned", diskName)
+		return nil, nil
+	}
+
+	zones := *disk.Zones
+	zoneID, err := strconv.Atoi(zones[0])
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse zone %v for AzureDisk %v: %v", zones, diskName, err)
+	}
+
+	zone := c.makeZone(zoneID)
+	glog.V(4).Infof("Get zone %q for Azure disk %q", zone, diskName)
+	labels := map[string]string{
+		kubeletapis.LabelZoneRegion:        c.Location,
+		kubeletapis.LabelZoneFailureDomain: zone,
+	}
+	return labels, nil
+}