ircd::allocator: Add allocator::scope device.

include/ircd/allocator.h

@@ -20,6 +20,7 @@
 namespace ircd::allocator
 {
 	struct state;
+	struct scope;
 	struct profile;
 	template<class T = char> struct dynamic;
 	template<class T = char, size_t = 512> struct fixed;
@@ -60,6 +61,41 @@ struct ircd::allocator::profile
 	static thread_local profile this_thread;
 };
 
+/// This object hooks and replaces global ::malloc() and family for the
+/// lifetime of the instance, redirecting those calls to the user's provided
+/// callbacks. This functionality may not be available on all platforms so it
+/// cannot be soley relied upon in a production release. It may still be used
+/// optimistically as an optimization in production.
+///
+/// This device is useful to control dynamic memory at level where specific
+/// class allocators are too fine-grained and replacing global new is too
+/// coarse (and far too intrusive to the whole process). Instead this works
+/// on the stack for everything further up the stack.
+///
+/// This class is friendly. It takes control from any other previous instance
+/// of allocator::scope and then restores their control after this goes out of
+/// scope. Once all instances of allocator::scope go out of scope, the previous
+/// global __malloc_hook is reinstalled.
+///
+struct ircd::allocator::scope
+{
+	using alloc_closure = std::function<void *(const size_t &)>;
+	using realloc_closure = std::function<void *(void *const &ptr, const size_t &)>;
+	using free_closure = std::function<void (void *const &ptr)>;
+
+	static scope *current;
+	scope *theirs;
+	alloc_closure user_alloc;
+	realloc_closure user_realloc;
+	free_closure user_free;
+
+  public:
+	scope(alloc_closure = {}, realloc_closure = {}, free_closure = {});
+	scope(const scope &) = delete;
+	scope(scope &&) = delete;
+	~scope();
+};
+
 /// Internal state structure for some of these tools. This is a very small and
 /// simple interface to a bit array representing the availability of an element
 /// in a pool of elements. The actual array of the proper number of bits must

ircd/allocator.cc

@@ -134,6 +134,242 @@ const
 	return this->next(n) < size;
 }
 
+//
+// allocator::scope
+//
+
+namespace ircd::allocator
+{
+	void *(*their_malloc_hook)(size_t, const void *);
+	static void *malloc_hook(size_t, const void *);
+	static void install_malloc_hook();
+	static void uninstall_malloc_hook();
+
+	void *(*their_realloc_hook)(void *, size_t, const void *);
+	static void *realloc_hook(void *, size_t, const void *);
+	static void install_realloc_hook();
+	static void uninstall_realloc_hook();
+
+	void (*their_free_hook)(void *, const void *);
+	static void free_hook(void *, const void *);
+	static void install_free_hook();
+	static void uninstall_free_hook();
+}
+
+decltype(ircd::allocator::scope::current)
+ircd::allocator::scope::current;
+
+ircd::allocator::scope::scope(alloc_closure ac,
+                              realloc_closure rc,
+                              free_closure fc)
+:theirs
+{
+	current
+}
+,user_alloc
+{
+	std::move(ac)
+}
+,user_realloc
+{
+	std::move(rc)
+}
+,user_free
+{
+	std::move(fc)
+}
+{
+	// If an allocator::scope instance already exists somewhere
+	// up the stack, *current will already be set. We only install
+	// our global hook handlers at the first instance ctor and
+	// uninstall it after that first instance dtors.
+	if(!current)
+	{
+		install_malloc_hook();
+		install_realloc_hook();
+		install_free_hook();
+	}
+
+	current = this;
+}
+
+ircd::allocator::scope::~scope()
+noexcept
+{
+	assert(current);
+	current = theirs;
+
+	// Reinstall the pre-existing hooks after our last scope instance
+	// has destructed (the first to have constructed). We know this when
+	// current becomes null.
+	if(!current)
+	{
+		uninstall_malloc_hook();
+		uninstall_realloc_hook();
+		uninstall_free_hook();
+	}
+}
+
+void
+ircd::allocator::free_hook(void *const ptr,
+                           const void *const caller)
+{
+	// Once we've hooked we put back their hook before calling the user
+	// so they can passthru to the function without hooking themselves.
+	uninstall_free_hook();
+	const unwind rehook_ours
+	{
+		install_free_hook
+	};
+
+	assert(scope::current);
+	if(scope::current->user_free)
+		scope::current->user_free(ptr);
+	else
+		::free(ptr);
+}
+
+void *
+ircd::allocator::realloc_hook(void *const ptr,
+                              size_t size,
+                              const void *const caller)
+{
+	// Once we've hooked we put back their hook before calling the user
+	// so they can passthru to the function without hooking themselves.
+	uninstall_realloc_hook();
+	const unwind rehook_ours
+	{
+		install_realloc_hook
+	};
+
+	assert(scope::current);
+	return scope::current->user_realloc?
+		scope::current->user_realloc(ptr, size):
+		::realloc(ptr, size);
+}
+
+void *
+ircd::allocator::malloc_hook(size_t size,
+                             const void *const caller)
+{
+	// Once we've hooked we put back their hook before calling the user
+	// so they can passthru to the function without hooking themselves.
+	uninstall_malloc_hook();
+	const unwind rehook_ours
+	{
+		install_malloc_hook
+	};
+
+	assert(scope::current);
+	return scope::current->user_alloc?
+		scope::current->user_alloc(size):
+		::malloc(size);
+}
+
+#if defined(__GNU_LIBRARY__) && defined(HAVE_MALLOC_H)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+void
+ircd::allocator::install_malloc_hook()
+{
+	assert(!their_malloc_hook);
+	their_malloc_hook = __malloc_hook;
+	__malloc_hook = malloc_hook;
+}
+#pragma GCC diagnostic pop
+#else
+void
+ircd::allocator::install_malloc_hook()
+{
+}
+#endif
+
+#if defined(__GNU_LIBRARY__) && defined(HAVE_MALLOC_H)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+void
+ircd::allocator::uninstall_malloc_hook()
+{
+	assert(their_malloc_hook);
+	__malloc_hook = their_malloc_hook;
+}
+#pragma GCC diagnostic pop
+#else
+void
+ircd::allocator::uninstall_malloc_hook()
+{
+}
+#endif
+
+#if defined(__GNU_LIBRARY__) && defined(HAVE_MALLOC_H)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+void
+ircd::allocator::install_realloc_hook()
+{
+	assert(!their_realloc_hook);
+	their_realloc_hook = __realloc_hook;
+	__realloc_hook = realloc_hook;
+}
+#pragma GCC diagnostic pop
+#else
+void
+ircd::allocator::install_realloc_hook()
+{
+}
+#endif
+
+#if defined(__GNU_LIBRARY__) && defined(HAVE_MALLOC_H)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+void
+ircd::allocator::uninstall_realloc_hook()
+{
+	assert(their_realloc_hook);
+	__realloc_hook = their_realloc_hook;
+}
+#else
+void
+ircd::allocator::uninstall_realloc_hook()
+{
+}
+#endif
+
+#if defined(__GNU_LIBRARY__) && defined(HAVE_MALLOC_H)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+void
+ircd::allocator::install_free_hook()
+{
+	assert(!their_free_hook);
+	their_free_hook = __free_hook;
+	__free_hook = free_hook;
+}
+#pragma GCC diagnostic pop
+#else
+void
+ircd::allocator::install_free_hook()
+{
+}
+#endif
+
+#if defined(__GNU_LIBRARY__) && defined(HAVE_MALLOC_H)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+void
+ircd::allocator::uninstall_free_hook()
+{
+	assert(their_free_hook);
+	__free_hook = their_free_hook;
+}
+#pragma GCC diagnostic pop
+#else
+void
+ircd::allocator::uninstall_free_hook()
+{
+}
+#endif
+
 //
 // allocator::profile
 //