From 1602df8596b38c0f77730c0d9ca6b2de026252b5 Mon Sep 17 00:00:00 2001
From: andyzhangx <xiazhang@microsoft.com>
Date: Thu, 14 Sep 2023 02:33:47 +0000
Subject: [PATCH] fix: add VolumeStats cache to avoid massive statfs calls

---
 pkg/blob/blob.go       | 10 ++++++++++
 pkg/blob/blob_test.go  |  1 +
 pkg/blob/nodeserver.go | 23 +++++++++++++++++++++--
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/pkg/blob/blob.go b/pkg/blob/blob.go
index 137f80c00..be67bbd7a 100644
--- a/pkg/blob/blob.go
+++ b/pkg/blob/blob.go
@@ -166,6 +166,7 @@ type DriverOptions struct {
 	MountPermissions                       uint64
 	KubeAPIQPS                             float64
 	KubeAPIBurst                           int
+	VolStatsCacheExpireInMinutes           int
 }
 
 // Driver implements all interfaces of CSI drivers
@@ -203,6 +204,8 @@ type Driver struct {
 	dataPlaneAPIVolCache *azcache.TimedCache
 	// a timed cache storing account search history (solve account list throttling issue)
 	accountSearchCache *azcache.TimedCache
+	// a timed cache storing volume stats <volumeID, volumeStats>
+	volStatsCache *azcache.TimedCache
 }
 
 // NewDriver Creates a NewCSIDriver object. Assumes vendor version is equal to driver version &
@@ -240,6 +243,13 @@ func NewDriver(options *DriverOptions) *Driver {
 	if d.dataPlaneAPIVolCache, err = azcache.NewTimedcache(10*time.Minute, getter); err != nil {
 		klog.Fatalf("%v", err)
 	}
+
+	if options.VolStatsCacheExpireInMinutes <= 0 {
+		options.VolStatsCacheExpireInMinutes = 10 // default expire in 10 minutes
+	}
+	if d.volStatsCache, err = azcache.NewTimedcache(time.Duration(options.VolStatsCacheExpireInMinutes)*time.Minute, getter); err != nil {
+		klog.Fatalf("%v", err)
+	}
 	return &d
 }
 
diff --git a/pkg/blob/blob_test.go b/pkg/blob/blob_test.go
index 54f355e76..53e074197 100644
--- a/pkg/blob/blob_test.go
+++ b/pkg/blob/blob_test.go
@@ -92,6 +92,7 @@ func TestNewDriver(t *testing.T) {
 	fakedriver.Version = driverVersion
 	fakedriver.accountSearchCache = driver.accountSearchCache
 	fakedriver.dataPlaneAPIVolCache = driver.dataPlaneAPIVolCache
+	fakedriver.volStatsCache = driver.volStatsCache
 	assert.Equal(t, driver, fakedriver)
 }
 
diff --git a/pkg/blob/nodeserver.go b/pkg/blob/nodeserver.go
index 0ef75af5f..06ba533c2 100644
--- a/pkg/blob/nodeserver.go
+++ b/pkg/blob/nodeserver.go
@@ -28,6 +28,7 @@ import (
 	"time"
 
 	volumehelper "sigs.k8s.io/blob-csi-driver/pkg/util"
+	azcache "sigs.k8s.io/cloud-provider-azure/pkg/cache"
 
 	"github.com/Azure/azure-sdk-for-go/storage"
 	"github.com/container-storage-interface/spec/lib/go/csi"
@@ -475,6 +476,17 @@ func (d *Driver) NodeGetVolumeStats(ctx context.Context, req *csi.NodeGetVolumeS
 		return nil, status.Error(codes.InvalidArgument, "NodeGetVolumeStats volume path was empty")
 	}
 
+	// check if the volume stats is cached
+	cache, err := d.volStatsCache.Get(req.VolumeId, azcache.CacheReadTypeDefault)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, err.Error())
+	}
+	if cache != nil {
+		resp := cache.(csi.NodeGetVolumeStatsResponse)
+		klog.V(6).Infof("NodeGetVolumeStats: volume stats for volume %s path %s is cached", req.VolumeId, req.VolumePath)
+		return &resp, nil
+	}
+
 	if _, err := os.Lstat(req.VolumePath); err != nil {
 		if os.IsNotExist(err) {
 			return nil, status.Errorf(codes.NotFound, "path %s does not exist", req.VolumePath)
@@ -482,6 +494,8 @@ func (d *Driver) NodeGetVolumeStats(ctx context.Context, req *csi.NodeGetVolumeS
 		return nil, status.Errorf(codes.Internal, "failed to stat file %s: %v", req.VolumePath, err)
 	}
 
+	klog.V(6).Infof("NodeGetVolumeStats: begin to get VolumeStats on volume %s path %s", req.VolumeId, req.VolumePath)
+
 	volumeMetrics, err := volume.NewMetricsStatFS(req.VolumePath).GetMetrics()
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed to get metrics: %v", err)
@@ -513,7 +527,7 @@ func (d *Driver) NodeGetVolumeStats(ctx context.Context, req *csi.NodeGetVolumeS
 		return nil, status.Errorf(codes.Internal, "failed to transform disk inodes used(%v)", volumeMetrics.InodesUsed)
 	}
 
-	return &csi.NodeGetVolumeStatsResponse{
+	resp := &csi.NodeGetVolumeStatsResponse{
 		Usage: []*csi.VolumeUsage{
 			{
 				Unit:      csi.VolumeUsage_BYTES,
@@ -528,7 +542,12 @@ func (d *Driver) NodeGetVolumeStats(ctx context.Context, req *csi.NodeGetVolumeS
 				Used:      inodesUsed,
 			},
 		},
-	}, nil
+	}
+
+	klog.V(6).Infof("NodeGetVolumeStats: volume stats for volume %s path %s is %v", req.VolumeId, req.VolumePath, resp)
+	// cache the volume stats per volume
+	d.volStatsCache.Set(req.VolumeId, *resp)
+	return resp, nil
 }
 
 // ensureMountPoint: create mount point if not exists