heap_periscope_agent 0.1.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 26c1519c5f34b72f63b341f5b3bd5bcdefaab803da97b88fd59e7971fcf9de7a
-  data.tar.gz: d5f3f63e496b04726790d16095298133ce410f991b2eb1dd8675245f7083aad2
+  metadata.gz: 6adb67f8381e1ed6c7fb6dbcd0fd8e7465027e7c1df3da699c7ffd88a8b89933
+  data.tar.gz: defa2d073ad2420c5f21756ccd0e092e820271ab9b74594138841310cee50477
 SHA512:
-  metadata.gz: ce334cdc9dbe12bb33b536b0a0331c9d2574637dee691997767adeff498f4347ef7495b41ed08c7c51057409d032b800a84606b7d0e0535a836b6ca51030db0b
-  data.tar.gz: a09edf18ed45aae75e8978440e6ec36e486280b9ad376ef802901cab933d67ddf83f188a8b429ffaaf3793a3fe4f7ea412e3aebb7f59e84ab79a60571e9393a0
+  metadata.gz: d19b5b9d2b8699a18d3615659cb91dc133c41cecf1ff423f571085a66d928f4e809d872021b49b1c30bcca98eed65213a495265c9272c548abc27b0ef230cec1
+  data.tar.gz: e9b9fe743a621a3233cbda8011a02094b89fa92c4b4b52ef719059384c7324c0cb1305070457f434cdee552362a86ec78554188c985061e2895359f3206c3eef

data/lib/generators/heap_periscope_agent/config_generator.rb

@@ -0,0 +1,15 @@
+require 'rails/generators/base'
+
+module HeapPeriscopeAgent
+  module Generators
+    class ConfigGenerator < Rails::Generators::Base
+      source_root File.expand_path("templates", __dir__)
+
+      def copy_initializer_file
+        copy_file "heap_periscope_agent.rb.tt", "config/initializers/heap_periscope_agent.rb"
+        puts "\nHeapPeriscopeAgent initializer created at config/initializers/heap_periscope_agent.rb"
+        puts "Please review and customize the configuration as needed for your environment."
+      end
+    end
+  end
+end

data/lib/generators/heap_periscope_agent/controller_instrumentation_generator.rb

@@ -0,0 +1,16 @@
+require 'rails/generators/base'
+
+module HeapPeriscopeAgent
+  module Generators
+    class ControllerInstrumentationGenerator < Rails::Generators::Base
+      source_root File.expand_path("templates", __dir__)
+
+      def copy_controller_initializer
+        template "heap_periscope_agent_controller.rb.tt", "config/initializers/heap_periscope_agent_controller.rb"
+        puts "\nHeapPeriscopeAgent controller instrumentation created at config/initializers/heap_periscope_agent_controller.rb"
+        puts "This will track living objects after each controller action."
+        puts "Remember to enable `config.enable_controller_instrumentation = true` in your main HeapPeriscopeAgent initializer."
+      end
+    end
+  end
+end

data/lib/generators/heap_periscope_agent/job_tracker_generator.rb

@@ -0,0 +1,30 @@
+require 'rails/generators/base'
+require 'active_support/core_ext/string/inflections'
+
+module HeapPeriscopeAgent
+  module Generators
+    class JobTrackerGenerator < Rails::Generators::Base
+      argument :job_class_name, type: :string, desc: "The name of the Sidekiq job class to track (e.g. MyJob or MyModule::MyJob)"
+
+      def add_tracker_to_job
+        job_file_path = find_job_file
+        if job_file_path && File.exist?(job_file_path)
+          # Ensure the tracker module is required.
+          insert_into_file job_file_path, "require 'heap_periscope_agent/sidekiq_job_tracker'\n\n", before: /class|module/
+
+          # Prepend the tracker module to the class.
+          inject_into_class job_file_path, job_class_name, "  prepend HeapPeriscopeAgent::SidekiqJobTracker\n"
+          say "Added HeapPeriscopeAgent::SidekiqJobTracker to #{job_class_name} in #{job_file_path}", :green
+        else
+          say "Could not find job file for '#{job_class_name}'. Expected at '#{job_file_path}'.", :red
+        end
+      end
+
+      private
+
+      def find_job_file
+        File.join("app", "jobs", "#{job_class_name.underscore}.rb")
+      end
+    end
+  end
+end

data/lib/generators/heap_periscope_agent/model_instrumentation_generator.rb

