[NTOS:MM] Add private pages to process working sets

ntoskrnl/include/internal/mm.h

@@ -1674,6 +1674,21 @@ VOID
 NTAPI
 MiInitializeWorkingSetList(_Inout_ PMMSUPPORT WorkingSet);
 
+_Requires_exclusive_lock_held_(Vm->WorkingSetMutex)
+VOID
+NTAPI
+MiInsertInWorkingSetList(
+    _Inout_ PMMSUPPORT Vm,
+    _In_ PVOID Address,
+    _In_ ULONG Protection);
+
+_Requires_exclusive_lock_held_(Vm->WorkingSetMutex)
+VOID
+NTAPI
+MiRemoveFromWorkingSetList(
+    _Inout_ PMMSUPPORT Vm,
+    _In_ PVOID Address);
+
 #ifdef __cplusplus
 } // extern "C"
 

ntoskrnl/mm/ARM3/pagfault.c

@@ -697,16 +697,6 @@ MiResolveDemandZeroFault(IN PVOID Address,
     /* Increment demand zero faults */
     KeGetCurrentPrcb()->MmDemandZeroCount++;
 
-    /* Do we have the lock? */
-    if (HaveLock)
-    {
-        /* Release it */
-        MiReleasePfnLock(OldIrql);
-
-        /* Update performance counters */
-        if (Process > HYDRA_PROCESS) Process->NumberOfPrivatePages++;
-    }
-
     /* Zero the page if need be */
     if (NeedZero) MiZeroPfn(PageFrameNumber);
 
@@ -743,6 +733,19 @@ MiResolveDemandZeroFault(IN PVOID Address,
         /* Windows does these sanity checks */
         ASSERT(Pfn1->u1.Event == 0);
         ASSERT(Pfn1->u3.e1.PrototypePte == 0);
+
+        /* Release it */
+        MiReleasePfnLock(OldIrql);
+
+        /* Update performance counters */
+        if (Process > HYDRA_PROCESS) Process->NumberOfPrivatePages++;
+    }
+
+    /* Add the page to our working set, if it's not a proto PTE */
+    if ((Process > HYDRA_PROCESS) && (PointerPte == MiAddressToPte(Address)))
+    {
+        /* FIXME: Also support session VM scenario */
+        MiInsertInWorkingSetList(&Process->Vm, Address, Protection);
     }
 
     //
@@ -962,6 +965,9 @@ MiResolvePageFileFault(_In_ BOOLEAN StoreInstruction,
         KeSetEvent(Pfn1->u1.Event, IO_NO_INCREMENT, FALSE);
     }
 
+    /* And we can insert this into the working set */
+    MiInsertInWorkingSetList(&CurrentProcess->Vm, FaultingAddress, Protection);
+
     return Status;
 }
 
