File: stack_prof.rb

package info (click to toggle)
ruby-test-prof 0.12.2%2Bdfsg-1
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye, forky, sid, trixie
  • size: 508 kB
  • sloc: ruby: 4,075; makefile: 4
file content (189 lines) | stat: -rw-r--r-- 4,563 bytes parent folder | download
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
# frozen_string_literal: true

module TestProf
  # StackProf wrapper.
  #
  # Has 2 modes: global and per-example.
  #
  # Example:
  #
  #   # To activate global profiling you can use env variable
  #   TEST_STACK_PROF=1 rspec ...
  #
  #   # or in your code
  #   TestProf::StackProf.run
  #
  # To profile a specific examples add :sprof tag to it:
  #
  #   it "is doing heavy stuff", :sprof do
  #     ...
  #   end
  #
  module StackProf
    # StackProf configuration
    class Configuration
      FORMATS = %w[html json].freeze

      attr_accessor :mode, :interval, :raw, :target, :format

      def initialize
        @mode = ENV.fetch("TEST_STACK_PROF_MODE", :wall).to_sym
        @target = ENV["TEST_STACK_PROF"] == "boot" ? :boot : :suite
        @raw = ENV["TEST_STACK_PROF_RAW"] != "0"
        @format =
          if FORMATS.include?(ENV["TEST_STACK_PROF_FORMAT"])
            ENV["TEST_STACK_PROF_FORMAT"]
          else
            "html"
          end

        sample_interval = ENV["TEST_STACK_PROF_INTERVAL"].to_i
        @interval = sample_interval > 0 ? sample_interval : nil
      end

      def raw?
        @raw == true
      end

      def boot?
        target == :boot
      end

      def suite?
        target == :suite
      end
    end

    class << self
      include Logging

      def config
        @config ||= Configuration.new
      end

      def configure
        yield config
      end

      # Run StackProf and automatically dump
      # a report when the process exits or when the application is booted.
      def run
        return unless profile

        @locked = true

        log :info, "StackProf#{config.raw? ? " (raw)" : ""} enabled globally: " \
                   "mode – #{config.mode}, target – #{config.target}"

        at_exit { dump("total") } if config.suite?
      end

      def profile(name = nil)
        if locked?
          log :warn, <<~MSG
            StackProf is activated globally, you cannot generate per-example report.

            Make sure you haven's set the TEST_STACK_PROF environmental variable.
          MSG
          return false
        end

        return false unless init_stack_prof

        options = {
          mode: config.mode,
          raw: config.raw
        }

        options[:interval] = config.interval if config.interval

        if block_given?
          options[:out] = build_path(name)
          ::StackProf.run(**options) { yield }
        else
          ::StackProf.start(**options)
        end
        true
      end

      def dump(name)
        ::StackProf.stop

        path = build_path(name)

        ::StackProf.results(path)

        log :info, "StackProf report generated: #{path}"

        return unless config.raw

        send("dump_#{config.format}_report", path)
      end

      private

      def build_path(name)
        TestProf.artifact_path(
          "stack-prof-report-#{config.mode}#{config.raw ? "-raw" : ""}-#{name}.dump"
        )
      end

      def locked?
        @locked == true
      end

      def init_stack_prof
        return @initialized if instance_variable_defined?(:@initialized)
        @locked = false
        @initialized = TestProf.require(
          "stackprof",
          <<~MSG
            Please, install 'stackprof' first:
               # Gemfile
              gem 'stackprof', '>= 0.2.9', require: false
          MSG
        ) { check_stack_prof_version }
      end

      def check_stack_prof_version
        if Utils.verify_gem_version("stackprof", at_least: "0.2.9")
          true
        else
          log :error, <<~MSG
            Please, upgrade 'stackprof' to version >= 0.2.9.
          MSG
          false
        end
      end

      def dump_html_report(path)
        html_path = path.gsub(/\.dump$/, ".html")

        log :info, <<~MSG
          Run the following command to generate a flame graph report:

          stackprof --flamegraph #{path} > #{html_path} && stackprof --flamegraph-viewer=#{html_path}
        MSG
      end

      def dump_json_report(path)
        report = ::StackProf::Report.new(
          Marshal.load(IO.binread(path)) # rubocop:disable Security/MarshalLoad
        )
        json_path = path.gsub(/\.dump$/, ".json")
        File.write(json_path, JSON.generate(report.data))

        log :info, <<~MSG
          StackProf JSON report generated: #{json_path}
        MSG
      end
    end
  end
end

require "test_prof/stack_prof/rspec" if TestProf.rspec?

# Hook to run StackProf globally
TestProf.activate("TEST_STACK_PROF") do
  TestProf::StackProf.run
end