sys: wait groups and notify implementation

src/sys/device.c

@@ -415,10 +415,11 @@ static NTSTATUS FspFsvolDeviceInit(PDEVICE_OBJECT DeviceObject)
         return Result;
     FsvolDeviceExtension->InitDoneEa = 1;
 
-    /* initialize the FSRTL Notify mechanism */
+    /* initialize the Volume Notify and FSRTL Notify mechanisms */
     Result = FspNotifyInitializeSync(&FsvolDeviceExtension->NotifySync);
     if (!NT_SUCCESS(Result))
         return Result;
+    FspWgroupInitialize(&FsvolDeviceExtension->VolumeNotifyWgroup);
     InitializeListHead(&FsvolDeviceExtension->NotifyList);
     FsvolDeviceExtension->InitDoneNotify = 1;
 
@@ -477,7 +478,7 @@ static VOID FspFsvolDeviceFini(PDEVICE_OBJECT DeviceObject)
     if (FsvolDeviceExtension->InitDoneStat)
         FspStatisticsDelete(FsvolDeviceExtension->Statistics);
 
-    /* uninitialize the FSRTL Notify mechanism */
+    /* uninitialize the Volume Notify and FSRTL Notify mechanisms */
     if (FsvolDeviceExtension->InitDoneNotify)
     {
         FspNotifyCleanupAll(

src/sys/driver.h

@@ -652,6 +652,28 @@ VOID FspIrpHookReset(PIRP Irp);
 PVOID FspIrpHookContext(PVOID Context);
 NTSTATUS FspIrpHookNext(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);
 
+/* utility: wait groups
+ *
+ * A Wait Group is a synchronization primitive that encapsulates a counter.
+ * A Wait Group is considered signaled when the counter is 0 and non-signaled
+ * when the counter is non-0. (Wait Group functionality is similar to Golang's
+ * sync.WaitGroup.)
+ *
+ * Wait Groups must always be allocated in non-paged storage.
+ */
+typedef struct
+{
+    KEVENT Event;
+    LONG Count;
+    KSPIN_LOCK SpinLock;
+} FSP_WGROUP;
+VOID FspWgroupInitialize(FSP_WGROUP *Wgroup);
+VOID FspWgroupIncrement(FSP_WGROUP *Wgroup);
+VOID FspWgroupDecrement(FSP_WGROUP *Wgroup);
+VOID FspWgroupSignalPermanently(FSP_WGROUP *Wgroup);
+NTSTATUS FspWgroupWait(FSP_WGROUP *Wgroup,
+    KPROCESSOR_MODE WaitMode, BOOLEAN Alertable, PLARGE_INTEGER PTimeout);
+
 /* silos */
 typedef struct
 {
@@ -1125,6 +1147,7 @@ typedef struct
     UINT64 InfoExpirationTime;
     FSP_FSCTL_VOLUME_INFO VolumeInfo;
     LONG VolumeNotifyLock;
+    FSP_WGROUP VolumeNotifyWgroup;
     PNOTIFY_SYNC NotifySync;
     LIST_ENTRY NotifyList;
     FSP_STATISTICS *Statistics;

src/sys/util.c

@@ -131,6 +131,12 @@ NTSTATUS FspIrpHook(PIRP Irp, PIO_COMPLETION_ROUTINE CompletionRoutine, PVOID Ow
 VOID FspIrpHookReset(PIRP Irp);
 PVOID FspIrpHookContext(PVOID Context);
 NTSTATUS FspIrpHookNext(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);
+VOID FspWgroupInitialize(FSP_WGROUP *Wgroup);
+VOID FspWgroupIncrement(FSP_WGROUP *Wgroup);
+VOID FspWgroupDecrement(FSP_WGROUP *Wgroup);
+VOID FspWgroupSignalPermanently(FSP_WGROUP *Wgroup);
+NTSTATUS FspWgroupWait(FSP_WGROUP *Wgroup,
+    KPROCESSOR_MODE WaitMode, BOOLEAN Alertable, PLARGE_INTEGER PTimeout);
 
 #ifdef ALLOC_PRAGMA
 #pragma alloc_text(PAGE, FspIsNtDdiVersionAvailable)
@@ -176,6 +182,11 @@ NTSTATUS FspIrpHookNext(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);
 #pragma alloc_text(PAGE, FspIrpHookReset)
 // !#pragma alloc_text(PAGE, FspIrpHookContext)
 // !#pragma alloc_text(PAGE, FspIrpHookNext)
+// !#pragma alloc_text(PAGE, FspWgroupInitialize)
+// !#pragma alloc_text(PAGE, FspWgroupIncrement)
+// !#pragma alloc_text(PAGE, FspWgroupDecrement)
+// !#pragma alloc_text(PAGE, FspWgroupSignalPermanently)
+// !#pragma alloc_text(PAGE, FspWgroupWait)
 #endif
 
 static const LONG Delays[] =
@@ -1488,3 +1499,56 @@ NTSTATUS FspIrpHookNext(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context)
 
     return Result;
 }
+
+VOID FspWgroupInitialize(FSP_WGROUP *Wgroup)
+{
+    // !PAGED_CODE();
+
+    KeInitializeEvent(&Wgroup->Event, NotificationEvent, TRUE);
+    Wgroup->Count = 0;
+    KeInitializeSpinLock(&Wgroup->SpinLock);
+}
+
+VOID FspWgroupIncrement(FSP_WGROUP *Wgroup)
+{
+    // !PAGED_CODE();
+
+    KIRQL Irql;
+
+    KeAcquireSpinLock(&Wgroup->SpinLock, &Irql);
+    if (0 <= Wgroup->Count && 1 == ++Wgroup->Count)
+        KeClearEvent(&Wgroup->Event);
+    KeReleaseSpinLock(&Wgroup->SpinLock, Irql);
+}
+
+VOID FspWgroupDecrement(FSP_WGROUP *Wgroup)
+{
+    // !PAGED_CODE();
+
+    KIRQL Irql;
+
+    KeAcquireSpinLock(&Wgroup->SpinLock, &Irql);
+    if (0 < Wgroup->Count && 0 == --Wgroup->Count)
+        KeSetEvent(&Wgroup->Event, 1, FALSE);
+    KeReleaseSpinLock(&Wgroup->SpinLock, Irql);
+}
+
+VOID FspWgroupSignalPermanently(FSP_WGROUP *Wgroup)
+{
+    // !PAGED_CODE();
+
+    KIRQL Irql;
+
+    KeAcquireSpinLock(&Wgroup->SpinLock, &Irql);
+    Wgroup->Count = -1;
+    KeSetEvent(&Wgroup->Event, 1, FALSE);
+    KeReleaseSpinLock(&Wgroup->SpinLock, Irql);
+}
+
+NTSTATUS FspWgroupWait(FSP_WGROUP *Wgroup,
+    KPROCESSOR_MODE WaitMode, BOOLEAN Alertable, PLARGE_INTEGER PTimeout)
+{
+    // !PAGED_CODE();
+
+    return KeWaitForSingleObject(&Wgroup->Event, Executive, WaitMode, Alertable, PTimeout);
+}

src/sys/volume.c

@@ -477,6 +477,7 @@ static VOID FspVolumeDeleteNoLock(
     }
 
     /* release the volume notify lock if held (so that any pending rename will abort) */
+    FspWgroupSignalPermanently(&FsvolDeviceExtension->VolumeNotifyWgroup);
     if (1 == InterlockedCompareExchange(&FsvolDeviceExtension->VolumeNotifyLock, 0, 1))
         FspFsvolDeviceFileRenameReleaseOwner(FsvolDeviceObject, FsvolDeviceObject);
 
@@ -1097,6 +1098,7 @@ NTSTATUS FspVolumeNotify(
     ASSERT(0 != IrpSp->FileObject->FsContext2);
 
     PDEVICE_OBJECT FsvolDeviceObject = IrpSp->FileObject->FsContext2;
+    FSP_FSVOL_DEVICE_EXTENSION *FsvolDeviceExtension = FspFsvolDeviceExtension(FsvolDeviceObject);
     PVOID InputBuffer = IrpSp->Parameters.FileSystemControl.Type3InputBuffer;
     ULONG InputBufferLength = IrpSp->Parameters.FileSystemControl.InputBufferLength;
     FSP_VOLUME_NOTIFY_WORK_ITEM *NotifyWorkItem = 0;
@@ -1131,6 +1133,8 @@ NTSTATUS FspVolumeNotify(
 
     ExInitializeWorkItem(&NotifyWorkItem->WorkItem, FspVolumeNotifyWork, NotifyWorkItem);
     NotifyWorkItem->FsvolDeviceObject = FsvolDeviceObject;
+
+    FspWgroupIncrement(&FsvolDeviceExtension->VolumeNotifyWgroup);
     ExQueueWorkItem(&NotifyWorkItem->WorkItem, DelayedWorkQueue);
 
     return STATUS_SUCCESS;
@@ -1196,6 +1200,7 @@ static VOID FspVolumeNotifyWork(PVOID NotifyWorkItem0)
     ULONG NotifyInfoSize;
     UNICODE_STRING FileName = { 0 }, StreamPart = { 0 }, AbsFileName = { 0 }, FullFileName = { 0 };
     ULONG StreamType = FspFileNameStreamTypeNone;
+    BOOLEAN Unlock = FALSE;
     NTSTATUS Result;
 
     /* iterate over notify information and invalidate/notify each file */
@@ -1206,9 +1211,7 @@ static VOID FspVolumeNotifyWork(PVOID NotifyWorkItem0)
 
         if (sizeof(FSP_FSCTL_NOTIFY_INFO) > NotifyInfoSize)
         {
-            /* !!!: we should wait until all pending notify work is done! */
+            Unlock = TRUE;
-            if (1 == InterlockedCompareExchange(&FsvolDeviceExtension->VolumeNotifyLock, 0, 1))
-                FspFsvolDeviceFileRenameReleaseOwner(FsvolDeviceObject, FsvolDeviceObject);
             break;
         }
 
@@ -1265,6 +1268,14 @@ static VOID FspVolumeNotifyWork(PVOID NotifyWorkItem0)
 
     FspFree(NotifyWorkItem);
 
+    FspWgroupDecrement(&FsvolDeviceExtension->VolumeNotifyWgroup);
+    if (Unlock)
+    {
+        FspWgroupWait(&FsvolDeviceExtension->VolumeNotifyWgroup, KernelMode, FALSE, 0);
+        if (1 == InterlockedCompareExchange(&FsvolDeviceExtension->VolumeNotifyLock, 0, 1))
+            FspFsvolDeviceFileRenameReleaseOwner(FsvolDeviceObject, FsvolDeviceObject);
+    }
+
     FspDeviceDereference(FsvolDeviceObject);
 
     FsRtlExitFileSystem();