@@ -0,0 +1,16 @@
+require 'rails/generators/base'
+
+module HeapPeriscopeAgent
+  module Generators
+    class ModelInstrumentationGenerator < Rails::Generators::Base
+      source_root File.expand_path("templates", __dir__)
+
+      def copy_model_initializer
+        template "heap_periscope_agent_model.rb.tt", "config/initializers/heap_periscope_agent_model.rb"
+        puts "\nHeapPeriscopeAgent model instrumentation created at config/initializers/heap_periscope_agent_model.rb"
+        puts "This will track living objects after each model create, update, and destroy."
+        puts "Remember to enable `config.enable_model_instrumentation = true` in your main HeapPeriscopeAgent initializer."
+      end
+    end
+  end
+end

data/lib/generators/heap_periscope_agent/rake_instrumentation_generator.rb

@@ -0,0 +1,15 @@
+require 'rails/generators/base'
+
+module HeapPeriscopeAgent
+  module Generators
+    class RakeInstrumentationGenerator < Rails::Generators::Base
+      source_root File.expand_path("templates", __dir__)
+
+      def copy_rake_initializer
+        template "heap_periscope_agent_rake.rb.tt", "config/initializers/heap_periscope_agent_rake.rb"
+        puts "\nHeapPeriscopeAgent Rake instrumentation created at config/initializers/heap_periscope_agent_rake.rb"
+        puts "This will track all Rake tasks invoked from the command line."
+      end
+    end
+  end
+end

data/lib/generators/heap_periscope_agent/rake_task_tracker_generator.rb

@@ -0,0 +1,43 @@
+require 'rails/generators/base'
+
+module HeapPeriscopeAgent
+  module Generators
+    class RakeTaskTrackerGenerator < Rails::Generators::Base
+      argument :task_name, type: :string, desc: "The name of the Rake task to track (e.g. my_namespace:my_task)"
+
+      def add_rake_tracker
+        rakefile_path = "lib/tasks/zzz_heap_periscope_agent_trackers.rake"
+
+        # Create the file with a header if it doesn't exist
+        unless File.exist?(rakefile_path)
+          create_file rakefile_path, "# This file is for auto-generated Rake task trackers for HeapPeriscopeAgent.\n# It is named with a 'zzz_' prefix to ensure it loads after other tasks are defined.\n\n"
+        end
+
+        tracker_code = <<~RUBY
+          # --- Tracker for #{task_name} task ---
+          if Rake::Task.task_defined?('#{task_name}')
+            task_to_track = Rake::Task['#{task_name}']
+            original_actions = task_to_track.actions.dup
+            task_to_track.clear_actions
+
+            task_to_track.enhance do |t, args|
+              begin
+                require 'heap_periscope_agent' # Ensure agent is loaded and configured
+                HeapPeriscopeAgent.log("Starting agent for Rake task: #{t.name}")
+                HeapPeriscopeAgent.start
+                original_actions.each { |action| action.call(t, args) } # Execute original task
+              ensure
+                HeapPeriscopeAgent.log("Stopping agent for Rake task: #{t.name}")
+                HeapPeriscopeAgent.stop
+              end
+            end
+          end
+        RUBY
+
+        append_to_file(rakefile_path, "\n#{tracker_code}\n")
+        say "Added tracker for Rake task '#{task_name}' in #{rakefile_path}", :green
+        say "NOTE: This method relies on this file being loaded *after* the task is defined. The 'zzz_' prefix helps, but is not a guarantee.", :yellow
+      end
+    end
+  end
+end

data/lib/generators/heap_periscope_agent/sidekiq_middleware_generator.rb

@@ -0,0 +1,15 @@
+require 'rails/generators/base'
+
+module HeapPeriscopeAgent
+  module Generators
+    class SidekiqMiddlewareGenerator < Rails::Generators::Base
+      source_root File.expand_path("templates", __dir__)
+
+      def copy_sidekiq_initializer
+        template "heap_periscope_agent_sidekiq.rb.tt", "config/initializers/heap_periscope_agent_sidekiq.rb"
+        puts "\nHeapPeriscopeAgent Sidekiq middleware created at config/initializers/heap_periscope_agent_sidekiq.rb"
+        puts "This will track all Sidekiq jobs. Please restart your Sidekiq server for changes to take effect."
+      end
+    end
+  end
+end

data/lib/generators/heap_periscope_agent/templates/heap_periscope_agent.rb.tt

