From 3d4a2873ae5ddf443b7ebe789795aac2a2f066ce Mon Sep 17 00:00:00 2001
From: bugclerk <40872210+bugclerk@users.noreply.github.com>
Date: Fri, 30 Aug 2024 09:17:26 -0700
Subject: [PATCH] NAS-130883 / 24.10-RC.1 / Make sure vm.query does not fail
 (by Qubad786) (#14384)

* Make sure vm.query does not fail

(cherry picked from commit dfbf5dae867b3818079ef3ee3c44668b1e93a588)

* Optimize getting counts of vms and pools in system

(cherry picked from commit e938d8b13fb54395d1a68115524da6c679f69f27)

---------

Co-authored-by: M. Rehan <mrehanlm93@gmail.com>
---
 src/middlewared/middlewared/plugins/reporting/rest.py | 5 +++--
 src/middlewared/middlewared/plugins/vm/connection.py  | 7 +++++--
 src/middlewared/middlewared/plugins/vm/vms.py         | 6 ++++--
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/middlewared/middlewared/plugins/reporting/rest.py b/src/middlewared/middlewared/plugins/reporting/rest.py
index 30edc30ea0397..83283df8656af 100644
--- a/src/middlewared/middlewared/plugins/reporting/rest.py
+++ b/src/middlewared/middlewared/plugins/reporting/rest.py
@@ -4,6 +4,7 @@
 from middlewared.service import accepts, Service
 from middlewared.schema import Str, Dict, Int
 from middlewared.utils.cpu import cpu_info
+from middlewared.utils.zfs import query_imported_fast_impl
 
 from .netdata import ClientConnectError, Netdata
 from .utils import calculate_disk_space_for_netdata, get_metrics_approximation, TIER_0_POINT_SIZE, TIER_1_POINT_SIZE
@@ -52,8 +53,8 @@ def calculated_metrics_count(self):
             len(self.middleware.call_sync('device.get_disks', False, True)),
             cpu_info()['core_count'],
             self.middleware.call_sync('interface.query', [], {'count': True}),
-            self.middleware.call_sync('zfs.pool.query', [], {'count': True}),
-            self.middleware.call_sync('vm.query', [], {'count': True}),
+            len(query_imported_fast_impl()),
+            self.middleware.call_sync('datastore.query', 'vm.vm', [], {'count': True}),
             len(glob.glob('/sys/fs/cgroup/**/*.service')),
         )
 
diff --git a/src/middlewared/middlewared/plugins/vm/connection.py b/src/middlewared/middlewared/plugins/vm/connection.py
index e415cb2e2b2e6..ab84cb9872147 100644
--- a/src/middlewared/middlewared/plugins/vm/connection.py
+++ b/src/middlewared/middlewared/plugins/vm/connection.py
@@ -65,7 +65,10 @@ def _check_connection_alive(self):
         if not self._is_libvirt_connection_alive():
             raise CallError('Failed to connect to libvirt')
 
-    def _check_setup_connection(self):
+    def _safely_check_setup_connection(self, timeout: int = 10):
         if not self._is_connection_alive():
-            self.middleware.call_sync('vm.setup_libvirt_connection', 10)
+            self.middleware.call_sync('vm.setup_libvirt_connection', timeout)
+
+    def _check_setup_connection(self):
+        self._safely_check_setup_connection()
         self._check_connection_alive()
diff --git a/src/middlewared/middlewared/plugins/vm/vms.py b/src/middlewared/middlewared/plugins/vm/vms.py
index f9213c888009a..4c01abf071108 100644
--- a/src/middlewared/middlewared/plugins/vm/vms.py
+++ b/src/middlewared/middlewared/plugins/vm/vms.py
@@ -101,9 +101,11 @@ def extend_context(self, rows, extra):
         status = {}
         kvm_supported = self._is_kvm_supported()
         if rows and kvm_supported:
-            self._check_setup_connection()
+            self._safely_check_setup_connection(5)
+
+        libvirt_running = self._is_connection_alive()
         for row in rows:
-            status[row['id']] = self.status_impl(row) if kvm_supported else get_default_status()
+            status[row['id']] = self.status_impl(row) if libvirt_running else get_default_status()
 
         return {
             'status': status,