Merge pull request #53 from ohadravid/feature/do-not-uninit-com

Stop calling `CoUninitialize`, make `COMLibrary: Copy + !Send`

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wmi"
-version = "0.10.0"
+version = "0.11.0"
 authors = ["Ohad Ravid <[email protected]>"]
 edition = "2018"
 license = "MIT OR Apache-2.0"

README.md

@@ -8,7 +8,7 @@ WMI (Windows Management Instrumentation) crate for rust.
 ```toml
 # Cargo.toml
 [dependencies]
-wmi = "0.10"
+wmi = "0.11"
 ```
 
 
@@ -60,7 +60,7 @@ If you prefer to use the `time` crate instead of the default `chrono`, include `
 
 ```toml
 [dependencies]
-wmi-rs = { version = "0.10", default-features = false, features = ["time"] }
+wmi-rs = { version = "0.11", default-features = false, features = ["time"] }
 ```
 
 and use the `WMIOffsetDateTime` wrapper instead of the the `WMIDateTime` wrapper.

src/connection.rs

@@ -1,9 +1,9 @@
 use crate::utils::{check_hres, WMIError};
 use crate::BStr;
 use log::debug;
+use std::marker::PhantomData;
 use std::ptr;
 use std::ptr::NonNull;
-use std::rc::Rc;
 use winapi::{
     shared::{
         ntdef::NULL,
@@ -14,30 +14,31 @@ use winapi::{
         wtypesbase::CLSCTX_INPROC_SERVER,
     },
     um::{
-        combaseapi::{
-            CoCreateInstance, CoInitializeEx, CoInitializeSecurity, CoSetProxyBlanket,
-            CoUninitialize,
-        },
+        combaseapi::{CoCreateInstance, CoInitializeEx, CoInitializeSecurity, CoSetProxyBlanket},
         objbase::COINIT_MULTITHREADED,
         objidl::EOAC_NONE,
         wbemcli::{CLSID_WbemLocator, IID_IWbemLocator, IWbemLocator, IWbemServices},
     },
 };
 
-pub struct COMLibrary {}
+/// A marker to indicate that the current thread was `CoInitialize`d.
+/// It can be freely copied within the same thread.
+#[derive(Clone, Copy)]
+pub struct COMLibrary {
+    // Force the type to be `!Send`, as each thread must be initialized separately.
+    _phantom: PhantomData<*mut ()>,
+}
 
 /// Initialize COM.
 ///
-/// COM will be `CoUninitialize`d after this object is dropped.
+/// `CoUninitialize` will NOT be called when dropped.
+/// See: https://github.com/microsoft/windows-rs/issues/1169#issuecomment-926877227
 ///
 impl COMLibrary {
     /// `CoInitialize`s the COM library for use by the calling thread.
     ///
     pub fn new() -> Result<Self, WMIError> {
-        unsafe { check_hres(CoInitializeEx(ptr::null_mut(), COINIT_MULTITHREADED))? }
-
-        let instance = Self {};
-
+        let instance = Self::without_security()?;
         instance.init_security()?;
 
         Ok(instance)
@@ -48,11 +49,38 @@ impl COMLibrary {
     pub fn without_security() -> Result<Self, WMIError> {
         unsafe { check_hres(CoInitializeEx(ptr::null_mut(), COINIT_MULTITHREADED))? }
 
-        let instance = Self {};
+        let instance = Self {
+            _phantom: PhantomData,
+        };
 
         Ok(instance)
     }
 
+    /// Assumes that COM was already initialized for this thread.
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe as it is the caller's responsibility to ensure that COM is initialized
+    /// and will not be uninitialized while any instance of object is in scope.
+    ///
+    /// ```edition2018
+    /// # fn main() -> Result<(), wmi::WMIError> {
+    /// # use wmi::*;
+    /// # use serde::Deserialize;
+    /// # let _actual_com = COMLibrary::new()?;
+    /// let initialized_com = unsafe { COMLibrary::assume_initialized() };
+    ///
+    /// // Later, in the same thread.
+    /// let wmi_con = WMIConnection::with_namespace_path("ROOT\\CIMV2", initialized_com)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub unsafe fn assume_initialized() -> Self {
+        Self {
+            _phantom: PhantomData,
+        }
+    }
+
     fn init_security(&self) -> Result<(), WMIError> {
         unsafe {
             check_hres(CoInitializeSecurity(
@@ -72,14 +100,14 @@ impl COMLibrary {
     }
 }
 
-impl Drop for COMLibrary {
-    fn drop(&mut self) {
-        unsafe { CoUninitialize() };
-    }
-}
+/// ```compile_fail
+/// let com = COMLibrary::new().unwrap();
+/// _test_com_lib_not_send(com);
+/// ```
+fn _test_com_lib_not_send(_s: impl Send) {}
 
 pub struct WMIConnection {
-    _com_con: Option<Rc<COMLibrary>>,
+    _com_con: COMLibrary,
     p_loc: Option<NonNull<IWbemLocator>>,
     p_svc: Option<NonNull<IWbemServices>>,
 }
@@ -100,7 +128,7 @@ impl WMIConnection {
     }
 
     /// Creates a connection with a default `CIMV2` namespace path.
-    pub fn new(com_lib: Rc<COMLibrary>) -> Result<Self, WMIError> {
+    pub fn new(com_lib: COMLibrary) -> Result<Self, WMIError> {
         Self::with_namespace_path("ROOT\\CIMV2", com_lib)
     }
 
@@ -116,10 +144,10 @@ impl WMIConnection {
     /// ```
     pub fn with_namespace_path(
         namespace_path: &str,
-        com_lib: Rc<COMLibrary>,
+        com_lib: COMLibrary,
     ) -> Result<Self, WMIError> {
         let mut instance = Self {
-            _com_con: Some(com_lib),
+            _com_con: com_lib,
             p_loc: None,
             p_svc: None,
         };
@@ -129,35 +157,6 @@ impl WMIConnection {
         Ok(instance)
     }
 
-    /// Like `with_namespace_path`, but assumes that COM is managed externally.
-    ///
-    /// # Safety
-    ///
-    /// This function is unsafe as it is the caller's responsibility to ensure that COM is initialized and will not be uninitialized before the connection object is dropped.
-    ///
-    /// ```edition2018
-    /// # fn main() -> Result<(), wmi::WMIError> {
-    /// # use wmi::*;
-    /// # use serde::Deserialize;
-    /// let _initialized_com = COMLibrary::new()?;
-    ///
-    /// // Later, in the same thread.
-    /// let wmi_con = unsafe { WMIConnection::with_initialized_com(Some("ROOT\\CIMV2"))? };
-    /// # Ok(())
-    /// # }
-    /// ```
-    pub unsafe fn with_initialized_com(namespace_path: Option<&str>) -> Result<Self, WMIError> {
-        let mut instance = Self {
-            _com_con: None,
-            p_loc: None,
-            p_svc: None,
-        };
-
-        instance.create_and_set_proxy(namespace_path)?;
-
-        Ok(instance)
-    }
-
     pub fn svc(&self) -> *mut IWbemServices {
         self.p_svc.unwrap().as_ptr()
     }

src/tests.rs

@@ -3,17 +3,15 @@ use crate::WMIConnection;
 
 pub mod fixtures {
     use super::*;
-    use lazy_static::lazy_static;
-
-    // This way we only setup COM security once during tests.
-    // We can't use `std::sync::Once` because we have to keep the `COM_LIB` object alive for the
-    // entire run.
-    lazy_static! {
-        static ref COM_LIB: COMLibrary = COMLibrary::new().unwrap();
+    // This way we only setup COM security once per thread during tests.
+    thread_local! {
+        static COM_LIB: COMLibrary = COMLibrary::without_security().unwrap();
     }
 
     pub fn wmi_con() -> WMIConnection {
-        let wmi_con = WMIConnection::new(COMLibrary::without_security().unwrap().into()).unwrap();
+        let com_lib = COM_LIB.with(|com| *com);
+
+        let wmi_con = WMIConnection::new(com_lib).unwrap();
 
         wmi_con
     }