From fd5089b471e12d327d7ab12c3de6069d9e1a9573 Mon Sep 17 00:00:00 2001
From: Vlad Brezae <brezaevlad@gmail.com>
Date: Fri, 31 Mar 2017 00:01:16 +0300
Subject: [PATCH] [sgen] Add debug option for valloc limit

This allows us to stress test OOM inside sgen (and potentially other areas).
---
 mono/sgen/sgen-gc.c              | 11 +++++++++++
 mono/utils/mono-mmap-internals.h |  6 ++++++
 mono/utils/mono-mmap-windows.c   |  6 ++++++
 mono/utils/mono-mmap.c           | 24 +++++++++++++++++++++++-
 mono/utils/mono-mmap.h           |  2 +-
 5 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/mono/sgen/sgen-gc.c b/mono/sgen/sgen-gc.c
index fb6082ab6ab..f6aad7eeb37 100644
--- a/mono/sgen/sgen-gc.c
+++ b/mono/sgen/sgen-gc.c
@@ -193,6 +193,7 @@
 #include "mono/utils/hazard-pointer.h"
 
 #include <mono/utils/memcheck.h>
+#include <mono/utils/mono-mmap-internals.h>
 
 #undef pthread_create
 #undef pthread_join
@@ -3215,6 +3216,15 @@ sgen_gc_init (void)
 			} else if (!strcmp (opt, "verify-before-allocs")) {
 				verify_before_allocs = 1;
 				has_per_allocation_action = TRUE;
+			} else if (g_str_has_prefix (opt, "max-valloc-size=")) {
+				size_t max_valloc_size;
+				char *arg = strchr (opt, '=') + 1;
+				if (*opt && mono_gc_parse_environment_string_extract_number (arg, &max_valloc_size)) {
+					mono_valloc_set_limit (max_valloc_size);
+				} else {
+					sgen_env_var_error (MONO_GC_DEBUG_NAME, NULL, "`max-valloc-size` must be an integer.");
+				}
+				continue;
 			} else if (g_str_has_prefix (opt, "verify-before-allocs=")) {
 				char *arg = strchr (opt, '=') + 1;
 				verify_before_allocs = atoi (arg);
@@ -3295,6 +3305,7 @@ sgen_gc_init (void)
 				fprintf (stderr, "Valid <option>s are:\n");
 				fprintf (stderr, "  collect-before-allocs[=<n>]\n");
 				fprintf (stderr, "  verify-before-allocs[=<n>]\n");
+				fprintf (stderr, "  max-valloc-size=N (where N is an integer, possibly with a k, m or a g suffix)\n");
 				fprintf (stderr, "  check-remset-consistency\n");
 				fprintf (stderr, "  check-mark-bits\n");
 				fprintf (stderr, "  check-nursery-pinned\n");
diff --git a/mono/utils/mono-mmap-internals.h b/mono/utils/mono-mmap-internals.h
index 8fa286337f5..8831c8bdadf 100644
--- a/mono/utils/mono-mmap-internals.h
+++ b/mono/utils/mono-mmap-internals.h
@@ -21,6 +21,12 @@ aligned_address (char *mem, size_t size, size_t alignment);
 void
 account_mem (MonoMemAccountType type, ssize_t size);
 
+gboolean
+mono_valloc_can_alloc (size_t size);
+
+void
+mono_valloc_set_limit (size_t size);
+
 int
 mono_pages_not_faulted (void *addr, size_t length);
 
diff --git a/mono/utils/mono-mmap-windows.c b/mono/utils/mono-mmap-windows.c
index 0f482d1a090..0ae90eee3b9 100644
--- a/mono/utils/mono-mmap-windows.c
+++ b/mono/utils/mono-mmap-windows.c
@@ -71,6 +71,9 @@ mono_valloc (void *addr, size_t length, int flags, MonoMemAccountType type)
 	int prot = mono_mmap_win_prot_from_flags (flags);
 	/* translate the flags */
 
+	if (!mono_valloc_can_alloc (length))
+		return NULL;
+
 	ptr = VirtualAlloc (addr, length, mflags, prot);
 
 	account_mem (type, (ssize_t)length);
@@ -88,6 +91,9 @@ mono_valloc_aligned (size_t length, size_t alignment, int flags, MonoMemAccountT
 	if (!mem)
 		return NULL;
 
+	if (!mono_valloc_can_alloc (length))
+		return NULL;
+
 	aligned = aligned_address (mem, length, alignment);
 
 	aligned = VirtualAlloc (aligned, length, MEM_COMMIT, prot);
diff --git a/mono/utils/mono-mmap.c b/mono/utils/mono-mmap.c
index ccfcb0b8ac0..d579d5a46a5 100644
--- a/mono/utils/mono-mmap.c
+++ b/mono/utils/mono-mmap.c
@@ -81,18 +81,36 @@ aligned_address (char *mem, size_t size, size_t alignment)
 	return aligned;
 }
 
-static volatile size_t allocation_count [MONO_MEM_ACCOUNT_MAX];
+static size_t allocation_count [MONO_MEM_ACCOUNT_MAX];
+static size_t total_allocation_count;
+static size_t alloc_limit;
 
 void
 account_mem (MonoMemAccountType type, ssize_t size)
 {
 #if SIZEOF_VOID_P == 4
 	InterlockedAdd ((volatile gint32*)&allocation_count [type], (gint32)size);
+	InterlockedAdd ((volatile gint32*)&total_allocation_count, (gint32)size);
 #else
 	InterlockedAdd64 ((volatile gint64*)&allocation_count [type], (gint64)size);
+	InterlockedAdd64 ((volatile gint64*)&total_allocation_count, (gint64)size);
 #endif
 }
 
+void
+mono_valloc_set_limit (size_t size)
+{
+	alloc_limit = size;
+}
+
+gboolean
+mono_valloc_can_alloc (size_t size)
+{
+	if (alloc_limit)
+		return (total_allocation_count + size) < alloc_limit;
+	return TRUE;
+}
+
 const char*
 mono_mem_account_type_name (MonoMemAccountType type)
 {
@@ -194,6 +212,10 @@ mono_valloc (void *addr, size_t length, int flags, MonoMemAccountType type)
 	void *ptr;
 	int mflags = 0;
 	int prot = prot_from_flags (flags);
+
+	if (!mono_valloc_can_alloc (length))
+		return NULL;
+
 	/* translate the flags */
 	if (flags & MONO_MMAP_FIXED)
 		mflags |= MAP_FIXED;
diff --git a/mono/utils/mono-mmap.h b/mono/utils/mono-mmap.h
index 72cf853389a..2e2c7cdf036 100644
--- a/mono/utils/mono-mmap.h
+++ b/mono/utils/mono-mmap.h
@@ -83,6 +83,6 @@ typedef void *(*mono_file_map_alloc_fn)   (size_t length);
 typedef void  (*mono_file_map_release_fn) (void *addr);
 
 MONO_API void mono_file_map_set_allocator (mono_file_map_alloc_fn alloc, mono_file_map_release_fn release);
-				  
+
 #endif /* __MONO_UTILS_MMAP_H__ */
 
-- 
2.25.1