Initial Support for Quota Pool Allocations.

Implemented ExAllocatePoolWithQuota, ExAllocatePoolWithQuotaTag, PsChargePoolQuota, PsChargeProcessNonPagedPoolQuota, PsChargeProcessPagedPoolQuota, PsChargeProcessPoolQuota.
Also implemented a basic ExAllocatePoolWithTagPriority, but needs Special Pool support for Driver Verifier functionality. The allocation works however.

svn path=/trunk/; revision=10429

reactos/ntoskrnl/include/internal/mm.h

@@ -13,6 +13,8 @@
 extern ULONG MiFreeSwapPages;
 extern ULONG MiUsedSwapPages;
 extern ULONG MmPagedPoolSize;
+extern ULONG MmTotalPagedPoolQuota;
+extern ULONG MmTotalNonPagedPoolQuota;
 
 struct _EPROCESS;
 
@@ -136,6 +138,17 @@ typedef struct _SECTION_OBJECT
 
 typedef struct _SECTION_OBJECT *PSECTION_OBJECT;
 
+typedef struct _EPROCESS_QUOTA_BLOCK {
+KSPIN_LOCK      QuotaLock;
+ULONG           ReferenceCount;
+ULONG           QuotaPeakPoolUsage[2];
+ULONG           QuotaPoolUsage[2];
+ULONG           QuotaPoolLimit[2];
+ULONG           PeakPagefileUsage;
+ULONG           PagefileUsage;
+ULONG           PagefileLimit;
+} EPROCESS_QUOTA_BLOCK, *PEPROCESS_QUOTA_BLOCK;
+
 #endif /* __USE_W32API */
 
 typedef struct
@@ -349,6 +362,16 @@ VOID MiInitializeNonPagedPool(VOID);
 
 PVOID MmGetMdlPageAddress(PMDL Mdl, PVOID Offset);
 
+/* pool.c *******************************************************************/
+
+BOOLEAN
+STDCALL
+MiRaisePoolQuota(
+    IN POOL_TYPE PoolType,
+    IN ULONG CurrentMaxQuota,
+    OUT PULONG NewMaxQuota
+    );
+
 /* mdl.c *********************************************************************/
 
 VOID MmBuildMdlFromPages(PMDL Mdl, PULONG Pages);

reactos/ntoskrnl/include/internal/pool.h

@@ -15,6 +15,14 @@ VOID STDCALL
 ExFreePagedPool(IN PVOID Block);
 VOID MmInitializePagedPool(VOID);
 
+PVOID
+STDCALL
+MiAllocateSpecialPool  (IN POOL_TYPE PoolType,
+                        IN SIZE_T NumberOfBytes,
+                        IN ULONG Tag,
+                        IN ULONG Underrun
+                        );
+
 extern PVOID MmPagedPoolBase;
 extern ULONG MmPagedPoolSize;
 

reactos/ntoskrnl/include/internal/ps.h

@@ -16,7 +16,7 @@
  *  along with this program; if not, write to the Free Software
  *  Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
  */
-/* $Id: ps.h,v 1.62 2004/07/23 07:44:26 jimtabor Exp $
+/* $Id: ps.h,v 1.63 2004/08/08 20:33:17 ion Exp $
  *
  * FILE:            ntoskrnl/ke/kthread.c
  * PURPOSE:         Process manager definitions
@@ -406,7 +406,7 @@ struct _EPROCESS
   HANDLE                SectionHandle;
   PPEB                  Peb;
   PVOID                 SectionBaseAddress;
-  PVOID                 QuotaBlock;
+  PEPROCESS_QUOTA_BLOCK QuotaBlock;
   NTSTATUS              LastThreadExitStatus;
   PVOID                 WorkingSetWatch;
   HANDLE                InheritedFromUniqueProcessId;

reactos/ntoskrnl/mm/npool.c

@@ -1,4 +1,4 @@
-/* $Id: npool.c,v 1.86 2004/08/01 07:24:58 hbirr Exp $
+/* $Id: npool.c,v 1.87 2004/08/08 20:33:17 ion Exp $
  *
  * COPYRIGHT:    See COPYING in the top level directory
  * PROJECT:      ReactOS kernel
@@ -151,6 +151,9 @@ ULONG EiFreeNonPagedPool = 0;
  */
 ULONG EiUsedNonPagedPool = 0;
 
