From a61a67961a9018501a009d8588002c8b0a8de631 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Fri, 29 May 2020 21:33:55 -0700
Subject: [PATCH] ircd::db: Lock database cache extents in RAM if possible.
 (closes #144)

---
 ircd/db.cc         | 52 ++++++++++++++++++++++++++++++++++++++++++++--
 ircd/db_database.h |  2 ++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/ircd/db.cc b/ircd/db.cc
index a9936c117..35b987986 100644
--- a/ircd/db.cc
+++ b/ircd/db.cc
@@ -8,6 +8,7 @@
 // copyright notice and this permission notice is present in all copies. The
 // full license for this software is available in the LICENSE file.
 
+#include <RB_INC_SYS_MMAN_H
 #include <RB_INC_JEMALLOC_H
 #include "db.h"
 
@@ -3528,6 +3529,15 @@ ircd::db::database::allocator::ALIGN_DEFAULT
 	#endif
 };
 
+decltype(ircd::db::database::allocator::mlock_limit)
+ircd::db::database::allocator::mlock_limit
+{
+	ircd::allocator::rlimit_memlock()
+};
+
+decltype(ircd::db::database::allocator::mlock_current)
+ircd::db::database::allocator::mlock_current;
+
 /// Handle to a jemalloc arena when non-zero. Used as the base arena for all
 /// cache allocators.
 decltype(ircd::db::database::allocator::cache_arena)
@@ -3610,7 +3620,22 @@ noexcept
 	};
 	#endif
 
-	return their_hooks.alloc(hooks, new_addr, size, alignment, zero, commit, arena_ind);
+	void *const ret
+	{
+		their_hooks.alloc(hooks, new_addr, size, alignment, zero, commit, arena_ind)
+	};
+
+	// This feature is only enabled when RLIMIT_MEMLOCK is unlimited. We don't
+	// want to deal with any limit at all.
+	#if defined(HAVE_MLOCK2) && defined(MLOCK_ONFAULT)
+	if(database::allocator::mlock_limit == -1)
+	{
+		syscall(::mlock2, ret, size, MLOCK_ONFAULT);
+		database::allocator::mlock_current += size;
+	}
+	#endif
+
+	return ret;
 }
 #endif
 
@@ -3638,7 +3663,21 @@ noexcept
 	};
 	#endif
 
-	return their_hooks.dalloc(hooks, ptr, size, committed, arena_ind);
+	const bool ret
+	{
+		their_hooks.dalloc(hooks, ptr, size, committed, arena_ind)
+	};
+
+	#if defined(HAVE_MLOCK2)
+	if(database::allocator::mlock_current && !ret)
+	{
+		syscall(::munlock, ptr, size);
+		assert(database::allocator::mlock_current >= size);
+		database::allocator::mlock_current -= size;
+	}
+	#endif
+
+	return ret;
 }
 #endif
 
@@ -3666,6 +3705,15 @@ noexcept
 	};
 	#endif
 
+	#if defined(HAVE_MLOCK2)
+	if(database::allocator::mlock_current)
+	{
+		syscall(::munlock, ptr, size);
+		assert(database::allocator::mlock_current >= size);
+		database::allocator::mlock_current -= size;
+	}
+	#endif
+
 	return their_hooks.destroy(hooks, ptr, size, committed, arena_ind);
 }
 #endif
diff --git a/ircd/db_database.h b/ircd/db_database.h
index 5dbf38c3e..929f60ba2 100644
--- a/ircd/db_database.h
+++ b/ircd/db_database.h
@@ -296,6 +296,8 @@ struct ircd::db::database::allocator final
 :rocksdb::MemoryAllocator
 {
 	static const size_t ALIGN_DEFAULT;
+	static const size_t mlock_limit;
+	static size_t mlock_current;
 	static unsigned cache_arena;
 
 	database *d {nullptr};