feat: keep NVML and NvAPI behind a refcount instead of global init

lact-daemon/src/server/gpu_controller.rs

@@ -17,7 +17,8 @@ use tokio::join;
 pub const VENDOR_AMD: &str = "1002";
 pub const VENDOR_NVIDIA: &str = "10DE";
 
-use crate::server::handler::{AMD_DRM, INTEL_DRM, NVML};
+use crate::config::Config;
+use crate::server::handler::{AMD_DRM, INTEL_DRM};
 use crate::server::opencl::get_opencl_info;
 use crate::server::vulkan::get_vulkan_info;
 use amdgpu_sysfs::gpu_handle::power_profile_mode::PowerProfileModesTable;
@@ -29,16 +30,12 @@ use lact_schema::{
 };
 use std::io;
 #[cfg(feature = "nvidia")]
-use std::sync::Arc;
+use std::sync::LazyLock;
+use std::sync::Mutex;
 use std::{collections::HashMap, fs, path::PathBuf, rc::Rc};
 use tokio::{sync::Notify, task::JoinHandle};
 use tracing::{error, warn};
 
-#[cfg(feature = "nvidia")]
-pub use nvidia::nvapi::NvApi;
-#[cfg(feature = "nvidia")]
-use nvml_wrapper::Nvml;
-
 pub type DynGpuController = Box<dyn GpuController>;
 type FanControlHandle = (Rc<Notify>, JoinHandle<()>);
 
@@ -166,11 +163,6 @@ pub struct PciSlotInfo {
     pub func: u16,
 }
 
