From 13643013514218711ceae436e41459cd69e2fcdd Mon Sep 17 00:00:00 2001
From: scorpion-26 <dev.scorpion26@gmail.com>
Date: Thu, 21 Sep 2023 19:18:06 +0800
Subject: [PATCH] Check nvidia gpu features on init

When a required feature is not supported, gBar would crash inside the query functions without a way to manually bypass NvidiaGPU. Now NvidiaGPU checks whether all features are supported and disables the module if any throw an error.
Fixes https://github.com/scorpion-26/gBar/issues/52
---
 src/NvidiaGPU.h | 63 +++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 51 insertions(+), 12 deletions(-)

diff --git a/src/NvidiaGPU.h b/src/NvidiaGPU.h
index 8952cbe..6ab6c20 100644
--- a/src/NvidiaGPU.h
+++ b/src/NvidiaGPU.h
@@ -7,10 +7,30 @@
 #ifdef WITH_NVIDIA
 namespace NvidiaGPU
 {
+    struct GPUUtilization
+    {
+        uint32_t gpu;
+        uint32_t vram;
+    };
+    struct VRAM
+    {
+        uint64_t totalB;
+        uint64_t freeB;
+        uint64_t usedB;
+    };
+
     static void* nvmldl;
 
     static void* nvmlGPUHandle;
 
+    // Preloaded Query functions 
+    typedef int (*PFN_nvmlDeviceGetUtilizationRates)(void*, struct GPUUtilization*);
+    static PFN_nvmlDeviceGetUtilizationRates nvmlDeviceGetUtilizationRates;
+    typedef int (*PFN_nvmlDeviceGetTemperature)(void*, uint32_t, uint32_t*);
+    static PFN_nvmlDeviceGetTemperature nvmlDeviceGetTemperature;
+    typedef int (*PFN_nvmlDeviceGetMemoryInfo)(void*, struct VRAM*);
+    static PFN_nvmlDeviceGetMemoryInfo nvmlDeviceGetMemoryInfo;
+
     inline void Init()
     {
         if (nvmldl || !RuntimeConfig::Get().hasNvidia)
@@ -40,6 +60,37 @@ namespace NvidiaGPU
         auto nvmlDeviceGetHandle = (PFN_nvmlDeviceGetHandle)dlsym(nvmldl, "nvmlDeviceGetHandleByIndex");
         res = nvmlDeviceGetHandle(0, &nvmlGPUHandle);
         ASSERT(res == 0, "Failed getting device (Error: " << res << ")!");
+
+        // Dynamically load functions
+        nvmlDeviceGetUtilizationRates = (PFN_nvmlDeviceGetUtilizationRates)dlsym(nvmldl, "nvmlDeviceGetUtilizationRates");
+        nvmlDeviceGetTemperature = (PFN_nvmlDeviceGetTemperature)dlsym(nvmldl, "nvmlDeviceGetTemperature");
+        nvmlDeviceGetMemoryInfo = (PFN_nvmlDeviceGetMemoryInfo)dlsym(nvmldl, "nvmlDeviceGetMemoryInfo");
+        
+        // Check if all information is available
+        GPUUtilization util;
+        int res = nvmlDeviceGetUtilizationRates(nvmlGPUHandle, &util);
+        if (res != 0)
+        {
+            LOG("Failed querying utilization rates (Error: " << res << "), disabling Nvidia GPU!");
+            RuntimeConfig::Get().hasNvidia = false;
+            return;
+        }
+        uint32_t temp;
+        res = nvmlDeviceGetTemperature(nvmlGPUHandle, 0, &temp);
+        if (res != 0)
+        {
+            LOG("Failed querying temperature (Error: " << res << "), disabling Nvidia GPU!");
+            RuntimeConfig::Get().hasNvidia = false;
+            return;
+        }
+        VRAM mem;
+        res = nvmlDeviceGetMemoryInfo(nvmlGPUHandle, &mem);
+        if (res != 0)
+        {
+            LOG("Failed querying VRAM (Error: " << res << "), disabling Nvidia GPU!");
+            RuntimeConfig::Get().hasNvidia = false;
+            return;
+        }
     }
 
     inline void Shutdown()
@@ -48,18 +99,6 @@ namespace NvidiaGPU
             dlclose(nvmldl);
     }
 
-    struct GPUUtilization
-    {
-        uint32_t gpu;
-        uint32_t vram;
-    };
-    struct VRAM
-    {
-        uint64_t totalB;
-        uint64_t freeB;
-        uint64_t usedB;
-    };
-
     inline GPUUtilization GetUtilization()
     {
         if (!RuntimeConfig::Get().hasNvidia)