From d4aa2b7ce4bf1d23843e3398f9c4bca813964559 Mon Sep 17 00:00:00 2001
From: fufesou <13586388+fufesou@users.noreply.github.com>
Date: Tue, 12 Nov 2024 20:12:05 +0800
Subject: [PATCH] fix: virtual display, headless, wait plug in done (#9895)

Signed-off-by: fufesou <linlong1266@gmail.com>
---
 src/virtual_display_manager.rs | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/virtual_display_manager.rs b/src/virtual_display_manager.rs
index 41e5b3fc8..b2791767e 100644
--- a/src/virtual_display_manager.rs
+++ b/src/virtual_display_manager.rs
@@ -529,12 +529,25 @@ pub mod amyuni_idd {
     }
 
     #[inline]
-    fn plug_monitor_(add: bool) -> Result<(), win_device::DeviceError> {
+    fn plug_monitor_(
+        add: bool,
+        wait_timeout: Option<Duration>,
+    ) -> Result<(), win_device::DeviceError> {
         let cmd = if add { 0x10 } else { 0x00 };
         let cmd = [cmd, 0x00, 0x00, 0x00];
+        let now = Instant::now();
+        let c1 = get_monitor_count();
         unsafe {
             win_device::device_io_control(&INTERFACE_GUID, PLUG_MONITOR_IO_CONTROL_CDOE, &cmd, 0)?;
         }
+        if let Some(wait_timeout) = wait_timeout {
+            while now.elapsed() < wait_timeout {
+                if get_monitor_count() != c1 {
+                    break;
+                }
+                std::thread::sleep(Duration::from_millis(30));
+            }
+        }
         // No need to consider concurrency here.
         if add {
             // If the monitor is plugged in, increase the count.
@@ -552,12 +565,16 @@ pub mod amyuni_idd {
 
     // `std::thread::sleep()` with a timeout is acceptable here.
     // Because user can wait for a while to plug in a monitor.
-    fn plug_in_monitor_(add: bool, is_driver_async_installed: bool) -> ResultType<()> {
+    fn plug_in_monitor_(
+        add: bool,
+        is_driver_async_installed: bool,
+        wait_timeout: Option<Duration>,
+    ) -> ResultType<()> {
         let timeout = Duration::from_secs(3);
         let now = Instant::now();
         let reg_connectivity_old = reg_display_settings::read_reg_connectivity();
         loop {
-            match plug_monitor_(add) {
+            match plug_monitor_(add, wait_timeout) {
                 Ok(_) => {
                     break;
                 }
@@ -622,7 +639,7 @@ pub mod amyuni_idd {
             bail!("Failed to install driver.");
         }
 
-        plug_in_monitor_(true, is_async)
+        plug_in_monitor_(true, is_async, Some(Duration::from_millis(3_000)))
     }
 
     pub fn plug_in_monitor() -> ResultType<()> {
@@ -636,7 +653,7 @@ pub mod amyuni_idd {
             bail!("There are already {VIRTUAL_DISPLAY_MAX_COUNT} monitors plugged in.");
         }
 
-        plug_in_monitor_(true, is_async)
+        plug_in_monitor_(true, is_async, None)
     }
 
     // `index` the display index to plug out. -1 means plug out all.
@@ -700,7 +717,7 @@ pub mod amyuni_idd {
         }
 
         for _i in 0..to_plug_out_count {
-            let _ = plug_monitor_(false);
+            let _ = plug_monitor_(false, None);
         }
         Ok(())
     }