+/* Total quota for Non Paged Pool */
+ULONG MmTotalNonPagedPoolQuota = 0;
+
 /*
  * Allocate a range of memory in the nonpaged pool
  */
@@ -1832,4 +1835,17 @@ MiInitializeNonPagedPool(VOID)
 #endif
 }
 
+PVOID
+STDCALL
+MiAllocateSpecialPool  (IN POOL_TYPE PoolType,
+                        IN SIZE_T NumberOfBytes,
+                        IN ULONG Tag,
+                        IN ULONG Underrun
+                        )
+{
+    /* FIXME: Special Pools not Supported */
+    DbgPrint("Special Pools not supported\n");
+    return 0;
+}
+
 /* EOF */

reactos/ntoskrnl/mm/pool.c

@@ -1,4 +1,4 @@
-/* $Id: pool.c,v 1.30 2004/07/17 03:03:52 ion Exp $
+/* $Id: pool.c,v 1.31 2004/08/08 20:33:17 ion Exp $
  * 
  * COPYRIGHT:    See COPYING in the top level directory
  * PROJECT:      ReactOS kernel
@@ -13,10 +13,17 @@
 #include <reactos/bugcodes.h>
 #include <internal/ntoskrnl.h>
 #include <internal/pool.h>
+#include <internal/mm.h>
+#include <pseh.h>
 
 #define NDEBUG
 #include <internal/debug.h>
 
+extern ULONG MiNonPagedPoolLength;
+extern ULONG MmTotalPagedPoolQuota;
+extern ULONG MmTotalNonPagedPoolQuota;
+extern MM_STATS MmStats;
+
 /* GLOBALS *****************************************************************/
 
 #define TAG_NONE (ULONG)(('N'<<0) + ('o'<<8) + ('n'<<16) + ('e'<<24))
@@ -147,7 +154,7 @@ ExAllocatePoolWithQuota (POOL_TYPE PoolType, ULONG NumberOfBytes)
 }
 
 /*
- * @unimplemented
+ * @implemented
  */
 PVOID
 STDCALL
@@ -158,28 +165,60 @@ ExAllocatePoolWithTagPriority(
     IN EX_POOL_PRIORITY Priority
     )
 {
-	UNIMPLEMENTED;
+    /* Check if this is one of the "Special" Flags, used by the Verifier */
+    if (Priority & 8) {
+        /* Check if this is a xxSpecialUnderrun */
+        if (Priority & 1) {
+            return MiAllocateSpecialPool(PoolType, NumberOfBytes, Tag, 1);
+        } else { /* xxSpecialOverrun */
+            return MiAllocateSpecialPool(PoolType, NumberOfBytes, Tag, 0);
+        }
+    }
+
+    /* FIXME: Do Ressource Checking Based on Priority and fail if resources too low*/
+
+    /* Do the allocation */
+    return ExAllocatePoolWithTag(PoolType, NumberOfBytes, Tag);
 }
