From c895312131addfea74bf9ef6ae7c9233acda8e61 Mon Sep 17 00:00:00 2001
From: Mehran Kholdi <semekh@hamravesh.com>
Date: Tue, 14 Jul 2020 15:55:00 +0430
Subject: [PATCH] Implement `GET_VOLUME_STATS` capability

---
 metrics.py          | 43 ++++++++++++++++++++++++++++++++-----------
 rawfile_servicer.py | 23 +++++++++++++++++++++++
 2 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/metrics.py b/metrics.py
index a58cc68..1defba7 100644
--- a/metrics.py
+++ b/metrics.py
@@ -10,6 +10,27 @@ import rawfile_util
 from rawfile_util import attached_loops
 
 
+def volume_stats(volume_id: str) -> dict:
+    img_file = rawfile_util.img_file(volume_id)
+    dev_stat = img_file.stat()
+    stats = {
+        "dev_size": dev_stat.st_size,
+        "dev_free": dev_stat.st_size - dev_stat.st_blocks * 512,
+    }
+    mountpoint = volume_to_mountpoint(img_file)
+    if mountpoint is not None:
+        fs_stat = os.statvfs(mountpoint)
+        stats.update(
+            {
+                "fs_size": fs_stat.f_frsize * fs_stat.f_blocks,
+                "fs_free": fs_stat.f_frsize * fs_stat.f_bfree,
+                "fs_files": fs_stat.f_files,
+                "fs_files_free": fs_stat.f_ffree,
+            }
+        )
+    return stats
+
+
 class VolumeStatsCollector(object):
     def collect(self):
         VOLUME_ID = "volume_id"
@@ -43,18 +64,18 @@ class VolumeStatsCollector(object):
         )
 
         for volume_id in rawfile_util.list_all_volumes():
-            img_file = rawfile_util.img_file(volume_id)
             labels = [volume_id]
-            dev_stat = img_file.stat()
-            dev_size.add_metric(labels, dev_stat.st_size)
-            dev_free.add_metric(labels, dev_stat.st_size - dev_stat.st_blocks * 512)
-            mountpoint = volume_to_mountpoint(img_file)
-            if mountpoint is not None:
-                fs_stat = os.statvfs(mountpoint)
-                fs_size.add_metric(labels, fs_stat.f_frsize * fs_stat.f_blocks)
-                fs_free.add_metric(labels, fs_stat.f_frsize * fs_stat.f_bfree)
-                fs_files.add_metric(labels, fs_stat.f_files)
-                fs_files_free.add_metric(labels, fs_stat.f_ffree)
+            key_to_gauge = {
+                "dev_size": dev_size,
+                "dev_free": dev_free,
+                "fs_size": fs_size,
+                "fs_free": fs_free,
+                "fs_files": fs_files,
+                "fs_files_free": fs_files_free,
+            }
+            stats = volume_stats(volume_id)
+            for key in stats.keys():
+                key_to_gauge[key].add_metric(labels, stats[key])
 
         return [fs_size, fs_free, fs_files, fs_files_free, dev_size, dev_free]
 
diff --git a/rawfile_servicer.py b/rawfile_servicer.py
index aa78cf9..dc39222 100644
--- a/rawfile_servicer.py
+++ b/rawfile_servicer.py
@@ -6,6 +6,7 @@ from google.protobuf.wrappers_pb2 import BoolValue
 import rawfile_util
 from csi import csi_pb2, csi_pb2_grpc
 from declarative import be_mounted, be_unmounted, be_symlink, be_absent
+from metrics import volume_stats
 from orchestrator.k8s import volume_to_node, run_on_node
 from rawfile_util import attach_loop, detach_loops
 from remote import init_rawfile, scrub, expand_rawfile
@@ -55,6 +56,7 @@ class RawFileNodeServicer(csi_pb2_grpc.NodeServicer):
         return csi_pb2.NodeGetCapabilitiesResponse(
             capabilities=[
                 Cap(rpc=Cap.RPC(type=Cap.RPC.STAGE_UNSTAGE_VOLUME)),
+                Cap(rpc=Cap.RPC(type=Cap.RPC.GET_VOLUME_STATS)),
                 Cap(rpc=Cap.RPC(type=Cap.RPC.EXPAND_VOLUME)),
             ]
         )
@@ -105,6 +107,27 @@ class RawFileNodeServicer(csi_pb2_grpc.NodeServicer):
         detach_loops(img_file)
         return csi_pb2.NodeUnstageVolumeResponse()
 
+    @log_grpc_request
+    def NodeGetVolumeStats(self, request, context):
+        volume_id = request.volume_id
+        stats = volume_stats(volume_id)
+        return csi_pb2.NodeGetVolumeStatsResponse(
+            usage=[
+                csi_pb2.VolumeUsage(
+                    available=stats["fs_free"],
+                    total=stats["fs_size"],
+                    used=stats["fs_size"] - stats["fs_free"],
+                    unit=csi_pb2.VolumeUsage.Unit.BYTES,
+                ),
+                csi_pb2.VolumeUsage(
+                    available=stats["fs_files_free"],
+                    total=stats["fs_files"],
+                    used=stats["fs_files"] - stats["fs_files_free"],
+                    unit=csi_pb2.VolumeUsage.Unit.INODES,
+                ),
+            ]
+        )
+
     @log_grpc_request
     def NodeExpandVolume(self, request, context):
         volume_id = request.volume_id