From 7d819bdaaf41e3ee1b3104ac221773f0f53e9cf2 Mon Sep 17 00:00:00 2001
From: Mike Larsson <mikelarssonftw@gmail.com>
Date: Mon, 24 Apr 2017 16:34:07 -0400
Subject: [PATCH] Add sidekiq profiler middleware that puts periodic heap dumps
 on s3

---
 config/initializers/sidekiq.rb |  9 +++++
 lib/sidekiq_profiler.rb        | 62 ++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)
 create mode 100644 lib/sidekiq_profiler.rb

diff --git a/config/initializers/sidekiq.rb b/config/initializers/sidekiq.rb
index ac5e34b6b1..9a5f79ddfe 100644
--- a/config/initializers/sidekiq.rb
+++ b/config/initializers/sidekiq.rb
@@ -7,4 +7,13 @@
   if database_url = ENV['DATABASE_URL']
     ActiveRecord::Base.establish_connection "#{database_url}?pool=25"
   end
+
+  if ENV["PROFILE_SIDEKIQ"].present?
+    require "sidekiq_profiler"
+    ObjectSpace.trace_object_allocations_start
+    Sidekiq.logger.info "allocations tracing enabled"
+    config.server_middleware do |chain|
+      chain.add Sidekiq::Middleware::Server::Profiler
+    end
+  end
 end
diff --git a/lib/sidekiq_profiler.rb b/lib/sidekiq_profiler.rb
new file mode 100644
index 0000000000..1254f52e93
--- /dev/null
+++ b/lib/sidekiq_profiler.rb
@@ -0,0 +1,62 @@
+require "objspace"
+require "tempfile"
+
+class HeapDumpUploader < CarrierWave::Uploader::Base
+  storage :fog
+
+  def store_dir
+    "heap_dumps"
+  end
+end
+
+module Sidekiq
+  module Middleware
+    module Server
+      class Profiler
+
+        # Number of jobs to process before reporting
+        JOBS = ENV.fetch("SIDEKIQ_PROFILER_REPORT_INCREMENT") { 100 }.to_i
+        STAMP = Time.now.to_i
+
+        class << self
+          mattr_accessor :counter
+          self.counter = 0
+
+          def synchronize(&block)
+            @lock ||= Mutex.new
+            @lock.synchronize(&block)
+          end
+        end
+
+        def call(worker_instance, item, queue)
+          begin
+            yield
+          ensure
+            self.class.synchronize do
+              self.class.counter += 1
+
+              if self.class.counter % JOBS == 0
+                Sidekiq.logger.info "reporting allocations after #{self.class.counter} jobs"
+                GC.start
+                basename = "heap-#{STAMP}-#{self.class.counter}-"
+                f = Tempfile.new([basename, '.json'])
+                begin
+                  ObjectSpace.dump_all(output: f)
+                  Sidekiq.logger.info "heap saved to #{f.path}"
+                  if ENV["UPLOAD_SIDEKIQ_HEAP_DUMPS"].present?
+                    uploader = HeapDumpUploader.new
+                    uploader.store!(f)
+                    Sidekiq.logger.info "heap uploaded to #{uploader.url}"
+                  end
+                ensure
+                  f.close
+                  f.unlink
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end