ruby_memprofiler_pprof 0.0.1

data/ext/ruby_memprofiler_pprof/vendor/upb/upbc/protoc-gen-upbdefs.cc

@@ -0,0 +1,193 @@
+// Copyright (c) 2009-2021, Google LLC
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//     * Redistributions of source code must retain the above copyright
+//       notice, this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above copyright
+//       notice, this list of conditions and the following disclaimer in the
+//       documentation and/or other materials provided with the distribution.
+//     * Neither the name of Google LLC nor the
+//       names of its contributors may be used to endorse or promote products
+//       derived from this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY DIRECT,
+// INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include <memory>
+
+#include "google/protobuf/compiler/code_generator.h"
+#include "google/protobuf/compiler/plugin.h"
+#include "google/protobuf/descriptor.h"
+#include "google/protobuf/descriptor.pb.h"
+#include "upbc/common.h"
+
+namespace upbc {
+namespace {
+
+namespace protoc = ::google::protobuf::compiler;
+namespace protobuf = ::google::protobuf;
+
+std::string DefInitSymbol(const protobuf::FileDescriptor* file) {
+  return ToCIdent(file->name()) + "_upbdefinit";
+}
+
+static std::string DefHeaderFilename(std::string proto_filename) {
+  return StripExtension(proto_filename) + ".upbdefs.h";
+}
+
+static std::string DefSourceFilename(std::string proto_filename) {
+  return StripExtension(proto_filename) + ".upbdefs.c";
+}
+
+void GenerateMessageDefAccessor(const protobuf::Descriptor* d, Output& output) {
+  output("UPB_INLINE const upb_MessageDef *$0_getmsgdef(upb_DefPool *s) {\n",
+         ToCIdent(d->full_name()));
+  output("  _upb_DefPool_LoadDefInit(s, &$0);\n", DefInitSymbol(d->file()));
+  output("  return upb_DefPool_FindMessageByName(s, \"$0\");\n",
+         d->full_name());
+  output("}\n");
+  output("\n");
+
+  for (int i = 0; i < d->nested_type_count(); i++) {
+    GenerateMessageDefAccessor(d->nested_type(i), output);
+  }
+}
+
+void WriteDefHeader(const protobuf::FileDescriptor* file, Output& output) {
+  EmitFileWarning(file, output);
+
+  output(
+      "#ifndef $0_UPBDEFS_H_\n"
+      "#define $0_UPBDEFS_H_\n\n"
+      "#include \"upb/def.h\"\n"
+      "#include \"upb/port_def.inc\"\n"
+      "#ifdef __cplusplus\n"
+      "extern \"C\" {\n"
+      "#endif\n\n",
+      ToPreproc(file->name()));
+
+  output("#include \"upb/def.h\"\n");
+  output("\n");
+  output("#include \"upb/port_def.inc\"\n");
+  output("\n");
+
+  output("extern _upb_DefPool_Init $0;\n", DefInitSymbol(file));
+  output("\n");
+
+  for (int i = 0; i < file->message_type_count(); i++) {
+    GenerateMessageDefAccessor(file->message_type(i), output);
+  }
+
+  output(
+      "#ifdef __cplusplus\n"
+      "}  /* extern \"C\" */\n"
+      "#endif\n"
+      "\n"
+      "#include \"upb/port_undef.inc\"\n"
+      "\n"
+      "#endif  /* $0_UPBDEFS_H_ */\n",
+      ToPreproc(file->name()));
+}
+
+void WriteDefSource(const protobuf::FileDescriptor* file, Output& output) {
+  EmitFileWarning(file, output);
+
+  output("#include \"upb/def.h\"\n");
+  output("#include \"$0\"\n", DefHeaderFilename(file->name()));
+  output("#include \"$0\"\n", HeaderFilename(file));
+  output("\n");
+
+  for (int i = 0; i < file->dependency_count(); i++) {
+    output("extern _upb_DefPool_Init $0;\n",
+           DefInitSymbol(file->dependency(i)));
+  }
+
+  protobuf::FileDescriptorProto file_proto;
+  file->CopyTo(&file_proto);
+  std::string file_data;
+  file_proto.SerializeToString(&file_data);
+
+  output("static const char descriptor[$0] = {", file_data.size());
+
+  // C90 only guarantees that strings can be up to 509 characters, and some
+  // implementations have limits here (for example, MSVC only allows 64k:
+  // https://docs.microsoft.com/en-us/cpp/error-messages/compiler-errors-1/fatal-error-c1091.
+  // So we always emit an array instead of a string.
+  for (size_t i = 0; i < file_data.size();) {
+    for (size_t j = 0; j < 25 && i < file_data.size(); ++i, ++j) {
+      output("'$0', ", absl::CEscape(file_data.substr(i, 1)));
+    }
+    output("\n");
+  }
+  output("};\n\n");
+
+  output("static _upb_DefPool_Init *deps[$0] = {\n",
+         file->dependency_count() + 1);
+  for (int i = 0; i < file->dependency_count(); i++) {
+    output("  &$0,\n", DefInitSymbol(file->dependency(i)));
+  }
+  output("  NULL\n");
+  output("};\n");
+  output("\n");
+
+  output("_upb_DefPool_Init $0 = {\n", DefInitSymbol(file));
+  output("  deps,\n");
+  output("  &$0,\n", FileLayoutName(file));
+  output("  \"$0\",\n", file->name());
+  output("  UPB_STRINGVIEW_INIT(descriptor, $0)\n", file_data.size());
+  output("};\n");
+}
+
+class Generator : public protoc::CodeGenerator {
+  ~Generator() override {}
+  bool Generate(const protobuf::FileDescriptor* file,
+                const std::string& parameter, protoc::GeneratorContext* context,
+                std::string* error) const override;
+  uint64_t GetSupportedFeatures() const override {
+    return FEATURE_PROTO3_OPTIONAL;
+  }
+};
+
+bool Generator::Generate(const protobuf::FileDescriptor* file,
+                         const std::string& parameter,
+                         protoc::GeneratorContext* context,
+                         std::string* error) const {
+  std::vector<std::pair<std::string, std::string>> params;
+  google::protobuf::compiler::ParseGeneratorParameter(parameter, &params);
+
+  for (const auto& pair : params) {
+    *error = "Unknown parameter: " + pair.first;
+    return false;
+  }
+
+  std::unique_ptr<protobuf::io::ZeroCopyOutputStream> h_output_stream(
+      context->Open(DefHeaderFilename(file->name())));
+  Output h_def_output(h_output_stream.get());
+  WriteDefHeader(file, h_def_output);
+
+  std::unique_ptr<protobuf::io::ZeroCopyOutputStream> c_output_stream(
+      context->Open(DefSourceFilename(file->name())));
+  Output c_def_output(c_output_stream.get());
+  WriteDefSource(file, c_def_output);
+
+  return true;
+}
+
+}  // namespace
+}  // namespace upbc
+
+int main(int argc, char** argv) {
+  std::unique_ptr<google::protobuf::compiler::CodeGenerator> generator(
+      new upbc::Generator());
+  return google::protobuf::compiler::PluginMain(argc, argv, generator.get());
+}