@@ -0,0 +1,36 @@
+HeapPeriscopeAgent.configure do |config|
+  # Interval in seconds for collecting and reporting metrics.
+  # This determines how frequently a full snapshot of heap and GC stats is sent.
+  # Default: 10
+  config.interval = 10
+
+  # Host of the server where metrics will be sent via UDP.
+  # Ensure this is reachable from your application environment.
+  # Default: '127.0.0.1'
+  config.host = '127.0.0.1'
+
+  # Port of the metrics server.
+  # This must match the port your UDP server is listening on.
+  # Default: 9000
+  config.port = 9000
+
+  # Enable verbose logging from the agent to STDOUT/STDERR.
+  # Useful for debugging the agent's behavior.
+  # Default: true
+  config.verbose = true
+
+  # Enable collection of detailed object allocation information (counts by class).
+  # This can have a higher performance overhead, especially during ObjectSpace.each_object calls.
+  # Use judiciously in production.
+  # Default: false
+  config.enable_detailed_objects = false
+
+  # If detailed_objects is enabled, this limits the number of
+  # distinct object types (by class name) to report on, sorted by count.
+  # Default: 20
+  config.detailed_objects_limit = 20
+
+  config.enable_model_instrumentation = true
+  config.enable_controller_instrumentation = true
+  config.enable_rake_instrumentation = true
+end

data/lib/generators/heap_periscope_agent/templates/heap_periscope_agent_controller.rb.tt

@@ -0,0 +1,40 @@
+# config/initializers/heap_periscope_agent_controller.rb
+# This file is automatically generated by the heap_periscope_agent gem.
+# It instruments Rails controller actions to report living objects.
+
+if defined?(Rails) && Rails.version.to_i >= 3
+  ActiveSupport.on_load(:action_controller_base) do
+    # Include this module to enable instrumentation for all controller actions.
+    # Ensure `config.enable_controller_instrumentation = true` is set in your
+    # main heap_periscope_agent initializer.
+    module HeapPeriscopeAgentControllerInstrumentation
+      extend ActiveSupport::Concern
+
+      included do
+        around_action :heap_periscope_agent_track_action
+      end
+
+      private
+
+      def heap_periscope_agent_track_action
+        if HeapPeriscopeAgent.configuration.enable_controller_instrumentation
+          yield # Execute the controller action
+
+          # After the action, collect detailed living objects
+          detailed_objects = HeapPeriscopeAgent::Collector.collect_detailed_living_objects_for_span(
+            HeapPeriscopeAgent.configuration.detailed_objects_limit
+          )
+
+          span_name = "#{self.class.name}##{action_name}"
+          HeapPeriscopeAgent::Reporter.add_span_report('controller', span_name, detailed_objects)
+        else
+          yield # Execute the controller action without tracking
+        end
+      end
+    end
+
+    # Include the instrumentation module into ActionController::Base
+    # This makes it apply to all controllers inheriting from ApplicationController
+    ActionController::Base.send :include, HeapPeriscopeAgentControllerInstrumentation
+  end
+end

data/lib/generators/heap_periscope_agent/templates/heap_periscope_agent_model.rb.tt

@@ -0,0 +1,56 @@
+# config/initializers/heap_periscope_agent_model.rb
+# This file is automatically generated by the heap_periscope_agent gem.
+# It instruments Rails model actions to report living objects.
+
+if defined?(Rails) && defined?(ActiveRecord) && Rails.version.to_i >= 3
+  ActiveSupport.on_load(:active_record) do
+    # Include this module to enable instrumentation for all model actions.
+    # Ensure `config.enable_model_instrumentation = true` is set in your
+    # main heap_periscope_agent initializer.
+    module HeapPeriscopeAgentModelInstrumentation
+      extend ActiveSupport::Concern
+
+      included do
+        around_save :heap_periscope_agent_track_model_save
+        around_destroy :heap_periscope_agent_track_model_destroy
+      end
+
+      private
+
+      def heap_periscope_agent_track_model_save
+        # This check is duplicated from track_model_destroy but is cheap.
+        # It avoids capturing the `is_new_record` flag unnecessarily.
+        return yield unless HeapPeriscopeAgent.configuration.enable_model_instrumentation
+
+        is_new_record = new_record?
+        yield # Execute the save operation
+
+        # After the action, collect detailed living objects
+        detailed_objects = HeapPeriscopeAgent::Collector.collect_detailed_living_objects_for_span(
+          HeapPeriscopeAgent.configuration.detailed_objects_limit
+        )
+
+        action_name = is_new_record ? 'create' : 'update'
+        span_name = "#{self.class.name}##{action_name}"
+        HeapPeriscopeAgent::Reporter.add_span_report('model', span_name, detailed_objects)
+      end
+
+      def heap_periscope_agent_track_model_destroy
+        return yield unless HeapPeriscopeAgent.configuration.enable_model_instrumentation
+
+        yield # Execute the destroy operation
+
+        detailed_objects = HeapPeriscopeAgent::Collector.collect_detailed_living_objects_for_span(
+          HeapPeriscopeAgent.configuration.detailed_objects_limit
+        )
+
+        span_name = "#{self.class.name}#destroy"
+        HeapPeriscopeAgent::Reporter.add_span_report('model', span_name, detailed_objects)
+      end
+    end
+
+    # Include the instrumentation module into ActiveRecord::Base
+    # This makes it apply to all models inheriting from it.
+    ActiveRecord::Base.send :include, HeapPeriscopeAgentModelInstrumentation
+  end
+end

