Refactor Library Load

src/bin/linux/init.c

@@ -39,7 +39,6 @@ Entry(
     void
     )
 {
-    CxPlatSystemLoad();
     MsQuicLibraryLoad();
 }
 
@@ -50,5 +49,4 @@ Exit(
     )
 {
     MsQuicLibraryUnload();
-    CxPlatSystemUnload();
 }

src/bin/winkernel/driver.c

@@ -15,7 +15,6 @@
 #include "driver.c.clog.h"
 #endif
 
-INITCODE
 _IRQL_requires_max_(PASSIVE_LEVEL)
 void
 MsQuicLibraryLoad(
@@ -83,8 +82,11 @@ Return Value:
     WDF_DRIVER_CONFIG Config;
     WDFDRIVER Driver;
 
-    CxPlatSystemLoad(DriverObject, RegistryPath);
-
+    //
+    // We explicitly load the MsQuic library upfront (instead of letting it
+    // delay load) because we need to be able to query performance counters at
+    // any time, even if there hasn't been a call to MsQuicOpen yet.
+    //
     MsQuicLibraryLoad();
 
     //
@@ -120,7 +122,6 @@ Return Value:
 
     if (!NT_SUCCESS(Status)) {
         MsQuicLibraryUnload();
-        CxPlatSystemUnload();
     }
 
     return Status;
@@ -152,5 +153,4 @@ Routine Description:
     PAGED_CODE();
     MsQuicPcwCleanup();
     MsQuicLibraryUnload();
-    CxPlatSystemUnload();
 }

src/bin/winuser/dllmain.c

@@ -42,13 +42,11 @@ DllMain(
 #else
         UNREFERENCED_PARAMETER(Instance);
 #endif
-        CxPlatSystemLoad();
         MsQuicLibraryLoad();
         break;
 
     case DLL_PROCESS_DETACH:
         MsQuicLibraryUnload();
-        CxPlatSystemUnload();
         break;
     }
 

src/bin/winuser_fuzz/dllmain.c

@@ -105,13 +105,11 @@ DllMain(
 #else
         UNREFERENCED_PARAMETER(Instance);
 #endif
-        CxPlatSystemLoad();
         MsQuicLibraryLoad();
         break;
 
     case DLL_PROCESS_DETACH:
         MsQuicLibraryUnload();
-        CxPlatSystemUnload();
         break;
     }
 

src/core/library.c

@@ -31,26 +31,31 @@ QuicLibraryEvaluateSendRetryState(
     );
 
 //
-// Initializes all global variables.
+// Initializes all global variables if not already done.
 //
-INITCODE
 _IRQL_requires_max_(PASSIVE_LEVEL)
 void
 MsQuicLibraryLoad(
     void
     )
 {
-    CxPlatLockInitialize(&MsQuicLib.Lock);
-    CxPlatDispatchLockInitialize(&MsQuicLib.DatapathLock);
-    CxPlatDispatchLockInitialize(&MsQuicLib.StatelessRetryKeysLock);
-    CxPlatListInitializeHead(&MsQuicLib.Registrations);
-    CxPlatListInitializeHead(&MsQuicLib.Bindings);
-    QuicTraceRundownCallback = QuicTraceRundown;
-    MsQuicLib.Loaded = TRUE;
+    if (InterlockedIncrement16(&MsQuicLib.LoadRefCount) == 1) {
+        //
+        // Load the library.
+        //
+        CxPlatSystemLoad();
+        CxPlatLockInitialize(&MsQuicLib.Lock);
+        CxPlatDispatchLockInitialize(&MsQuicLib.DatapathLock);
+        CxPlatDispatchLockInitialize(&MsQuicLib.StatelessRetryKeysLock);
+        CxPlatListInitializeHead(&MsQuicLib.Registrations);
+        CxPlatListInitializeHead(&MsQuicLib.Bindings);
+        QuicTraceRundownCallback = QuicTraceRundown;
+        MsQuicLib.Loaded = TRUE;
+    }
 }
 
 //
-// Uninitializes global variables.
+// Uninitializes global variables if necessary.
 //
 _IRQL_requires_max_(PASSIVE_LEVEL)
 void
@@ -59,12 +64,15 @@ MsQuicLibraryUnload(
     )
 {
     CXPLAT_FRE_ASSERT(MsQuicLib.Loaded);
-    QUIC_LIB_VERIFY(MsQuicLib.RefCount == 0);
-    QUIC_LIB_VERIFY(!MsQuicLib.InUse);
-    MsQuicLib.Loaded = FALSE;
-    CxPlatDispatchLockUninitialize(&MsQuicLib.StatelessRetryKeysLock);
-    CxPlatDispatchLockUninitialize(&MsQuicLib.DatapathLock);
-    CxPlatLockUninitialize(&MsQuicLib.Lock);
+    if (InterlockedDecrement16(&MsQuicLib.LoadRefCount) == 0) {
+        QUIC_LIB_VERIFY(MsQuicLib.OpenRefCount == 0);
+        QUIC_LIB_VERIFY(!MsQuicLib.InUse);
+        MsQuicLib.Loaded = FALSE;
+        CxPlatDispatchLockUninitialize(&MsQuicLib.StatelessRetryKeysLock);
+        CxPlatDispatchLockUninitialize(&MsQuicLib.DatapathLock);
+        CxPlatLockUninitialize(&MsQuicLib.Lock);
+        CxPlatSystemUnload();
+    }
 }
 
 void