data/lib/ruby_memprofiler_pprof/atfork.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+# A mechanism for adding atfork handlers, so that background profile flushing threads
+# can re-create themselves in forked children.
+# Mostly ripped off from the Datadog one here:
+# https://github.com/DataDog/dd-trace-rb/blob/master/lib/datadog/profiling/ext/forking.rb
+
+module MemprofilerPprof
+  module Atfork
+    class Handler
+      attr_accessor :block
+      attr_accessor :stage
+
+      def call
+        block.call
+      end
+
+      def remove!
+        Atfork.at_fork_handlers.reject! { |h| h === self }
+      end
+    end
+
+    module_function
+    def at_fork_handlers
+      @@at_fork_handlers ||= []
+    end
+
+    def at_fork(stage = :prepare, &block)
+      handler = Handler.new.tap do |h|
+        h.block = block
+        h.stage = stage
+      end
+      at_fork_handlers << handler
+      handler
+    end
+
+    def fork
+      # If a block is provided, it must be wrapped to trigger callbacks.
+      child_block = if block_given?
+        proc do
+          # Trigger :child callback
+          at_fork_handlers.select { |h| h.stage == :child }.each(&:call)
+
+          # Invoke original block
+          yield
+        end
+      end
+
+      # Trigger :prepare callback
+      at_fork_handlers.select { |h| h.stage == :prepare }.each(&:call)
+
+      # Start fork
+      # If a block is provided, use the wrapped version.
+      result = child_block.nil? ? super : super(&child_block)
+
+      # Trigger correct callbacks depending on whether we're in the parent or child.
+      # If we're in the fork, result = nil: trigger child callbacks.
+      # If we're in the parent, result = fork PID: trigger parent callbacks.
+      # rubocop:disable Style/IfInsideElse
+      if result.nil?
+        # Trigger :child callback
+        at_fork_handlers.select { |h| h.stage == :child }.each(&:call)
+      else
+        # Trigger :parent callback
+        at_fork_handlers.select { |h| h.stage == :parent }.each(&:call)
+      end
+      # rubocop:enable Style/IfInsideElse
+
+      # Return PID from #fork
+      result
+    end
+  end
+end
+
+::Process.singleton_class.prepend(MemprofilerPprof::Atfork)
+::Kernel.singleton_class.prepend(MemprofilerPprof::Atfork)
+TOPLEVEL_BINDING.receiver.class.prepend(MemprofilerPprof::Atfork)