data/lib/generators/heap_periscope_agent/templates/heap_periscope_agent_rake.rb.tt

@@ -0,0 +1,28 @@
+# This file was generated by `rails generate heap_periscope_agent:rake_instrumentation`
+# It instruments Rake to track memory usage for the entire duration of a Rake command.
+
+# We only want to apply this patch when running Rake tasks, not in a Rails server or console.
+is_rake_execution = (defined?(Rake) && Rake.application.top_level_tasks.any?)
+is_server_or_console = (defined?(Rails::Server) || defined?(Rails::Console))
+
+if is_rake_execution && !is_server_or_console
+  HeapPeriscopeAgent.log("Setting up Rake instrumentation for HeapPeriscopeAgent.")
+
+  module HeapPeriscopeAgent
+    module RakeApplicationInstrumentation
+      def top_level
+        begin
+          HeapPeriscopeAgent.log("Rake execution started. Starting agent...")
+          HeapPeriscopeAgent.start
+          super
+        ensure
+          HeapPeriscopeAgent.log("Rake execution finished. Stopping agent.")
+          HeapPeriscopeAgent.stop
+        end
+      end
+    end
+  end
+
+  Rake::Application.prepend(HeapPeriscopeAgent::RakeApplicationInstrumentation)
+  HeapPeriscopeAgent.log("Rake::Application is now instrumented by HeapPeriscopeAgent.")
+end

data/lib/generators/heap_periscope_agent/templates/heap_periscope_agent_sidekiq.rb.tt

@@ -0,0 +1,24 @@
+# This file was generated by `rails generate heap_periscope_agent:sidekiq_middleware`
+# It sets up Sidekiq server middleware to automatically track memory usage for all jobs.
+
+if defined?(Sidekiq) && Sidekiq.server?
+  HeapPeriscopeAgent.log("Setting up Sidekiq server middleware for HeapPeriscopeAgent.")
+
+  class HeapPeriscopeAgent::SidekiqMiddleware
+    def call(worker, job, queue)
+      # The agent uses a reference counter, so it's safe to call start/stop for every job.
+      # It will only run one reporting thread per process.
+      HeapPeriscopeAgent.start
+      yield
+    ensure
+      HeapPeriscopeAgent.stop
+    end
+  end
+
+  Sidekiq.configure_server do |config|
+    config.server_middleware do |chain|
+      HeapPeriscopeAgent.log("Adding HeapPeriscopeAgent::SidekiqMiddleware to Sidekiq.")
+      chain.add HeapPeriscopeAgent::SidekiqMiddleware
+    end
+  end
+end

data/lib/heap_periscope_agent/collector.rb

@@ -1,12 +1,38 @@
-require 'time'
-require 'json'
-require 'socket'
 require 'objspace'
-require 'concurrent/atomic/atomic_boolean'
-require 'thread'
 
 module HeapPeriscopeAgent
   class Collector
+    PLATFORM_CLASS_IDENTIFIERS = [
+      # Ruby Core Classes
+      "String",
+      "Array",
+      "Hash",
+      "Integer",
+      "Float",
+      "Symbol",
+      "Regexp",
+      "Time",
+      "Proc",
+      "Thread",
+      "Range",
+      "Set",
+      "Enumerator",
+      "TrueClass",
+      "FalseClass",
+      "NilClass",
+      # Rails Framework Namespaces
+      "ActiveRecord::",
+      "ActiveSupport::",
+      "ActionView::",
+      "ActionController::",
+      "ActionDispatch::",
+      "ActionMailer::",
+      "ActiveJob::",
+      "ActiveModel::",
+      "Rails::",
+      "Sprockets::" # Often included with Rails asset pipeline
+    ].freeze
+
     def self.collect_snapshot(detailed_mode = false)
       data = {
         gc_stats: GC.stat,
@@ -14,21 +40,95 @@ module HeapPeriscopeAgent
       }
 
       if detailed_mode