+
 /*
- * @unimplemented
+ * @implemented
  */
 PVOID STDCALL
 ExAllocatePoolWithQuotaTag (IN POOL_TYPE PoolType,
                             IN ULONG  NumberOfBytes,
                             IN ULONG  Tag)
 {
-#if 0
+    PVOID Block;
-   PVOID Block;
+    PEPROCESS Process;
-   Block = EiAllocatePool(PoolType,
-                          NumberOfBytes,
-                          Tag,
-                          (PVOID)__builtin_return_address(0));
-   return(Block);
-#else
 
-   UNIMPLEMENTED;
+    /* Allocate the Pool First */
-   return(NULL);
+    Block = EiAllocatePool(PoolType,
-#endif
+                            NumberOfBytes,
+                            Tag,
+                            &ExAllocatePoolWithQuotaTag);
+
+    /* "Quota is not charged to the thread for allocations >= PAGE_SIZE" - OSR Docs */
+    if (!(NumberOfBytes >= PAGE_SIZE)) {
+
+        /* Get the Current Process */
+        Process = PsGetCurrentProcess();
+
+        /* PsChargePoolQuota returns an exception, so this needs SEH */
+        _SEH_FILTER(FreeAndGoOn) {
+            /* Couldn't charge, so free the pool and let the caller SEH manage */
+            ExFreePool(Block);
+            return EXCEPTION_CONTINUE_SEARCH;
+        } _SEH_TRY_FILTER(FreeAndGoOn) {
+            //* FIXME: Is there a way to get the actual Pool size allocated from the pool header? */
+            PsChargePoolQuota(Process, PoolType, NumberOfBytes);
+        } _SEH_HANDLE {
+            /* Quota Exceeded and the caller had no SEH! */
+            KeBugCheck(STATUS_QUOTA_EXCEEDED);
+        } _SEH_END;
+    }
+
+    return Block;
 }
 
 /*
@@ -250,8 +289,42 @@ MmFreeMappingAddress (
 	UNIMPLEMENTED;
 }
 
+BOOLEAN
+STDCALL
+MiRaisePoolQuota(
+    IN POOL_TYPE PoolType,
+    IN ULONG CurrentMaxQuota,
+    OUT PULONG NewMaxQuota
+    )
+{
+    /* Different quota raises depending on the type (64K vs 512K) */
+    if (PoolType == PagedPool) {
+
+        /* Make sure that 4MB is still left */
+        if ((MM_PAGED_POOL_SIZE >> 12) < ((MmPagedPoolSize + 4194304) >> 12)) {
+            return FALSE;
+        }
+
+        /* Increase Paged Pool Quota by 512K */
+        MmTotalPagedPoolQuota += 524288;
+        *NewMaxQuota = CurrentMaxQuota + 524288;
+        return TRUE;
+
+    } else { /* Nonpaged Pool */
+
+        /* Check if we still have 200 pages free*/
+        if (MmStats.NrFreePages < 200) return FALSE;
+
+        /* Check that 4MB is still left */
+        if ((MM_NONPAGED_POOL_SIZE >> 12) < ((MiNonPagedPoolLength + 4194304) >> 12)) {
+            return FALSE;
+        }
+
+        /* Increase Non Paged Pool Quota by 64K */
+        MmTotalNonPagedPoolQuota += 65536;
+        *NewMaxQuota = CurrentMaxQuota + 65536;
+        return TRUE;
+    }
+}
+
 /* EOF */
-
-
-
-

reactos/ntoskrnl/mm/ppool.c

@@ -1,4 +1,4 @@
-/* $Id: ppool.c,v 1.28 2004/04/10 22:35:25 gdalsnes Exp $
+/* $Id: ppool.c,v 1.29 2004/08/08 20:33:17 ion Exp $
  *
  * COPYRIGHT:       See COPYING in the top level directory
  * PROJECT:         ReactOS kernel
@@ -52,6 +52,7 @@ MM_PPOOL_USED_BLOCK_HEADER, *PMM_PPOOL_USED_BLOCK_HEADER;
 
 PVOID MmPagedPoolBase;
 ULONG MmPagedPoolSize;
+ULONG MmTotalPagedPoolQuota = 0;
 static FAST_MUTEX MmPagedPoolLock;
 static PMM_PPOOL_FREE_BLOCK_HEADER MmPagedPoolFirstFreeBlock;
 #if MM_PPOOL_REDZONE_BYTES

reactos/ntoskrnl/ps/process.c

@@ -1,4 +1,4 @@
-/* $Id: process.c,v 1.137 2004/07/22 17:22:38 jimtabor Exp $
+/* $Id: process.c,v 1.138 2004/08/08 20:33:17 ion Exp $
  *
  * COPYRIGHT:         See COPYING in the top level directory
  * PROJECT:           ReactOS kernel
@@ -1991,7 +1991,7 @@ PsSetProcessWin32WindowStation(
 
 /* Pool Quotas */
 /*
- * @unimplemented
+ * @implemented
  */
 VOID
 STDCALL