data/lib/ruby_memprofiler_pprof/block_flusher.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module MemprofilerPprof
+  class BlockFlusher
+    attr_reader :collector
+
+    def initialize(collector, interval: 30, logger: nil, on_flush: nil)
+      @collector = collector
+      @interval = interval
+      @logger = logger
+      @thread = nil
+      @on_flush = on_flush
+    end
+
+    def start!
+      stop!
+      @thread = Thread.new { flusher_thread }
+      @atfork_handler = MemprofilerPprof::Atfork.at_fork(:child, &method(:at_fork_in_child))
+    end
+
+    def stop!
+      @thread&.kill
+      @thread&.join
+      @thread = nil
+      @atfork_handler&.remove!
+      @atfork_handler = nil
+    end
+
+    def run
+      start!
+      begin
+        yield
+      ensure
+        stop!
+      end
+    end
+
+    private
+
+    def flusher_thread
+      prev_flush_duration = 0
+      loop do
+        sleep([0, @interval - prev_flush_duration].max)
+        t1 = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        begin
+          profile_data = @collector.flush
+          @on_flush&.call(profile_data)
+        rescue => e
+          @logger&.error("BaseFlusher: failed to flush profiling data: #{e.inspect}")
+        ensure
+          t2 = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          prev_flush_duration = (t2 - t1)
+        end
+      end
+    end
+
+    def at_fork_in_child
+      start! if @thread && !@thread.alive?
+    end
+  end
+end

data/lib/ruby_memprofiler_pprof/file_flusher.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require 'fileutils'
+require 'forwardable'
+
+module MemprofilerPprof
+  class FileFlusher
+    extend Forwardable
+
+    def initialize(collector, pattern: 'tmp/profiles/mem-%{pid}-%{isotime}.pprof', interval: 30, logger: nil)
+      @logger = logger
+      @pattern = pattern
+      @profile_counter = 0
+      @block_flusher = BlockFlusher.new(collector, interval: interval, logger: logger, on_flush: method(:on_flush))
+    end
+
+    def_delegators :@block_flusher, :start!, :stop!, :run
+    attr_accessor :pattern
+
+    private
+
+    def on_flush(profile_data)
+      fname = template_string(@pattern)
+      dirname = File.dirname(fname)
+      FileUtils.mkdir_p dirname
+      # Need to explicitly specify the encoding, because some applications might do exotic
+      # things to File#default_external/#default_internal that would attempt to convert
+      # our protobuf to UTF-8.
+      File.write(template_string(@pattern), profile_data.pprof_data, encoding: 'ASCII-8BIT')
+      @profile_counter += 1
+    rescue => e
+      @logger&.error("FileFlusher: failed to flush profiling data: #{e.inspect}")
+    end
+
+    def template_string(tmpl)
+      vars = {
+        pid: Process.pid,
+        isotime: Time.now.iso8601,
+        unixtime: Time.now.to_i,
+        index: @profile_counter,
+      }
+      sprintf(tmpl, vars)
+    end
+  end
+end