@@ -124,7 +132,7 @@ QuicLibrarySumPerfCountersExternal(
 {
     CxPlatLockAcquire(&MsQuicLib.Lock);
 
-    if (MsQuicLib.RefCount == 0) {
+    if (MsQuicLib.OpenRefCount == 0) {
         CxPlatZeroMemory(Buffer, BufferLength);
     } else {
         QuicLibrarySumPerfCounters(Buffer, BufferLength);
@@ -575,10 +583,10 @@ MsQuicAddRef(
     // Increment global ref count, and if this is the first ref, initialize all
     // the global library state.
     //
-    if (++MsQuicLib.RefCount == 1) {
+    if (++MsQuicLib.OpenRefCount == 1) {
         Status = MsQuicLibraryInitialize();
         if (QUIC_FAILED(Status)) {
-            MsQuicLib.RefCount--;
+            MsQuicLib.OpenRefCount--;
             goto Error;
         }
     }
@@ -607,12 +615,12 @@ MsQuicRelease(
     // last ref.
     //
 
-    CXPLAT_FRE_ASSERT(MsQuicLib.RefCount > 0);
+    CXPLAT_FRE_ASSERT(MsQuicLib.OpenRefCount > 0);
     QuicTraceEvent(
         LibraryRelease,
         "[ lib] Release");
 
-    if (--MsQuicLib.RefCount == 0) {
+    if (--MsQuicLib.OpenRefCount == 0) {
         MsQuicLibraryUninitialize();
     }
 
@@ -1290,6 +1298,9 @@ MsQuicOpen(
     )
 {
     QUIC_STATUS Status;
+    BOOLEAN ReleaseRefOnFailure = FALSE;
+
+    MsQuicLibraryLoad();
 
     if (QuicApi == NULL) {
         QuicTraceLogVerbose(
@@ -1307,11 +1318,12 @@ MsQuicOpen(
     if (QUIC_FAILED(Status)) {
         goto Exit;
     }
+    ReleaseRefOnFailure = TRUE;
 
     QUIC_API_TABLE* Api = CXPLAT_ALLOC_NONPAGED(sizeof(QUIC_API_TABLE), QUIC_POOL_API);
     if (Api == NULL) {
         Status = QUIC_STATUS_OUT_OF_MEMORY;
-        goto Error;
+        goto Exit;
     }
 
     Api->SetContext = MsQuicSetContext;
@@ -1353,19 +1365,21 @@ MsQuicOpen(
 
     *QuicApi = Api;
 
-Error:
-
-    if (QUIC_FAILED(Status)) {
-        MsQuicRelease();
-    }
-
 Exit:
 
     QuicTraceLogVerbose(
         LibraryMsQuicOpenExit,
         "[ api] MsQuicOpen, status=0x%x",
         Status);
 
+    if (QUIC_FAILED(Status)) {
+        if (ReleaseRefOnFailure) {
+            MsQuicRelease();
+        }
+
+        MsQuicLibraryUnload();
+    }
+
     return Status;
 }
 
@@ -1396,6 +1410,7 @@ MsQuicClose(
             "[ api] MsQuicClose");
         CXPLAT_FREE(QuicApi, QUIC_POOL_API);
         MsQuicRelease();
+        MsQuicLibraryUnload();
     }
 }
 
@@ -1805,7 +1820,7 @@ QuicTraceRundown(
 
     CxPlatLockAcquire(&MsQuicLib.Lock);
 
-    if (MsQuicLib.RefCount > 0) {
+    if (MsQuicLib.OpenRefCount > 0) {
         QuicTraceEvent(
             LibraryRundown,
             "[ lib] Rundown, PartitionCount=%u DatapathFeatures=%u",

src/core/library.h

@@ -119,9 +119,14 @@ typedef struct QUIC_LIBRARY {
     CXPLAT_DISPATCH_LOCK DatapathLock;
 
     //
-    // Total outstanding references on the library.
+    // Total outstanding references from calls to MsQuicLoadLibrary.
     //
-    uint32_t RefCount;
+    volatile short LoadRefCount;
+
+    //
+    // Total outstanding references from calls to MsQuicOpen.
+    //
+    uint16_t OpenRefCount;
 
     //
     // Number of processors currently being used.

src/inc/quic_platform.h

@@ -178,6 +178,64 @@ DEFINE_ENUM_FLAG_OPERATORS(CXPLAT_THREAD_FLAGS);
 #error "Unsupported Platform"
 #endif
 
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+//
+// Library Initialization
+//
+
+//
+// Called in main, DLLMain or DriverEntry.
+//
+PAGEDX
+_IRQL_requires_max_(PASSIVE_LEVEL)
+void
+CxPlatSystemLoad(
+    void
+    );
+
+//
+// Called in main (exit), DLLMain or DriverUnload.
+//
+PAGEDX
+_IRQL_requires_max_(PASSIVE_LEVEL)
+void
+CxPlatSystemUnload(
+    void
+    );
+
+//
+// Initializes the PAL library. Calls to this and
+// CxPlatformUninitialize must be serialized and cannot overlap.
+//
+PAGEDX
+_IRQL_requires_max_(PASSIVE_LEVEL)
+QUIC_STATUS
+CxPlatInitialize(
+    void
+    );
+
+//
+// Uninitializes the PAL library. Calls to this and
+// CxPlatformInitialize must be serialized and cannot overlap.
+//
+PAGEDX
+_IRQL_requires_max_(PASSIVE_LEVEL)
+void
+CxPlatUninitialize(
+    void
+    );
+
+#if defined(__cplusplus)
+}
+#endif
+
+//
+// List Abstraction
+//
+
 #define QuicListEntryValidate(Entry) \
     CXPLAT_DBG_ASSERT( \
         (((Entry->Flink)->Blink) == Entry) && \