-        data[:living_objects_by_class] = collect_detailed_living_objects(HeadPeriscopeAgent.configuration.detailed_objects_limit)
+        data[:living_objects_by_class] = collect_detailed_living_objects(HeapPeriscopeAgent.configuration.detailed_objects_limit)
       end
 
       data
     end
 
+    def self.is_platform_class?(class_name_str)
+      return false if class_name_str.nil? || class_name_str.empty?
+      PLATFORM_CLASS_IDENTIFIERS.any? do |identifier|
+        if identifier.end_with?("::")
+          class_name_str.start_with?(identifier)
+        else
+          class_name_str == identifier
+        end
+      end
+    end
+
+    def self.human_readable_bytes(bytes)
+      units = ['B', 'KB', 'MB', 'GB', 'TB']
+      return '0 B' if bytes == 0
+      i = (Math.log(bytes) / Math.log(1024)).floor
+      i = [i, units.length - 1].min
+      "#{'%.2f' % (bytes.to_f / (1024 ** i))} #{units[i]}"
+    end
+
+    def self.collect_detailed_living_objects_for_span(limit)
+      raw_data = _collect_raw_object_data
+
+      # Convert raw_data to the desired array format for spans
+      # [{ name: "String", type: "native", count: 10, size: "100KB" }, ...]
+      sorted_data = raw_data.sort_by { |_, data| -data[:count] }.first(limit)
+
+      sorted_data.map do |class_name, data|
+        {
+          name: class_name,
+          type: is_platform_class?(class_name) ? 'native' : 'application',
+          count: data[:count],
+          size: human_readable_bytes(data[:total_size])
+        }
+      end
+    end
+
     private
 
-    def self.collect_detailed_living_objects(limit)
-      counts = Hash.new(0)
+    def self._collect_raw_object_data
+      raw_data = Hash.new do |h, k|
+        h[k] = { count: 0, total_size: 0 }
+      end
+
       ObjectSpace.each_object do |obj|
-        class_name = obj.class.name rescue 'Anonymous/Singleton Class'
-        counts[class_name] += 1
+        begin
+          # Get class name
+          obj_class = obj.class
+          class_name = obj_class.name
+          class_name = obj_class.inspect if class_name.nil? || class_name.empty?
+
+          raw_data[class_name][:count] += 1
+
+          # Attempt to get object size.
+          begin
+            size = ObjectSpace.memsize_of(obj)
+            raw_data[class_name][:total_size] += size
+          rescue TypeError, NoMethodError
+            # Ignore objects that don't have a measurable size (e.g., immediates, BasicObject)
+          rescue => e
+            HeapPeriscopeAgent.log("Error getting memsize_of for #{class_name}: #{e.message}", level: :warn)
+          end
+        rescue NoMethodError # Handles BasicObject where .class is not defined
+          class_name = obj.inspect rescue 'Uninspectable BasicObject'
+          raw_data[class_name][:count] += 1
+        rescue => e
+          error_key = 'Error Collecting Object Data'
+          raw_data[error_key][:count] += 1
+          HeapPeriscopeAgent.log("Error during object iteration: #{e.message}", level: :warn)
+        end
       end
-      counts.sort_by { |_, count| -count }.first(limit).to_h
+      raw_data
+    end
+
+    # Collects detailed living objects for the main snapshot report.
+    # This method is kept for the `living_objects_by_class` part of the payload.
+    def self.collect_detailed_living_objects(limit)
+      raw_data = _collect_raw_object_data
+
+      formatted_data = raw_data.map do |class_name, data|
+        [class_name, { count: data[:count], is_platform_class: is_platform_class?(class_name) }]
+      end.to_h
+
+      formatted_data.sort_by { |_, data| -data[:count] }.first(limit).to_h
     end
   end
 end

data/lib/heap_periscope_agent/configuration.rb

@@ -1,6 +1,7 @@
 module HeapPeriscopeAgent
   class Configuration