@@ -979,6 +985,8 @@ MiResolveTransitionFault(IN BOOLEAN StoreInstruction,
     PMMPFN Pfn1;
     MMPTE TempPte;
     PMMPTE PointerToPteForProtoPage;
+    ULONG Protection;
+
     DPRINT("Transition fault on 0x%p with PTE 0x%p in process %s\n",
             FaultingAddress, PointerPte, CurrentProcess->ImageFileName);
 
@@ -1072,8 +1080,9 @@ MiResolveTransitionFault(IN BOOLEAN StoreInstruction,
     ASSERT(PointerPte->u.Hard.Valid == 0);
     ASSERT(PointerPte->u.Trans.Prototype == 0);
     ASSERT(PointerPte->u.Trans.Transition == 1);
+    Protection = TempPte.u.Trans.Protection;
     TempPte.u.Long = (PointerPte->u.Long & ~0xFFF) |
-                     (MmProtectToPteMask[PointerPte->u.Trans.Protection]) |
+                     (MmProtectToPteMask[Protection]) |
                      MiDetermineUserGlobalPteMask(PointerPte);
 
     /* Is the PTE writeable? */
@@ -1093,6 +1102,10 @@ MiResolveTransitionFault(IN BOOLEAN StoreInstruction,
     /* Write the valid PTE */
     MI_WRITE_VALID_PTE(PointerPte, TempPte);
 
+    /* If this was a user fault, add it to the working set */
+    if (CurrentProcess > HYDRA_PROCESS)
+        MiInsertInWorkingSetList(&CurrentProcess->Vm, FaultingAddress, Protection);
+
     /* Return success */
     return STATUS_PAGE_FAULT_TRANSITION;
 }
@@ -1234,6 +1247,9 @@ MiResolveProtoPteFault(IN BOOLEAN StoreInstruction,
         Pfn1 = MI_PFN_ELEMENT(PageFrameIndex);
         MiInitializePfn(PageFrameIndex, PointerPte, TRUE);
 
+        /* The caller expects us to release the PFN lock */
+        MiReleasePfnLock(OldIrql);
+
         /* Fix the protection */
         Protection &= ~MM_WRITECOPY;
         Protection |= MM_READWRITE;
@@ -1251,8 +1267,12 @@ MiResolveProtoPteFault(IN BOOLEAN StoreInstruction,
         /* And finally, write the valid PTE */
         MI_WRITE_VALID_PTE(PointerPte, PteContents);
 
-        /* The caller expects us to release the PFN lock */
+        /* Add the page to our working set */
-        MiReleasePfnLock(OldIrql);
+        if (Process > HYDRA_PROCESS)
+        {
+            /* FIXME: Also support session VM scenario */
+            MiInsertInWorkingSetList(&Process->Vm, Address, Protection);
+        }
         return Status;
     }
 
@@ -2211,6 +2231,7 @@ UserFault:
             {
                 PFN_NUMBER PageFrameIndex, OldPageFrameIndex;
                 PMMPFN Pfn1;
+                ProtectionCode = TempPte.u.Soft.Protection;
 
                 LockIrql = MiAcquirePfnLock();
 
@@ -2234,13 +2255,18 @@ UserFault:
 
                 /* And make a new shiny one with our page */
                 MiInitializePfn(PageFrameIndex, PointerPte, TRUE);
+
+                MiReleasePfnLock(LockIrql);
+
                 TempPte.u.Hard.PageFrameNumber = PageFrameIndex;
                 TempPte.u.Hard.Write = 1;
                 TempPte.u.Hard.CopyOnWrite = 0;
 
                 MI_WRITE_VALID_PTE(PointerPte, TempPte);
 
-                MiReleasePfnLock(LockIrql);
+                /* We can now add it to our working set */
+                MiInsertInWorkingSetList(&CurrentProcess->Vm, Address, ProtectionCode);
+
 
                 /* Return the status */
                 MiUnlockProcessWorkingSet(CurrentProcess, CurrentThread);
@@ -2357,6 +2383,7 @@ UserFault:
                 TempPte.u.Soft.Protection = ProtectionCode;
                 MI_WRITE_INVALID_PTE(PointerPte, TempPte);
             }
+            ProtectionCode = PointerPte->u.Soft.Protection;
 
             /* Lock the PFN database since we're going to grab a page */
             OldIrql = MiAcquirePfnLock();
@@ -2388,14 +2415,14 @@ UserFault:
             /* Initialize the PFN entry now */
             MiInitializePfn(PageFrameIndex, PointerPte, 1);
 
+            /* And be done with the lock */
+            MiReleasePfnLock(OldIrql);
+
             /* Increment the count of pages in the process */
             CurrentProcess->NumberOfPrivatePages++;
 
             /* One more demand-zero fault */
-            KeGetCurrentPrcb()->MmDemandZeroCount++;
+            InterlockedIncrement(&KeGetCurrentPrcb()->MmDemandZeroCount);
-
-            /* And we're done with the lock */
-            MiReleasePfnLock(OldIrql);
 
             /* Fault on user PDE, or fault on user PTE? */
             if (PointerPte <= MiHighestUserPte)
@@ -2423,6 +2450,9 @@ UserFault:
             Pfn1 = MI_PFN_ELEMENT(PageFrameIndex);
             ASSERT(Pfn1->u1.Event == NULL);
 
+            /* We can now insert it into the working set */
+            MiInsertInWorkingSetList(&CurrentProcess->Vm, Address, ProtectionCode);
+
             /* Demand zero */
             ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
             MiUnlockProcessWorkingSet(CurrentProcess, CurrentThread);

ntoskrnl/mm/ARM3/virtual.c

@@ -398,9 +398,6 @@ MiDeletePte(IN PMMPTE PointerPte,
     PFN_NUMBER PageFrameIndex;
     PMMPDE PointerPde;
 
-    /* PFN lock must be held */
-    MI_ASSERT_PFN_LOCK_HELD();
-
     /* Capture the PTE */
     TempPte = *PointerPte;
 
@@ -425,6 +422,8 @@ MiDeletePte(IN PMMPTE PointerPte,
             /* Destroy the PTE */
             MI_ERASE_PTE(PointerPte);
 
+            KIRQL OldIrql = MiAcquirePfnLock();
+
             /* Drop the reference on the page table. */
             MiDecrementShareCount(MiGetPfnEntry(Pfn1->u4.PteFrame), Pfn1->u4.PteFrame);
 
@@ -444,6 +443,8 @@ MiDeletePte(IN PMMPTE PointerPte,
                 MI_SET_PFN_DELETED(Pfn1);
                 MiDecrementReferenceCount(Pfn1, PageFrameIndex);
             }
+
+            MiReleasePfnLock(OldIrql);
             return;
         }
     }
@@ -474,6 +475,8 @@ MiDeletePte(IN PMMPTE PointerPte,
 #if (_MI_PAGING_LEVELS == 2)
         }
 #endif
+
+        KIRQL OldIrql = MiAcquirePfnLock();
         /* Drop the share count on the page table */
         PointerPde = MiPteToPde(PointerPte);
         MiDecrementShareCount(MiGetPfnEntry(PointerPde->u.Hard.PageFrameNumber),
@@ -481,6 +484,7 @@ MiDeletePte(IN PMMPTE PointerPte,
 
         /* Drop the share count */
         MiDecrementShareCount(Pfn1, PageFrameIndex);
+        MiReleasePfnLock(OldIrql);
 
         /* Either a fork, or this is the shared user data page */
         if ((PointerPte <= MiHighestUserPte) && (PrototypePte != Pfn1->PteAddress))
@@ -503,6 +507,9 @@ MiDeletePte(IN PMMPTE PointerPte,
     }
     else
     {
+        /* Remove this address from the WS list */
+        MiRemoveFromWorkingSetList(&CurrentProcess->Vm, VirtualAddress);
+
         /* Make sure the saved PTE address is valid */
         if ((PMMPTE)((ULONG_PTR)Pfn1->PteAddress & ~0x1) != PointerPte)
         {
@@ -517,6 +524,7 @@ MiDeletePte(IN PMMPTE PointerPte,
         /* Erase the PTE */
         MI_ERASE_PTE(PointerPte);
 
+        KIRQL OldIrql = MiAcquirePfnLock();
         /* There should only be 1 shared reference count */
         ASSERT(Pfn1->u2.ShareCount == 1);
 
@@ -526,6 +534,7 @@ MiDeletePte(IN PMMPTE PointerPte,
         /* Mark the PFN for deletion and dereference what should be the last ref */
         MI_SET_PFN_DELETED(Pfn1);
         MiDecrementShareCount(Pfn1, PageFrameIndex);
+        MiReleasePfnLock(OldIrql);
 
         /* We should eventually do this */
         //CurrentProcess->NumberOfPrivatePages--;
@@ -2339,24 +2348,23 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
                 if ((NewAccessProtection & PAGE_NOACCESS) ||
                     (NewAccessProtection & PAGE_GUARD))
                 {
-                    KIRQL OldIrql = MiAcquirePfnLock();
+                    /* Remove this from the working set */
+                    MiRemoveFromWorkingSetList(AddressSpace, MiPteToAddress(PointerPte));
 
                     /* Mark the PTE as transition and change its protection */
                     PteContents.u.Hard.Valid = 0;
                     PteContents.u.Soft.Transition = 1;
                     PteContents.u.Trans.Protection = ProtectionMask;
                     /* Decrease PFN share count and write the PTE */
+                    KIRQL OldIrql = MiAcquirePfnLock();
                     MiDecrementShareCount(Pfn1, PFN_FROM_PTE(&PteContents));
-                    // FIXME: remove the page from the WS
+                    MiReleasePfnLock(OldIrql);
                     MI_WRITE_INVALID_PTE(PointerPte, PteContents);
 #ifdef CONFIG_SMP
                     // FIXME: Should invalidate entry in every CPU TLB
                     ASSERT(FALSE);
 #endif
                     KeInvalidateTlbEntry(MiPteToAddress(PointerPte));
-
-                    /* We are done for this PTE */
-                    MiReleasePfnLock(OldIrql);
                 }
                 else
                 {
@@ -2485,41 +2493,44 @@ MiMakePdeExistAndMakeValid(IN PMMPDE PointerPde,
 
 VOID
 NTAPI
-MiProcessValidPteList(IN PMMPTE *ValidPteList,
+MiProcessValidPteList(
-                      IN ULONG Count)
+    _Inout_ PMMSUPPORT Vm,
+    _Inout_ PMMPTE *ValidPteList,
+    _In_ ULONG Count)
 {
-    KIRQL OldIrql;
     ULONG i;
-    MMPTE TempPte;
-    PFN_NUMBER PageFrameIndex;
-    PMMPFN Pfn1, Pfn2;
-
     //
-    // Acquire the PFN lock and loop all the PTEs in the list
+    // Loop all the PTEs in the list
     //
-    OldIrql = MiAcquirePfnLock();
     for (i = 0; i != Count; i++)
     {
         //
         // The PTE must currently be valid
         //
-        TempPte = *ValidPteList[i];
+        MMPTE TempPte = *ValidPteList[i];
         ASSERT(TempPte.u.Hard.Valid == 1);
 
+        //
+        // We can now remove this addres from the working set
+        //
+        MiRemoveFromWorkingSetList(Vm, MiPteToAddress(ValidPteList[i]));
+
         //
         // Get the PFN entry for the page itself, and then for its page table
         //
-        PageFrameIndex = PFN_FROM_PTE(&TempPte);
+        PFN_NUMBER PageFrameIndex = PFN_FROM_PTE(&TempPte);
-        Pfn1 = MiGetPfnEntry(PageFrameIndex);
+        PMMPFN Pfn1 = MiGetPfnEntry(PageFrameIndex);
-        Pfn2 = MiGetPfnEntry(Pfn1->u4.PteFrame);
+        PMMPFN Pfn2 = MiGetPfnEntry(Pfn1->u4.PteFrame);
 
         //
         // Decrement the share count on the page table, and then on the page
         // itself
         //
+        KIRQL OldIrql = MiAcquirePfnLock();
         MiDecrementShareCount(Pfn2, Pfn1->u4.PteFrame);
         MI_SET_PFN_DELETED(Pfn1);
         MiDecrementShareCount(Pfn1, PageFrameIndex);
+        MiReleasePfnLock(OldIrql);
 
         //
         // Make the page decommitted
@@ -2532,7 +2543,6 @@ MiProcessValidPteList(IN PMMPTE *ValidPteList,
     // and then release the PFN lock
     //
     KeFlushCurrentTb();
-    MiReleasePfnLock(OldIrql);
 }
 
 ULONG
@@ -2547,7 +2557,6 @@ MiDecommitPages(IN PVOID StartingAddress,
     ULONG CommitReduction = 0;
     PMMPTE ValidPteList[256];
     ULONG PteCount = 0;
-    PMMPFN Pfn1;
     MMPTE PteContents;
     PETHREAD CurrentThread = PsGetCurrentThread();
 
@@ -2579,10 +2588,10 @@ MiDecommitPages(IN PVOID StartingAddress,
             // such, and does not flush the entire TLB all the time, but right
             // now we have bigger problems to worry about than TLB flushing.
             //
-            PointerPde = MiAddressToPde(StartingAddress);
+            PointerPde = MiPteToPde(PointerPte);
             if (PteCount)
             {
-                MiProcessValidPteList(ValidPteList, PteCount);
+                MiProcessValidPteList(&Process->Vm, ValidPteList, PteCount);
                 PteCount = 0;
             }
 
@@ -2617,21 +2626,13 @@ MiDecommitPages(IN PVOID StartingAddress,
                 //Process->NumberOfPrivatePages--;
                 if (PteContents.u.Hard.Valid)
                 {
-                    //
-                    // It's valid. At this point make sure that it is not a ROS
-                    // PFN. Also, we don't support ProtoPTEs in this code path.
-                    //
-                    Pfn1 = MiGetPfnEntry(PteContents.u.Hard.PageFrameNumber);
-                    ASSERT(MI_IS_ROS_PFN(Pfn1) == FALSE);
-                    ASSERT(Pfn1->u3.e1.PrototypePte == FALSE);
-
                     //
                     // Flush any pending PTEs that we had not yet flushed, if our
                     // list has gotten too big, then add this PTE to the flush list.
                     //
                     if (PteCount == 256)
                     {
-                        MiProcessValidPteList(ValidPteList, PteCount);
+                        MiProcessValidPteList(&Process->Vm, ValidPteList, PteCount);
                         PteCount = 0;
                     }
                     ValidPteList[PteCount++] = PointerPte;
@@ -2661,7 +2662,7 @@ MiDecommitPages(IN PVOID StartingAddress,
             // This used to be a zero PTE and it no longer is, so we must add a
             // reference to the pagetable.
             //
-            MiIncrementPageTableReferences(StartingAddress);
+            MiIncrementPageTableReferences(MiPteToAddress(PointerPte));
 
             //
             // Next, we account for decommitted PTEs and make the PTE as such
@@ -2671,17 +2672,16 @@ MiDecommitPages(IN PVOID StartingAddress,
         }
 
         //
-        // Move to the next PTE and the next address
+        // Move to the next PTE
         //
         PointerPte++;
-        StartingAddress = (PVOID)((ULONG_PTR)StartingAddress + PAGE_SIZE);
     }
 
     //
     // Flush any dangling PTEs from the loop in the last page table, and then
     // release the working set and return the commit reduction accounting.
     //
-    if (PteCount) MiProcessValidPteList(ValidPteList, PteCount);
+    if (PteCount) MiProcessValidPteList(&Process->Vm, ValidPteList, PteCount);
     MiUnlockProcessWorkingSetUnsafe(Process, CurrentThread);
     return CommitReduction;
 }

ntoskrnl/mm/marea.c

@@ -572,15 +572,14 @@ MmDeleteProcessAddressSpace(PEPROCESS Process)
 
 #if (_MI_PAGING_LEVELS == 2)
     {
-        KIRQL OldIrql;
         PVOID Address;
         PMMPDE pointerPde;
+        KAPC_STATE ApcState;
 
         /* Attach to Process */
-        KeAttachProcess(&Process->Pcb);
+        KeStackAttachProcess(&Process->Pcb, &ApcState);
 
-        /* Acquire PFN lock */
+        MiLockProcessWorkingSet(Process, PsGetCurrentThread());
-        OldIrql = MiAcquirePfnLock();
 
         for (Address = MI_LOWEST_VAD_ADDRESS;
              Address < MM_HIGHEST_VAD_ADDRESS;
@@ -604,11 +603,10 @@ MmDeleteProcessAddressSpace(PEPROCESS Process)
             ASSERT(pointerPde->u.Hard.Valid == 0);
         }
 
-        /* Release lock */
+        MiUnlockProcessWorkingSet(Process, PsGetCurrentThread());
-        MiReleasePfnLock(OldIrql);
 
         /* Detach */
-        KeDetachProcess();
+        KeUnstackDetachProcess(&ApcState);
     }
 #endif