libwnbd: add WnbdPollDiskNumber function

After mapping a WNBD disk, we're informimg the Storport driver that the bus changed, expecting it to rescan the bus and then expose the disk. This happens asynchronously and usually takes a few milliseconds, however it can take a few seconds under huge load. WNBD clients need to wait for the disk to become available before being able to use it. For convenience, we'll add a function that performs the necessary polling, waiting for a disk number to be assigned. It optionally opens the disk, ensuring that it's ready. Signed-off-by: Lucian Petrut <[email protected]>
cloudbase · Aug 24, 2023 · b5cffd5 · b5cffd5
1 parent 7d87bdb
commit b5cffd5
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 62 deletions.
diff --git a/include/wnbd.h b/include/wnbd.h
@@ -186,6 +186,16 @@ DWORD WnbdSetDiskSize(
 // Cleanup the PWNBD_DISK structure. This should be called after stopping
 // the IO dispatchers.
 VOID WnbdClose(PWNBD_DISK Disk);
+// Waits for the disk to become available and returns the associated disk
+// number. Returns ERROR_TIMEOUT if the timeout is exceeded and zero if the
+// operation succeeded.
+DWORD WnbdPollDiskNumber(
+    const char* InstanceName,
+    BOOLEAN ExpectMapped,
+    BOOLEAN TryOpen,
+    DWORD TimeoutMs,
+    DWORD RetryInterval,
+    PDWORD DiskNumber);
 
 DWORD WnbdList(
     PWNBD_CONNECTION_LIST ConnectionList,

diff --git a/libwnbd/libwnbd.cpp b/libwnbd/libwnbd.cpp
@@ -658,6 +658,80 @@ void WnbdClose(PWNBD_DISK Disk)
     free(Disk);
 }
 
+DWORD WnbdPollDiskNumber(
+    const char* InstanceName,
+    BOOLEAN ExpectMapped,
+    BOOLEAN TryOpen,
+    DWORD TimeoutMs,
+    DWORD RetryInterval,
+    PDWORD DiskNumber)
+{
+    LARGE_INTEGER StartTime, CurrTime, ElapsedMs, CounterFreq;
+    QueryPerformanceFrequency(&CounterFreq);
+    QueryPerformanceCounter(&StartTime);
+    ElapsedMs = { 0 };
+
+    do {
+        QueryPerformanceCounter(&CurrTime);
+        ElapsedMs.QuadPart = CurrTime.QuadPart - StartTime.QuadPart;
+        ElapsedMs.QuadPart *= 1000;
+        ElapsedMs.QuadPart /= CounterFreq.QuadPart;
+
+        WNBD_CONNECTION_INFO ConnectionInfo = {0};
+        DWORD Status = WnbdShow(InstanceName, &ConnectionInfo);
+        if (Status) {
+            if (ExpectMapped ||
+                    (Status != ERROR_NO_SUCH_DEVICE &&
+                     Status != ERROR_FILE_NOT_FOUND)) {
+                LogError("Couldn't retrieve WNBD disk info. Error: %d", Status);
+                return Status;
+            }
+
+            // The disk isn't available yet.
+            if (TimeoutMs > ElapsedMs.QuadPart) {
+                Sleep(RetryInterval);
+            }
+            continue;
+        }
+
+        if (ConnectionInfo.DiskNumber != -1) {
+            std::string DiskPath = "\\\\.\\PhysicalDrive" + std::to_string(
+                ConnectionInfo.DiskNumber);
+
+            HANDLE DiskHandle = CreateFileA(
+                DiskPath.c_str(),
+                GENERIC_READ,
+                FILE_SHARE_READ | FILE_SHARE_WRITE,
+                NULL,
+                OPEN_EXISTING,
+                NULL,
+                NULL);
+            if (INVALID_HANDLE_VALUE != DiskHandle) {
+                // Disk available
+                CloseHandle(DiskHandle);
+
+                *DiskNumber = ConnectionInfo.DiskNumber;
+                return 0;
+            }
+
+            DWORD Status = GetLastError();
+            if (Status != ERROR_NO_SUCH_DEVICE &&
+                    Status != ERROR_FILE_NOT_FOUND) {
+                LogError("Unable to open WNBD disk, error: %d", Status);
+                return Status;
+            }
+        }
+
+        // Disk not available yet
+        if (TimeoutMs > ElapsedMs.QuadPart) {
+            Sleep(RetryInterval);
+        }
+    } while (TimeoutMs > ElapsedMs.QuadPart);
+
+    LogError("Timed out waiting for WNBD disk to become available.");
+    return WAIT_TIMEOUT;
+}
+
 VOID WnbdSignalStopped(PWNBD_DISK Disk)
 {
     LogDebug("Marking device as stopped.");

diff --git a/libwnbd/libwnbd.def b/libwnbd/libwnbd.def
@@ -7,6 +7,7 @@ EXPORTS
     WnbdRemoveEx
     WnbdSetDiskSize
     WnbdClose
+    WnbdPollDiskNumber
     WnbdList
     WnbdShow
     WnbdGetUserspaceStats

diff --git a/tests/libwnbd_tests/utils.cpp b/tests/libwnbd_tests/utils.cpp
@@ -49,70 +49,22 @@ std::string GetDiskPath(const char* InstanceName, bool ExpectMapped)
 {
     DWORD TimeoutMs = 10 * 1000;
     DWORD RetryInterval = 500;
-    LARGE_INTEGER StartTime, CurrTime, ElapsedMs, CounterFreq;
-    QueryPerformanceFrequency(&CounterFreq);
-    QueryPerformanceCounter(&StartTime);
-    ElapsedMs = { 0 };
 
-    do {
-        QueryPerformanceCounter(&CurrTime);
-        ElapsedMs.QuadPart = CurrTime.QuadPart - StartTime.QuadPart;
-        ElapsedMs.QuadPart *= 1000;
-        ElapsedMs.QuadPart /= CounterFreq.QuadPart;
-
-        WNBD_CONNECTION_INFO ConnectionInfo = {0};
-        NTSTATUS Status = WnbdShow(InstanceName, &ConnectionInfo);
-        if (Status) {
-            if (ExpectMapped ||
-                    (Status != ERROR_NO_SUCH_DEVICE &&
-                     Status != ERROR_FILE_NOT_FOUND)) {
-                std::string Msg = "WnbdShow failed, error: " +
-                    std::to_string(Status);
-                throw std::runtime_error(Msg);
-            }
-
-            // The disk isn't available yet.
-            if (TimeoutMs > ElapsedMs.QuadPart) {
-                Sleep(RetryInterval);
-            }
-            continue;
-        }
-
-        if (ConnectionInfo.DiskNumber != -1) {
-            std::string DiskPath = "\\\\.\\PhysicalDrive" + std::to_string(
-                ConnectionInfo.DiskNumber);
-
-            HANDLE DiskHandle = CreateFileA(
-                DiskPath.c_str(),
-                GENERIC_READ,
-                FILE_SHARE_READ | FILE_SHARE_WRITE,
-                NULL,
-                OPEN_EXISTING,
-                NULL,
-                NULL);
-            if (INVALID_HANDLE_VALUE != DiskHandle) {
-                // Disk available
-                CloseHandle(DiskHandle);
-                return DiskPath;
-            }
-
-            DWORD Status = GetLastError();
-            if (Status != ERROR_NO_SUCH_DEVICE &&
-                    Status != ERROR_FILE_NOT_FOUND) {
-                std::string Msg = "unable to open WNBD disk, error: " +
-                    std::to_string(Status);
-                throw std::runtime_error(Msg);
-            }
-        }
-
-        // Disk not available yet
-        if (TimeoutMs > ElapsedMs.QuadPart) {
-            Sleep(RetryInterval);
-        }
-    } while (TimeoutMs > ElapsedMs.QuadPart);
+    DWORD DiskNumber;
+    DWORD Status = WnbdPollDiskNumber(
+        InstanceName, ExpectMapped,
+        TRUE, // TryOpen
+        TimeoutMs, RetryInterval,
+        &DiskNumber);
+    if (Status) {
+        std::string Msg = "Couldn't retrieve disk number, error: " +
+            std::to_string(Status);
+        throw std::runtime_error(Msg);
+    }
 
-    std::string Msg = "couln't retrieve WNBD disk info, timed out";
-    throw std::runtime_error(Msg);
+    std::string DiskPath =
+        "\\\\.\\PhysicalDrive" + std::to_string(DiskNumber);
+    return DiskPath;
 }
 
 void SetDiskWritable(std::string InstanceName)