@@ -2001,39 +2001,47 @@ PsChargePoolQuota(
     IN ULONG_PTR Amount
     )
 {
-	UNIMPLEMENTED;
+    NTSTATUS Status;
+
+    /* Charge the usage */
+    Status = PsChargeProcessPoolQuota(Process, PoolType, Amount);
+
+    /* Raise Exception */
+    if (!NT_SUCCESS(Status)) {
+        ExRaiseStatus(Status);
+    }
 }
 
 /*
- * @unimplemented
+ * @implemented
  */
 NTSTATUS
 STDCALL
 PsChargeProcessNonPagedPoolQuota (
-    	IN PEPROCESS Process,
+    IN PEPROCESS Process,
-    	IN ULONG_PTR Amount
+    IN ULONG_PTR Amount
-	)
+    )
 {
-	UNIMPLEMENTED;	
+    /* Call the general function */
-	return STATUS_NOT_IMPLEMENTED;
+    return PsChargeProcessPoolQuota(Process, NonPagedPool, Amount);
 }
 
 /*
- * @unimplemented
+ * @implemented
  */
 NTSTATUS
 STDCALL
 PsChargeProcessPagedPoolQuota (
-    	IN PEPROCESS Process,
+    IN PEPROCESS Process,
-    	IN ULONG_PTR Amount
+    IN ULONG_PTR Amount
-	)
+    )
 {
-	UNIMPLEMENTED;	
+    /* Call the general function */
-	return STATUS_NOT_IMPLEMENTED;
+    return PsChargeProcessPoolQuota(Process, PagedPool, Amount);
 }
 
 /*
- * @unimplemented
+ * @implemented
  */
 NTSTATUS
 STDCALL
@@ -2043,8 +2051,53 @@ PsChargeProcessPoolQuota(
     IN ULONG_PTR Amount
     )
 {
-	UNIMPLEMENTED;	
+    PEPROCESS_QUOTA_BLOCK QuotaBlock;
-	return STATUS_NOT_IMPLEMENTED;
+    KIRQL OldValue;
+    ULONG NewUsageSize;
+    ULONG NewMaxQuota;
+
+    /* Get current Quota Block */
+    QuotaBlock = Process->QuotaBlock;
+
+    /* Quota Operations are not to be done on the SYSTEM Process */
+    if (Process == PsInitialSystemProcess) return STATUS_SUCCESS;
+
+    /* Acquire Spinlock */
+    KeAcquireSpinLock(&QuotaBlock->QuotaLock, &OldValue);
+
+    /* New Size in use */
+    NewUsageSize = QuotaBlock->QuotaPoolUsage[PoolType] + Amount;
+
+    /* Does this size respect the quota? */
+    if (NewUsageSize > QuotaBlock->QuotaPoolLimit[PoolType]) {
+
+        /* It doesn't, so keep raising the Quota */
+        while (MiRaisePoolQuota(PoolType, QuotaBlock->QuotaPoolLimit[PoolType], &NewMaxQuota)) {
+            /* Save new Maximum Quota */
+            QuotaBlock->QuotaPoolLimit[PoolType] = NewMaxQuota;
+
+            /* See if the new Maximum Quota fulfills our need */
+            if (NewUsageSize <= NewMaxQuota) goto QuotaChanged;
+        }
+
+        KeReleaseSpinLock(&QuotaBlock->QuotaLock, OldValue);
+        return STATUS_QUOTA_EXCEEDED;
+    }
+
+QuotaChanged:
+    /* Save new Usage */
+    QuotaBlock->QuotaPoolUsage[PoolType] = NewUsageSize;
+
+    /* Is this a new peak? */
+    if (NewUsageSize > QuotaBlock->QuotaPeakPoolUsage[PoolType]) {
+        QuotaBlock->QuotaPeakPoolUsage[PoolType] = NewUsageSize;
+    }
+
+    /* Release spinlock */
+    KeReleaseSpinLock(&QuotaBlock->QuotaLock, OldValue);
+
+    /* All went well */
+    return STATUS_SUCCESS;
 }
 
 /*