-    attr_accessor :interval, :host, :port, :verbose, :enable_detailed_objects, :detailed_objects_limit
+    attr_accessor :interval, :host, :port, :verbose, :enable_detailed_objects,
+                  :detailed_objects_limit, :service_name, :enable_controller_instrumentation, :enable_model_instrumentation
 
     def initialize
       @interval = 10 # seconds
@@ -9,6 +10,9 @@ module HeapPeriscopeAgent
       @verbose = true
       @enable_detailed_objects = false
       @detailed_objects_limit = 20
+      @service_name = nil
+      @enable_controller_instrumentation = false
+      @enable_model_instrumentation = false
     end
   end
 end

data/lib/heap_periscope_agent/railtie.rb

@@ -0,0 +1,10 @@
+require 'heap_periscope_agent'
+
+module HeapPeriscopeAgent
+  class Railtie < Rails::Railtie
+    initializer "heap_periscope_agent.start_agent", after: :load_config_initializers do |app|
+      HeapPeriscopeAgent.log("Rails Initializer: Starting HeapPeriscopeAgent...")
+      HeapPeriscopeAgent.start
+    end
+  end
+end

data/lib/heap_periscope_agent/reporter.rb

@@ -2,64 +2,80 @@ require 'time'
 require 'json'
 require 'socket'
 require 'objspace'
-require 'concurrent/atomic/atomic_boolean'
+require 'concurrent/atomic/atomic_fixnum'
 require 'thread'
 
 module HeapPeriscopeAgent
   class Reporter
-    def self.start
-      @config = HeapPeriscopeAgent.configuration
-      @running = Concurrent::AtomicBoolean.new(false)
+    @lock = Mutex.new
+    @active_count = Concurrent::AtomicFixnum.new(0)
+    @thread = nil
+    @last_gc_total_time = 0
+    @span_reports = {}
 
-      return if @running.true?
-      @running.make_true
+    def self.start
+      @config ||= HeapPeriscopeAgent.configuration
 
-      log("Reporter starting with interval: #{@config.interval}s.")
+      if @active_count.increment == 1
+        @lock.synchronize do
+          return if @thread&.alive?
 
-      # Enable the standard GC Profiler
-      GC::Profiler.enable
-      # Store the initial total time to calculate deltas later
-      @last_gc_total_time = GC::Profiler.total_time
+          HeapPeriscopeAgent.log("Reporter starting with interval: #{@config.interval}s.")
 
-      @thread = Thread.new do
-        @socket = UDPSocket.new
-        last_snapshot_time = Time.now
+          GC::Profiler.enable
+          @last_gc_total_time = GC::Profiler.total_time
 
-        while @running.true?
-          # Periodically send a full snapshot
-          if Time.now - last_snapshot_time >= @config.interval
-            send_snapshot_report
+          @thread = Thread.new do
+            @socket = UDPSocket.new
             last_snapshot_time = Time.now
-          end
 
-          # Check for new GC activity since the last check
-          send_gc_profiler_report
+            while @active_count.value > 0
+              if Time.now - last_snapshot_time >= @config.interval
+                send_snapshot_report
+                last_snapshot_time = Time.now
+              end
 
-          sleep(1) # Main loop delay
-        end
+              send_gc_profiler_report
 
-        log("Reporter loop finished.")
-        GC::Profiler.disable # Clean up the profiler
-        @socket.close
-      end
+              sleep(1)
+            end
+
+            HeapPeriscopeAgent.log("Reporter loop finished.")
+            GC::Profiler.disable
+            @socket.close
+          end
 
-      @thread.report_on_exception = true
-      log("Reporter thread initiated.")
+          @thread.report_on_exception = true
+          HeapPeriscopeAgent.log("Reporter thread initiated.")
+        end
+      else
+        HeapPeriscopeAgent.log("Reporter already running. Active count: #{@active_count.value}")
+      end
     end
 
     def self.stop
-      return unless @running&.true?
+      return if @active_count.value.zero?
+
+      if @active_count.decrement == 0
+        @lock.synchronize do
+          return unless @active_count.value.zero?
+
+          thread_to_join = @thread
+          @thread = nil
+
+          HeapPeriscopeAgent.log("Stopping reporter...")
 
-      log("Stopping reporter...")
-      @running.make_false
-      if @thread&.join(5)
-        log("Reporter thread stopped gracefully.")
+          if thread_to_join&.join(5)
+            HeapPeriscopeAgent.log("Reporter thread stopped gracefully.")
+          else
+            HeapPeriscopeAgent.log("Reporter thread did not stop in time, killing.", level: :warn)
+            thread_to_join&.kill
+          end
+          HeapPeriscopeAgent.log("Reporter stop process complete.")
+        end
       else