-#[cfg(feature = "nvidia")]
-pub type NvidiaLibs = (Arc<Nvml>, Arc<Option<NvApi>>);
-#[cfg(not(feature = "nvidia"))]
-pub type NvidiaLibs = ();
-
 pub(crate) fn build_controller_info(
     path: PathBuf,
     pci_db: &pciid_parser::Database,
@@ -246,9 +238,59 @@ pub(crate) fn build_controller_info(
 
 pub(crate) fn init_controller(
     common: CommonControllerInfo,
+    config: &Config,
 ) -> anyhow::Result<Box<dyn GpuController>> {
+    static INIT_MUTEX: Mutex<()> = Mutex::new(());
+
     #[cfg(not(feature = "nvidia"))]
-    let _ = NVML;
+    let _ = config;
+
+    // SAFETY: We use global LazyLock to make sure it's safe.
+    // Loading of shared libaries is unsafe
+    // https://docs.rs/libloading/0.8.8/libloading/struct.Library.html#method.new
+    let _guard = INIT_MUTEX.lock().unwrap();
+
+    #[cfg(feature = "nvidia")]
+    #[allow(unused_unsafe)]
+    let nvml = LazyLock::new(|| unsafe {
+        use nvml_wrapper::Nvml;
+        use tracing::info;
+
+        Nvml::init()
+            .map(|nvml| {
+                info!(
+                    "Nvidia management library {} loaded",
+                    nvml.sys_nvml_version()
+                        .unwrap_or_else(|err| err.to_string())
+                );
+                Rc::new(nvml)
+            })
+            .inspect_err(|err| {
+                error!("could not load Nvidia management library: {err}");
+            })
+            .ok()
+    });
+
+    #[cfg(feature = "nvidia")]
+    let nvapi = LazyLock::new(|| {
+        use crate::server::gpu_controller::nvidia::nvapi::NvApi;
+        use tracing::info;
+
+        if config.daemon.disable_nvapi == Some(true) {
+            info!("NvAPI support is disabled");
+            None
+        } else {
+            NvApi::new()
+                .map(|nvapi| {
+                    info!("NvAPI library loaded");
+                    Rc::new(nvapi)
+                })
+                .inspect_err(|err| {
+                    warn!("could not load NvAPI library: {err:#}");
+                })
+                .ok()
+        }
+    });
 
     match common.driver.as_str() {
         "amdgpu" | "radeon" => {
@@ -269,8 +311,8 @@ pub(crate) fn init_controller(
         }
         #[cfg(feature = "nvidia")]
         "nvidia" => {
-            if let Some((nvml, nvapi)) = NVML.as_ref() {
-                match NvidiaGpuController::new(common.clone(), nvml, nvapi.as_ref().as_ref()) {
+            if let Some(nvml) = nvml.clone() {
+                match NvidiaGpuController::new(common.clone(), nvml, nvapi.clone()) {
                     Ok(controller) => {
                         return Ok(Box::new(controller));
                     }

lact-daemon/src/server/gpu_controller/nvidia.rs

@@ -5,7 +5,6 @@ use super::{CommonControllerInfo, FanControlHandle, GpuController};
 use crate::{
     bindings::nvidia::NvPhysicalGpuHandle,
     server::gpu_controller::{
-        NvApi,
         common::{fan_control::FanCurveExt, resolve_process_name},
         nvidia::nvapi::{CLOCK_CLIENT_CLK_VF_POINT_TYPE_PROG, ClockClientClkVfPointInfoV1},
     },
@@ -26,6 +25,7 @@ use lact_schema::{
     ProcessType, ProcessUtilizationType, TemperatureEntry, VoltageStats, VramStats,
     config::{CurvePoint, FanControlSettings, FanCurve, GpuConfig},
 };
+use nvapi::NvApi;
 use nvml_wrapper::{
     Device, Nvml,
     bitmasks::device::{PowerMizerModes, ThrottleReasons},
@@ -52,12 +52,12 @@ const SUPPORTED_UTIL_TYPES: &[ProcessUtilizationType] = &[
 ];
 
 pub struct NvidiaGpuController {
-    nvml: &'static Nvml,
+    nvml: Rc<Nvml>,
     common: CommonControllerInfo,
     fan_control_handle: RefCell<Option<FanControlHandle>>,
     initial_target_temp: Option<u32>,
 
-    nvapi: Option<(&'static NvApi, NvPhysicalGpuHandle)>,
+    nvapi: Option<(Rc<NvApi>, NvPhysicalGpuHandle)>,
     driver_handle: Option<DriverHandle>,
     nvapi_thermals_mask: Option<i32>,
 
@@ -75,8 +75,8 @@ pub struct NvidiaGpuController {
 impl NvidiaGpuController {
     pub fn new(
         common: CommonControllerInfo,
-        nvml: &'static Nvml,
-        nvapi: Option<&'static NvApi>,
+        nvml: Rc<Nvml>,
+        nvapi: Option<Rc<NvApi>>,
     ) -> anyhow::Result<Self> {
         let device = nvml
             .device_by_pci_bus_id(common.pci_slot_name.as_str())
@@ -205,7 +205,7 @@ impl NvidiaGpuController {
         let notify = Rc::new(Notify::new());
         let task_notify = notify.clone();
 
-        let nvml = self.nvml;
+        let nvml = self.nvml.clone();
         let pci_slot_id = self.common.pci_slot_name.clone();
         debug!("spawning new fan control task");
 

lact-daemon/src/server/gpu_controller/nvidia/nvapi.rs

@@ -16,6 +16,7 @@ use std::{
     mem::{self, transmute},
     ptr,
 };
+use tracing::error;
 
 const LIBARY_NAME: &str = "libnvidia-api.so.1";
 const QUERY_INTERFACE_FN: &[u8] = b"nvapi_QueryInterface\0";
@@ -371,7 +372,10 @@ impl Drop for NvApi {
         unsafe {
             let unload = self.query_interface(QUERY_NVAPI_UNLOAD).unwrap();
             let unload: unsafe extern "C" fn() -> NvAPI_Status = transmute(unload);
-            unload();
+            let status = unload();
+            if let Err(err) = self.handle_status(status) {
+                error!("could not unload NvAPI: {err}");
+            }
         }
     }
 }

lact-daemon/src/server/handler.rs

@@ -5,6 +5,7 @@ use super::{
 };
 #[cfg(feature = "display-info")]
 use crate::server::display;
+use crate::system::run_command;
 use crate::{
     bindings::intel::IntelDrm,
     config::Config,
@@ -15,7 +16,6 @@ use crate::{
         system::DAEMON_VERSION,
     },
 };
-use crate::{server::gpu_controller::NvidiaLibs, system::run_command};
 use amdgpu_sysfs::gpu_handle::{
     PerformanceLevel, PowerLevelKind, power_profile_mode::PowerProfileModesTable,
 };
@@ -33,12 +33,8 @@ use lact_schema::{
 use libdrm_amdgpu_sys::LibDrmAmdgpu;
 use libflate::gzip;
 use nix::libc;
-#[cfg(all(not(test), feature = "nvidia"))]
-use nvml_wrapper::Nvml;
 use pciid_parser::Database;
 use serde_json::json;
-#[cfg(all(not(test), feature = "nvidia"))]
-use std::sync::Arc;
 use std::{
     cell::{Cell, RefCell},
     collections::{BTreeMap, HashMap},
@@ -132,7 +128,7 @@ impl<'a> Handler {
         // For such scenarios there is a retry logic when no GPUs were found,
         // or if some of the PCI devices don't have a drm entry yet.
         for i in 1..=CONTROLLERS_LOAD_RETRY_ATTEMPTS {
-            controllers = load_controllers(base_path, pci_db, &[])?;
+            controllers = load_controllers(base_path, pci_db, &[], &config)?;
 
             let mut should_retry = false;
 
@@ -245,7 +241,7 @@ impl<'a> Handler {
 
         let base_path = drm_base_path();
         let pci_db = read_pci_db();
-        match load_controllers(&base_path, &pci_db, &detached_ids) {
+        match load_controllers(&base_path, &pci_db, &detached_ids, &config) {
             Ok(new_controllers) => {
                 info!(
                     "GPU list reloaded with {} devices, reapplying configuration",
@@ -1342,48 +1338,6 @@ pub(crate) fn read_pci_db() -> Database {
     })
 }
 
-#[cfg(any(test, not(feature = "nvidia")))]
-pub(crate) static NVML: LazyLock<Option<NvidiaLibs>> = LazyLock::new(|| None);
-
-#[cfg(all(not(test), feature = "nvidia"))]
-// SAFETY: We use global LazyLock to make sure it's safe.
-// Loading of shared libaries is unsafe
-// https://docs.rs/libloading/0.8.8/libloading/struct.Library.html#method.new
-#[allow(unused_unsafe)]
-pub(crate) static NVML: LazyLock<Option<NvidiaLibs>> =
-    LazyLock::new(|| match unsafe { Nvml::init() } {
-        Ok(nvml) => {
-            use crate::server::gpu_controller::NvApi;
-
-            // The config has to be re-read here, because a LazyLock cannot capture external variables into the init closure
-            let disable_nvapi = Config::load()
-                .ok()
-                .flatten()
-                .and_then(|config| config.daemon.disable_nvapi);
-
-            info!("Nvidia management library loaded");
-            let nvapi = if disable_nvapi == Some(true) {
-                info!("NvAPI support is disabled");
-                None
-            } else {
-                NvApi::new()
-                    .inspect(|_| {
-                        info!("NvAPI library loaded");
-                    })
-                    .inspect_err(|err| {
-                        warn!("could not load NvAPI library: {err:#}");
-                    })
-                    .ok()
-            };
-
-            Some((Arc::new(nvml), Arc::new(nvapi)))
-        }
-        Err(err) => {
-            error!("could not load Nvidia management library: {err}");
-            None
-        }
-    });
-
 pub(crate) static AMD_DRM: LazyLock<Option<LibDrmAmdgpu>> = LazyLock::new(|| {
     // SAFETY: We use global LazyLock to make sure it's safe.
     #[allow(unused_unsafe)]
@@ -1421,6 +1375,7 @@ fn load_controllers(
     base_path: &Path,
     pci_db: &Database,
     detached_ids: &[String],
+    config: &Config,
 ) -> anyhow::Result<BTreeMap<String, DynGpuController>> {
     let mut controllers = BTreeMap::new();
 
@@ -1455,7 +1410,7 @@ fn load_controllers(
                 continue;
             }
 
-            match init_controller(info) {
+            match init_controller(info, config) {
                 Ok(controller) => {
                     let info = controller.controller_info();