data/lib/ruby_memprofiler_pprof/profile_app.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+# This file gets loaded by libexec/ruby_memprofiler_pprof_profile as part
+# of RUBYOPTS for a child process. This means we're going to run in some target
+# process _before_ any of the app code is run. Our job here is to set up
+# heap profiling in some default way.
+
+require 'logger'
+require 'ruby_memprofiler_pprof'
+
+collector = MemprofilerPprof::Collector.new
+collector.sample_rate = ENV.fetch('RUBY_MEMPROFILER_PPROF_SAMPLE_RATE', '1').to_f
+collector.allocation_retain_rate = ENV.fetch('RUBY_MEMPROFILER_PPROF_ALLOC_RETAIN_RATE', '1').to_f
+if ENV.key?('RUBY_MEMPROFILER_PPROF_MAX_ALLOC_SAMPLES')
+  collector.max_allocation_samples = ENV['RUBY_MEMPROFILER_PPROF_MAX_ALLOC_SAMPLES'].to_i
+end
+if ENV.key?('RUBY_MEMPROFILER_PPROF_MAX_HEAP_SAMPLES')
+  collector.max_heap_samples = ENV['RUBY_MEMPROFILER_PPROF_MAX_HEAP_SAMPLES'].to_i
+end
+
+kwargs = {
+  logger: Logger.new(STDERR)
+}
+if ENV.key?('RUBY_MEMPROFILER_PPROF_FILE_PATTERN')
+  kwargs[:pattern] = ENV['RUBY_MEMPROFILER_PPROF_FILE_PATTERN']
+end
+
+flusher = MemprofilerPprof::FileFlusher.new(collector, **kwargs)
+collector.start!
+flusher.start!

data/lib/ruby_memprofiler_pprof/profile_data.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module MemprofilerPprof
+  class ProfileData
+    attr_accessor :pprof_data, :allocation_samples_count, :heap_samples_count,
+      :dropped_samples_nolock, :dropped_samples_allocation_bufsize, :dropped_samples_heap_bufsize
+
+    def to_s
+      "<MemprofilerPprof::ProfileData:#{object_id.to_s(16)} (sample counts: " +
+        "allocation=#{allocation_samples_count}, " +
+        "heap=#{heap_samples_count}, " +
+        "dropped_nolock=#{dropped_samples_nolock}, " +
+        "dropped_allocation_bufsize=#{dropped_samples_allocation_bufsize}, " +
+        "dropped_heap_bufsize=#{dropped_samples_heap_bufsize}" +
+        ")>"
+    end
+  end
+end

data/lib/ruby_memprofiler_pprof/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module MemprofilerPprof
+  VERSION = "0.0.1"
+end

data/lib/ruby_memprofiler_pprof.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require "ruby_memprofiler_pprof/version"
+require "ruby_memprofiler_pprof/profile_data"
+require "ruby_memprofiler_pprof/ruby_memprofiler_pprof_ext"
+require 'ruby_memprofiler_pprof/atfork'
+require 'ruby_memprofiler_pprof/block_flusher'
+require 'ruby_memprofiler_pprof/file_flusher'

data/libexec/ruby_memprofiler_pprof_profile

@@ -0,0 +1,16 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+# Simply starts the provided program as-is, but sets RUBYLIB so that it will be
+# profiled by ruby_memprofiler_pprof.
+
+require 'shellwords'
+
+base_path = File.expand_path("../lib", __dir__)
+profile_lib = File.join(base_path, 'ruby_memprofiler_pprof', 'profile_app.rb')
+opts = [
+    "-I#{Shellwords.escape base_path}",
+    "-r#{Shellwords.escape profile_lib}"
+]
+ENV['RUBYOPT'] = ([ENV.fetch('RUBYOPT', nil)] + opts).compact.join(" ")
+exec *ARGV