From 2eb85eed7535f0452a9fc68ff2aadc2133573b79 Mon Sep 17 00:00:00 2001
From: illustris <rharikrishnan95@gmail.com>
Date: Thu, 5 Sep 2024 11:38:43 +0530
Subject: [PATCH] collect disk size

---
 src/pvemon/__init__.py  | 10 ++++++++++
 src/qmblock/__init__.py | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/src/pvemon/__init__.py b/src/pvemon/__init__.py
index 803dfe1..bb7e6fa 100644
--- a/src/pvemon/__init__.py
+++ b/src/pvemon/__init__.py
@@ -41,6 +41,8 @@ gauge_settings = [
     ('kvm_io_write_chars', 'Number of bytes written including buffers', ['id']),
 
     ('kvm_nic_queues', 'Number of queues in multiqueue config', ['id', 'ifname']),
+
+    ('kvm_disk_size', 'Size of virtual disk', ['id', 'disk_name']),
 ]
 
 label_flags = [ "-id", "-name", "-cpu" ]
@@ -163,6 +165,8 @@ def collect_kvm_metrics():
         except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
             continue
 
+    # Try to find ID to pool mapping here
+
     for proc, cmdline, id in procs:
         # Extract vm labels from cmdline
         info_label_dict = {get_label_name(l): flag_to_label_value(cmdline,l) for l in label_flags}
@@ -221,6 +225,12 @@ def collect_kvm_metrics():
                 disk_labels = {"id": id, "disk_name": disk_name}
                 prom_disk_info = create_or_get_info("kvm_disk", disk_labels.keys(), dynamic_infos, info_lock)
                 prom_disk_info.add_metric(disk_labels.values(), disk_info)
+                disk_size = qmblock.get_disk_size(disk_info["disk_path"], disk_info["disk_type"])
+                if disk_size == None and disk_info["disk_type"] != "qcow2":
+                    logging.debug(f"collect_kvm_metrics: failed to get disk size for {disk_info=}")
+                else:
+                    gauge_dict["kvm_disk_size"].add_metric([id, disk_name], qmblock.get_disk_size(disk_info["disk_path"], disk_info["disk_type"]))
+
 
         list(executor.map(map_netstat_proc, [ proc[2] for proc in procs ]))
         list(executor.map(map_disk_proc, [ proc[2] for proc in procs ]))
diff --git a/src/qmblock/__init__.py b/src/qmblock/__init__.py
index 86fd513..641b16c 100644
--- a/src/qmblock/__init__.py
+++ b/src/qmblock/__init__.py
@@ -2,6 +2,7 @@ import pexpect
 import re
 import os
 import json
+import stat
 
 import pvecommon
 
@@ -94,6 +95,24 @@ def extract_disk_info_from_monitor(vm_id, retries = 0):
                 disks_map[disk_name]["detect_zeroes"] = "on"
     return disks_map
 
+def get_disk_size(disk_path, disk_type):
+    if stat.S_ISBLK(os.stat(disk_path).st_mode):
+        disk_name = os.path.basename(os.path.realpath(disk_path))
+        size_file_path = f"/sys/block/{disk_name}/size"
+        sector_size_file_path = f"/sys/block/{disk_name}/queue/hw_sector_size"
+
+        with open(size_file_path, 'r') as f:
+            sectors = int(f.read().strip())
+
+            with open(sector_size_file_path, 'r') as sector_size_file:
+                sector_size = int(sector_size_file.read().strip())
+
+                size_in_bytes = sectors * sector_size
+    else:
+        size_in_bytes = os.path.getsize(disk_path)
+
+    return size_in_bytes
+
 if __name__ == "__main__":
     import json
     import sys