From 21c8a63689f6e0cf2d31aaf09c0452f3ea1de625 Mon Sep 17 00:00:00 2001
From: Pengfei Ni <peni@microsoft.com>
Date: Wed, 7 Feb 2018 21:57:11 +0800
Subject: [PATCH] Add cache for network security groups

---
 pkg/cloudprovider/providers/azure/azure.go    | 11 +++++-
 .../providers/azure/azure_backoff.go          |  7 +++-
 .../providers/azure/azure_loadbalancer.go     |  2 +-
 .../providers/azure/azure_test.go             |  1 +
 .../providers/azure/azure_wrap.go             | 37 ++++++++++++++++++-
 5 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/pkg/cloudprovider/providers/azure/azure.go b/pkg/cloudprovider/providers/azure/azure.go
index 6b40abdfeec..1fd3eed2fc0 100644
--- a/pkg/cloudprovider/providers/azure/azure.go
+++ b/pkg/cloudprovider/providers/azure/azure.go
@@ -135,8 +135,9 @@ type Cloud struct {
 	VirtualMachineScaleSetsClient   VirtualMachineScaleSetsClient
 	VirtualMachineScaleSetVMsClient VirtualMachineScaleSetVMsClient
 
-	vmCache *timedCache
-	lbCache *timedCache
+	vmCache  *timedCache
+	lbCache  *timedCache
+	nsgCache *timedCache
 
 	*BlobDiskController
 	*ManagedDiskController
@@ -259,6 +260,12 @@ func NewCloud(configReader io.Reader) (cloudprovider.Interface, error) {
 	}
 	az.lbCache = lbCache
 
+	nsgCache, err := az.newNSGCache()
+	if err != nil {
+		return nil, err
+	}
+	az.nsgCache = nsgCache
+
 	if err := initDiskControllers(&az); err != nil {
 		return nil, err
 	}
diff --git a/pkg/cloudprovider/providers/azure/azure_backoff.go b/pkg/cloudprovider/providers/azure/azure_backoff.go
index ddc2f69ff52..64c212d40a9 100644
--- a/pkg/cloudprovider/providers/azure/azure_backoff.go
+++ b/pkg/cloudprovider/providers/azure/azure_backoff.go
@@ -131,7 +131,12 @@ func (az *Cloud) CreateOrUpdateSGWithRetry(sg network.SecurityGroup) error {
 		resp := <-respChan
 		err := <-errChan
 		glog.V(10).Infof("SecurityGroupsClient.CreateOrUpdate(%s): end", *sg.Name)
-		return processRetryResponse(resp.Response, err)
+		done, err := processRetryResponse(resp.Response, err)
+		if done && err == nil {
+			// Invalidate the cache right after updating
+			az.lbCache.Delete(*sg.Name)
+		}
+		return done, err
 	})
 }
 
diff --git a/pkg/cloudprovider/providers/azure/azure_loadbalancer.go b/pkg/cloudprovider/providers/azure/azure_loadbalancer.go
index 8b4387cf63a..65b6f9c6287 100644
--- a/pkg/cloudprovider/providers/azure/azure_loadbalancer.go
+++ b/pkg/cloudprovider/providers/azure/azure_loadbalancer.go
@@ -819,7 +819,7 @@ func (az *Cloud) reconcileSecurityGroup(clusterName string, service *v1.Service,
 		ports = []v1.ServicePort{}
 	}
 
-	sg, err := az.SecurityGroupsClient.Get(az.ResourceGroup, az.SecurityGroupName, "")
+	sg, err := az.getSecurityGroup()
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/cloudprovider/providers/azure/azure_test.go b/pkg/cloudprovider/providers/azure/azure_test.go
index 8e9c1d8bbf6..8226815efa2 100644
--- a/pkg/cloudprovider/providers/azure/azure_test.go
+++ b/pkg/cloudprovider/providers/azure/azure_test.go
@@ -880,6 +880,7 @@ func getTestCloud() (az *Cloud) {
 	az.vmSet = newAvailabilitySet(az)
 	az.vmCache, _ = az.newVMCache()
 	az.lbCache, _ = az.newLBCache()
+	az.nsgCache, _ = az.newNSGCache()
 
 	return az
 }
diff --git a/pkg/cloudprovider/providers/azure/azure_wrap.go b/pkg/cloudprovider/providers/azure/azure_wrap.go
index c75bd62a70e..792946cc907 100644
--- a/pkg/cloudprovider/providers/azure/azure_wrap.go
+++ b/pkg/cloudprovider/providers/azure/azure_wrap.go
@@ -17,6 +17,7 @@ limitations under the License.
 package azure
 
 import (
+	"fmt"
 	"net/http"
 	"time"
 
@@ -29,8 +30,9 @@ import (
 )
 
 var (
-	vmCacheTTL = time.Minute
-	lbCacheTTL = 2 * time.Minute
+	vmCacheTTL  = time.Minute
+	lbCacheTTL  = 2 * time.Minute
+	nsgCacheTTL = 2 * time.Minute
 )
 
 // checkExistsFromError inspects an error and returns a true if err is nil,
@@ -152,6 +154,19 @@ func (az *Cloud) getAzureLoadBalancer(name string) (lb network.LoadBalancer, exi
 	return *(cachedLB.(*network.LoadBalancer)), true, nil
 }
 
+func (az *Cloud) getSecurityGroup() (nsg network.SecurityGroup, err error) {
+	securityGroup, err := az.nsgCache.Get(az.SecurityGroupName)
+	if err != nil {
+		return nsg, err
+	}
+
+	if securityGroup == nil {
+		return nsg, fmt.Errorf("nsg %q not found", az.SecurityGroupName)
+	}
+
+	return *(securityGroup.(*network.SecurityGroup)), nil
+}
+
 func (az *Cloud) newVMCache() (*timedCache, error) {
 	getter := func(key string) (interface{}, error) {
 		vm, err := az.VirtualMachinesClient.Get(az.ResourceGroup, key, compute.InstanceView)
@@ -187,3 +202,21 @@ func (az *Cloud) newLBCache() (*timedCache, error) {
 
 	return newTimedcache(lbCacheTTL, getter)
 }
+
+func (az *Cloud) newNSGCache() (*timedCache, error) {
+	getter := func(key string) (interface{}, error) {
+		nsg, err := az.SecurityGroupsClient.Get(az.ResourceGroup, key, "")
+		exists, realErr := checkResourceExistsFromError(err)
+		if realErr != nil {
+			return nil, realErr
+		}
+
+		if !exists {
+			return nil, nil
+		}
+
+		return &nsg, nil
+	}
+
+	return newTimedcache(nsgCacheTTL, getter)
+}