-        log("Reporter thread did not stop in time, killing.", level: :warn)
-        @thread&.kill
+        HeapPeriscopeAgent.log("Jobs still active. Active count: #{@active_count.value}")
       end
-      @thread = nil
-      log("Reporter stop process complete.")
     end
 
     def self.report_once!
@@ -69,31 +85,51 @@ module HeapPeriscopeAgent
       @socket.close
     end
 
+    def self.add_span_report(span_type, span_name, objects_data)
+      @span_reports[span_type] ||= []
+      @span_reports[span_type] << { name: span_name, live_objects: objects_data }
+    end
+
     private
 
+    def self.service_name
+      if @config && @config.service_name
+        return @config.service_name
+      end
+
+      if defined?(::Sidekiq) && ::Sidekiq.server?
+        'Sidekiq'
+      elsif File.basename($0) == 'rake'
+        'Rake'
+      elsif defined?(::Rails)
+        'Rails'
+      else
+        File.basename($0)
+      end
+    end
+
     def self.send_snapshot_report
-      log("Collecting periodic snapshot...")
-      stats = HeapPeriscopeAgent.Collector.collect_snapshot(@config.enable_detailed_objects)
-      send_payload(stats, "snapshot")
+      HeapPeriscopeAgent.log("Collecting periodic snapshot...")
+      stats = HeapPeriscopeAgent::Collector.collect_snapshot(@config.enable_detailed_objects)
+
+      payload_data = stats.merge(living_objects_by_spans: @span_reports)
+      send_payload(payload_data, "snapshot")
+
+      @span_reports = {}
     end
 
     def self.send_gc_profiler_report
       current_gc_total_time = GC::Profiler.total_time
-      # Calculate time spent in GC since our last check
       gc_time_delta = current_gc_total_time - @last_gc_total_time
 
-      # If there was GC activity, report it
       if gc_time_delta > 0
-        log("Detected GC activity. Sending GC profiler report.")
+        HeapPeriscopeAgent.log("Detected GC activity. Sending GC profiler report.")
         payload = {
-          # Convert from seconds to milliseconds
           gc_duration_since_last_check_ms: (gc_time_delta * 1000).round(2),
-          gc_invocation_count: GC.count, # Total number of GCs so far
           latest_gc_info: GC.latest_gc_info,
         }
         send_payload(payload, "gc_profiler_report")
         
-        # Update the last known time
         @last_gc_total_time = current_gc_total_time
       end
     end
@@ -102,21 +138,15 @@ module HeapPeriscopeAgent
       payload = {
         type: type,
         process_id: Process.pid,
+        service_name: self.service_name,
         reported_at: Time.now.utc.iso8601,
         payload: data
       }.to_json
 
       @socket.send(payload, 0, @config.host, @config.port)
-      log("Sent #{type} payload.")
+      HeapPeriscopeAgent.log("Sent #{type} payload.")
     rescue => e
-      log("Failed to send payload: #{e.message}", level: :error)
-    end
-
-    def self.log(message, level: :info)
-      return unless @config&.verbose
-      timestamp = Time.now.strftime('%Y-%m-%d %H:%M:%S')
-      output = "[HeapPeriscopeAgent][#{level.to_s.upcase}] #{timestamp}: #{message}"
-      level == :error ? warn(output) : puts(output)
+      HeapPeriscopeAgent.log("Failed to send payload: #{e.message}", level: :error)
     end
   end
 end

data/lib/heap_periscope_agent/sidekiq_job_tracker.rb

@@ -0,0 +1,12 @@
+module HeapPeriscopeAgent
+    module SidekiqJobTracker
+      def perform(*args)
+        HeapPeriscopeAgent.log("Starting agent for #{self.class.name}...")
+        HeapPeriscopeAgent.start
+        super
+      ensure
+        HeapPeriscopeAgent.log("Stopping agent for #{self.class.name}.")
+        HeapPeriscopeAgent.stop
+      end
+    end
+end

data/lib/heap_periscope_agent/version.rb

@@ -1,3 +1,3 @@
 module HeapPeriscopeAgent
-  VERSION = "0.1.0"
+  VERSION = "0.3.0"
 end

data/lib/heap_periscope_agent.rb

