1 files changed, 88 insertions, 0 deletions
diff --git a/lib/gitlab/memory/jemalloc.rb b/lib/gitlab/memory/jemalloc.rb
new file mode 100644
index 00000000000..454c54569de
--- /dev/null
+++ b/lib/gitlab/memory/jemalloc.rb
@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+require 'fiddle'
+
+module Gitlab
+  module Memory
+    module Jemalloc
+      extend self
+
+      STATS_FORMATS = {
+        json: { options: 'J', extension: 'json' },
+        text: { options: '', extension: 'txt' }
+      }.freeze
+
+      STATS_DEFAULT_FORMAT = :json
+
+      # Return jemalloc stats as a string.
+      def stats(format: STATS_DEFAULT_FORMAT)
+        verify_format!(format)
+
+        with_malloc_stats_print do |stats_print|
+          StringIO.new.tap { |io| write_stats(stats_print, io, STATS_FORMATS[format]) }.string
+        end
+      end
+
+      # Write jemalloc stats to the given directory.
+      def dump_stats(path:, format: STATS_DEFAULT_FORMAT)
+        verify_format!(format)
+
+        with_malloc_stats_print do |stats_print|
+          format_settings = STATS_FORMATS[format]
+          File.open(File.join(path, file_name(format_settings[:extension])), 'wb') do |io|
+            write_stats(stats_print, io, format_settings)
+          end
+        end
+      end
+
+      private
+
+      def verify_format!(format)
+        raise "format must be one of #{STATS_FORMATS.keys}" unless STATS_FORMATS.key?(format)
+      end
+
+      def with_malloc_stats_print
+        fiddle_func = malloc_stats_print
+        return unless fiddle_func
+
+        yield fiddle_func
+      end
+
+      def malloc_stats_print
+        method = Fiddle::Handle.sym("malloc_stats_print")
+
+        Fiddle::Function.new(
+          method,
+          # C signature:
+          # void (write_cb_t *write_cb, void *cbopaque, const char *opts)
+          #   arg1: callback function pointer (see below)
+          #   arg2: pointer to cbopaque holding additional callback data; always NULL here
+          #   arg3: options string, affects output format (text or JSON)
+          #
+          # Callback signature (write_cb_t):
+          # void (void *, const char *)
+          #   arg1: pointer to cbopaque data (see above; unused)
+          #   arg2: pointer to string buffer holding textual output
+          [Fiddle::TYPE_VOIDP, Fiddle::TYPE_VOIDP, Fiddle::TYPE_VOIDP],
+          Fiddle::TYPE_VOID
+        )
+      rescue Fiddle::DLError
+        # This means the Fiddle::Handle to jemalloc was not open (jemalloc wasn't loaded)
+        # or already closed. Eiher way, return nil.
+      end
+
+      def write_stats(stats_print, io, format)
+        callback = Fiddle::Closure::BlockCaller.new(
+          Fiddle::TYPE_VOID, [Fiddle::TYPE_VOIDP, Fiddle::TYPE_VOIDP]) do |_, fragment|
+          io << fragment
+        end
+
+        stats_print.call(callback, nil, format[:options])
+      end
+
+      def file_name(extension)
+        "jemalloc_stats.#{$$}.#{Time.current.to_i}.#{extension}"
+      end
+    end
+  end
+end