summaryrefslogtreecommitdiff
path: root/lib/gitlab/memory/jemalloc.rb
blob: e20e186cab9be193497ce59674406b17cdf7f640 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
# 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

      FILENAME_PREFIX = 'jemalloc_stats'

      # 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
      # @param [String] path Directory path the dump will be put into
      # @param [String] tmp_dir Directory path the dump will be streaming to. It is moved to `path` when finished.
      # @param [String] format `json` or `txt`
      # @param [String] filename_label Optional custom string that will be injected into the file name, e.g. `worker_0`
      # @return [String] Full path to the resulting dump file
      def dump_stats(path:, tmp_dir: Dir.tmpdir, format: STATS_DEFAULT_FORMAT, filename_label: nil)
        verify_format!(format)

        format_settings = STATS_FORMATS[format]
        tmp_file_path = File.join(tmp_dir, file_name(format_settings[:extension], filename_label))
        file_path = File.join(path, file_name(format_settings[:extension], filename_label))

        with_malloc_stats_print do |stats_print|
          File.open(tmp_file_path, 'wb') do |io|
            write_stats(stats_print, io, format_settings)
          end
        end

        # On OSX, `with_malloc_stats_print` is no-op, and, as result, no file will be written
        return unless File.exist?(tmp_file_path)

        FileUtils.mv(tmp_file_path, file_path)
        file_path
      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, filename_label)
        [FILENAME_PREFIX, $$, filename_label, Time.current.to_i, extension].reject(&:blank?).join('.')
      end
    end
  end
end