@@ -1,14 +1,54 @@
-require_relative './heap_periscope_agent/version.rb'
-require_relative './heap_periscope_agent/configuration.rb'
-require_relative './heap_periscope_agent/collector.rb'
-require_relative './heap_periscope_agent/reporter.rb'
+require_relative "heap_periscope_agent/version"
+require_relative "heap_periscope_agent/configuration"
+require_relative "heap_periscope_agent/collector"
+require_relative "heap_periscope_agent/reporter"
+require_relative "heap_periscope_agent/sidekiq_job_tracker"
 
 module HeapPeriscopeAgent
-  def self.configuration
-    @configuration ||= HeapPeriscopeAgent::Configuration.new
+  class << self
+    attr_writer :configuration
+
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    def configure
+      yield(configuration)
+    end
+
+    def start
+      log("HeapPeriscopeAgent starting...")
+      Reporter.start
+      log("HeapPeriscopeAgent started successfully.")
+    rescue => e
+      log("Failed to start HeapPeriscopeAgent: #{e.message} #{e.backtrace.join("\n")}", level: :error)
+    end
+
+    def stop
+      log("HeapPeriscopeAgent stopping...")
+      Reporter.stop
+      log("HeapPeriscopeAgent stopped.")
+    rescue => e
+      log("Failed to stop HeapPeriscopeAgent: #{e.message}", level: :error)
+    end
+
+    def report_once!
+      Reporter.report_once!
+    end
+
+    def log(message, level: :info)
+      return unless configuration&.verbose
+      timestamp = Time.now.strftime('%Y-%m-%d %H:%M:%S')
+      output = "[HeapPeriscopeAgent][#{level.to_s.upcase}] #{timestamp}: #{message}"
+      if level == :error || level == :warn
+        Kernel.warn(output)
+      else
+        Kernel.puts(output)
+      end
+    end
   end
 
-  def self.configure
-    yield(configuration)
+  if defined?(Rails::Railtie)
+    require_relative "heap_periscope_agent/railtie"
   end
-end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: heap_periscope_agent
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.3.0
 platform: ruby
 authors:
 - Jonathan Natanael Siahaan
@@ -13,16 +13,36 @@ description: |
   Heap Periscope Agent offers deep insights into your Ruby application's memory behavior.
   It collects and reports real-time Garbage Collection (GC) statistics and object
   allocation patterns, empowering developers to identify memory leaks, optimize usage,
-  and enhance performance. Highly configurable and designed for minimal overhead.
+  and enhance performance. This gem is the backend agent for memory monitoring. To
+  visualize the collected data, you must also install the companion gem,
+  heap_periscope_ui
+
+  The agent's visualizer is available here:
+  - Gem: https://rubygems.org/gems/heap_periscope_ui
+  - Repository: https://github.com/codepawpaw/heap_periscope_ui
 email: js.jonathan.n@gmail.com
 executables: []
 extensions: []
 extra_rdoc_files: []
 files:
+- lib/generators/heap_periscope_agent/config_generator.rb
+- lib/generators/heap_periscope_agent/controller_instrumentation_generator.rb
+- lib/generators/heap_periscope_agent/job_tracker_generator.rb
+- lib/generators/heap_periscope_agent/model_instrumentation_generator.rb
+- lib/generators/heap_periscope_agent/rake_instrumentation_generator.rb
+- lib/generators/heap_periscope_agent/rake_task_tracker_generator.rb
+- lib/generators/heap_periscope_agent/sidekiq_middleware_generator.rb
+- lib/generators/heap_periscope_agent/templates/heap_periscope_agent.rb.tt
+- lib/generators/heap_periscope_agent/templates/heap_periscope_agent_controller.rb.tt
+- lib/generators/heap_periscope_agent/templates/heap_periscope_agent_model.rb.tt
+- lib/generators/heap_periscope_agent/templates/heap_periscope_agent_rake.rb.tt
+- lib/generators/heap_periscope_agent/templates/heap_periscope_agent_sidekiq.rb.tt
 - lib/heap_periscope_agent.rb
 - lib/heap_periscope_agent/collector.rb
 - lib/heap_periscope_agent/configuration.rb
+- lib/heap_periscope_agent/railtie.rb
 - lib/heap_periscope_agent/reporter.rb
+- lib/heap_periscope_agent/sidekiq_job_tracker.rb
 - lib/heap_periscope_agent/version.rb
 homepage: https://github.com/codepawpaw/heap_periscope_agent
 licenses: