gitlab-build-images/patches/ruby/3.0/thread-memory-allocations-3.0.patch

From 97f14ebfd8d24d71e10c450e0a90b6322f9c0d59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Tue, 22 Dec 2020 15:33:08 +0100
Subject: [PATCH] Expose `Thread#memory_allocations` counters

This provides currently a per-thread GC heap slots
and malloc allocations statistics.

This is designed to measure a memory allocations
in a multi-threaded environments (concurrent requests
processing) with an accurate information about allocated
memory within a given execution context.

Example: Measure memory pressure generated by a given
requests to easier find requests with a lot of allocations.

Ref: https://gitlab.com/gitlab-org/gitlab/-/issues/296530
---
 gc.c                                          | 20 ++++++
 .../test_thread_trace_memory_allocations.rb   | 67 +++++++++++++++++++
 thread.c                                      | 55 +++++++++++++++
 vm_core.h                                     | 17 +++++
 4 files changed, 159 insertions(+)
 create mode 100644 test/ruby/test_thread_trace_memory_allocations.rb

diff --git a/gc.c b/gc.c
index 27cf65b196a3..280c62fbe341 100644
--- a/gc.c
+++ b/gc.c
@@ -2123,6 +2123,13 @@ newobj_init(VALUE klass, VALUE flags, int wb_protected, rb_objspace_t *objspace,
     // TODO: make it atomic, or ractor local
     objspace->total_allocated_objects++;
 
+#if THREAD_TRACE_MEMORY_ALLOCATIONS
+    rb_thread_t *th = ruby_threadptr_for_trace_memory_allocations();
+    if (th) {
+	ATOMIC_SIZE_INC(th->memory_allocations.total_allocated_objects);
+    }
+#endif
+
 #if RGENGC_PROFILE
     if (wb_protected) {
         objspace->profile.total_generated_normal_object_count++;
@@ -10487,6 +10494,19 @@ objspace_malloc_increase(rb_objspace_t *objspace, void *mem, size_t new_size, si
 #endif
     }
 
+#if THREAD_TRACE_MEMORY_ALLOCATIONS
+    rb_thread_t *th = ruby_threadptr_for_trace_memory_allocations();
+    if (th) {
+	if (new_size > old_size) {
+	    ATOMIC_SIZE_ADD(th->memory_allocations.total_malloc_bytes, new_size - old_size);
+	}
+
+	if (type == MEMOP_TYPE_MALLOC) {
+	    ATOMIC_SIZE_INC(th->memory_allocations.total_mallocs);
+	}
+    }
+#endif
+
     if (type == MEMOP_TYPE_MALLOC) {
       retry:
 	if (malloc_increase > malloc_limit && ruby_native_thread_p() && !dont_gc_val()) {
diff --git a/test/ruby/test_thread_trace_memory_allocations.rb b/test/ruby/test_thread_trace_memory_allocations.rb
new file mode 100644
index 000000000000..2e281513578b
--- /dev/null
+++ b/test/ruby/test_thread_trace_memory_allocations.rb
@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require 'test/unit'
+
+class TestThreadTraceMemoryAllocations < Test::Unit::TestCase
+  def test_disabled_trace_memory_allocations
+    Thread.trace_memory_allocations = false
+
+    assert_predicate Thread.current.memory_allocations, :nil?
+  end
+
+  def test_enabled_trace_memory_allocations
+    Thread.trace_memory_allocations = true
+
+    assert_not_nil(Thread.current.memory_allocations)
+  end
+
+  def test_only_this_thread_allocations_are_counted
+    changed = {
+      total_allocated_objects: 1000,
+      total_malloc_bytes: 1_000_000,
+      total_mallocs: 100
+    }
+
+    Thread.trace_memory_allocations = true
+
+    assert_less_than(changed) do
+      Thread.new do
+       assert_greater_than(changed) do
+          # This will allocate: 5k objects, 5k mallocs, 5MB
+          allocate(5000, 1000)
+       end
+      end.join
+
+      # This will allocate: 50 objects, 50 mallocs, 500 bytes
+      allocate(50, 10)
+    end
+  end
+
+  private
+
+  def allocate(slots, bytes)
+    Array.new(slots).map do
+      '0' * bytes
+    end
+  end
+
+  def assert_greater_than(keys)
+    before = Thread.current.memory_allocations
+    yield
+    after = Thread.current.memory_allocations
+
+    keys.each do |key, by|
+      assert_operator(by, :<=, after[key]-before[key], "expected the #{key} to change more than #{by}")
+    end
+  end
+
+  def assert_less_than(keys)
+    before = Thread.current.memory_allocations
+    yield
+    after = Thread.current.memory_allocations
+
+    keys.each do |key, by|
+      assert_operator(by, :>, after[key]-before[key], "expected the #{key} to change less than #{by}")
+    end
+  end
+end
diff --git a/thread.c b/thread.c
index dce181d24e02..247440766cdf 100644
--- a/thread.c
+++ b/thread.c
@@ -5412,6 +5412,55 @@ Init_Thread_Mutex(void)
     rb_native_mutex_initialize(&th->interrupt_lock);
 }
 
+#if THREAD_TRACE_MEMORY_ALLOCATIONS
+rb_thread_t *
+ruby_threadptr_for_trace_memory_allocations(void)
+{
+    // The order of this checks is important due
+    // to how Ruby VM is initialized
+    if (GET_VM()->thread_trace_memory_allocations && GET_EC() != NULL) {
+        return GET_THREAD();
+    }
+
+    return NULL;
+}
+
+static VALUE
+rb_thread_s_trace_memory_allocations(VALUE _)
+{
+    return GET_THREAD()->vm->thread_trace_memory_allocations ? Qtrue : Qfalse;
+}
+
+static VALUE
+rb_thread_s_trace_memory_allocations_set(VALUE self, VALUE val)
+{
+    GET_THREAD()->vm->thread_trace_memory_allocations = RTEST(val);
+    return val;
+}
+
+static VALUE
+rb_thread_memory_allocations(VALUE self)
+{
+    rb_thread_t *th = rb_thread_ptr(self);
+
+    if (!th->vm->thread_trace_memory_allocations) {
+        return Qnil;
+    }
+
+    VALUE ret = rb_hash_new();
+
+	VALUE total_allocated_objects = ID2SYM(rb_intern_const("total_allocated_objects"));
+	VALUE total_malloc_bytes = ID2SYM(rb_intern_const("total_malloc_bytes"));
+	VALUE total_mallocs = ID2SYM(rb_intern_const("total_mallocs"));
+
+	rb_hash_aset(ret, total_allocated_objects, SIZET2NUM(th->memory_allocations.total_allocated_objects));
+	rb_hash_aset(ret, total_malloc_bytes, SIZET2NUM(th->memory_allocations.total_malloc_bytes));
+	rb_hash_aset(ret, total_mallocs, SIZET2NUM(th->memory_allocations.total_mallocs));
+
+    return ret;
+}
+#endif
+
 /*
  *  Document-class: ThreadError
  *
@@ -5497,6 +5546,12 @@ Init_Thread(void)
     rb_define_method(rb_cThread, "to_s", rb_thread_to_s, 0);
     rb_define_alias(rb_cThread, "inspect", "to_s");
 
+#if THREAD_TRACE_MEMORY_ALLOCATIONS
+    rb_define_singleton_method(rb_cThread, "trace_memory_allocations", rb_thread_s_trace_memory_allocations, 0);
+    rb_define_singleton_method(rb_cThread, "trace_memory_allocations=", rb_thread_s_trace_memory_allocations_set, 1);
+    rb_define_method(rb_cThread, "memory_allocations", rb_thread_memory_allocations, 0);
+#endif
+
     rb_vm_register_special_exception(ruby_error_stream_closed, rb_eIOError,
 				     "stream closed in another thread");
 
diff --git a/vm_core.h b/vm_core.h
index 5f8d4ab87670..ac15f72fa25b 100644
--- a/vm_core.h
+++ b/vm_core.h
@@ -97,6 +97,13 @@
 # define VM_INSN_INFO_TABLE_IMPL 2
 #endif
 
+/*
+ * track a per thread memory allocations
+ */
+#ifndef THREAD_TRACE_MEMORY_ALLOCATIONS
+# define THREAD_TRACE_MEMORY_ALLOCATIONS 1
+#endif
+
 #if defined(NSIG_MAX)           /* POSIX issue 8 */
 # undef NSIG
 # define NSIG NSIG_MAX
@@ -606,6 +613,7 @@ typedef struct rb_vm_struct {
     unsigned int thread_abort_on_exception: 1;
     unsigned int thread_report_on_exception: 1;
     unsigned int thread_ignore_deadlock: 1;
+    unsigned int thread_trace_memory_allocations: 1;
 
     /* object management */
     VALUE mark_object_ary;
@@ -981,6 +989,14 @@ typedef struct rb_thread_struct {
 
     struct rb_waiting_list *join_list;
 
+#if THREAD_TRACE_MEMORY_ALLOCATIONS
+    struct {
+        size_t total_allocated_objects;
+        size_t total_malloc_bytes;
+        size_t total_mallocs;
+    } memory_allocations;
+#endif
+
     union {
         struct {
             VALUE proc;
@@ -1901,6 +1917,7 @@ void rb_threadptr_interrupt(rb_thread_t *th);
 void rb_threadptr_unlock_all_locking_mutexes(rb_thread_t *th);
 void rb_threadptr_pending_interrupt_clear(rb_thread_t *th);
 void rb_threadptr_pending_interrupt_enque(rb_thread_t *th, VALUE v);
+rb_thread_t *ruby_threadptr_for_trace_memory_allocations(void);
 VALUE rb_ec_get_errinfo(const rb_execution_context_t *ec);
 void rb_ec_error_print(rb_execution_context_t * volatile ec, volatile VALUE errinfo);
 void rb_execution_context_update(const rb_execution_context_t *ec);