Implement Unicorn hooks that are settable from within Gitlab source codes

3 files changed, 101 insertions, 0 deletions

diff --git a/config/initializers/unicorn.rb b/config/initializers/unicorn.rb
new file mode 100644
index 00000000000..18ad88f195a
--- /dev/null
+++ b/config/initializers/unicorn.rb
@@ -0,0 +1,4 @@
+Gitlab::Unicorn::Hook.after_fork do |server, worker|
+  ## TODO: `Prometheus::Client.store_worker_id(worker.id)`
+  Prometheus::Client.reinitialize_on_pid_change
+end
diff --git a/lib/gitlab/unicorn/hook.rb b/lib/gitlab/unicorn/hook.rb
new file mode 100644
index 00000000000..e15f3a7f8a9
--- /dev/null
+++ b/lib/gitlab/unicorn/hook.rb
@@ -0,0 +1,25 @@
+module Gitlab
+  module Unicorn
+    module Hook
+      class << self
+        attr_writer :before_fork, :after_fork
+
+        def before_fork(&block)
+          @before_fork = block
+        end
+
+        def after_fork(&block)
+          @after_fork = block
+        end
+
+        def run_before_fork(server, worker)
+          @before_fork.call(server, worker) if @before_fork
+        end
+
+        def run_after_fork(server, worker)
+          @after_fork.call(server, worker) if @after_fork
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/unicorn/hook_spec.rb b/spec/lib/gitlab/unicorn/hook_spec.rb
new file mode 100644
index 00000000000..99918510413
--- /dev/null
+++ b/spec/lib/gitlab/unicorn/hook_spec.rb
@@ -0,0 +1,72 @@
+require 'spec_helper'
+
+describe Gitlab::Unicorn::Hook do
+  let(:canary) { double(:canary) }
+  let(:server) { double(:server) }
+  let(:worker) { double(:worker) }
+
+  before do
+    described_class.before_fork = nil
+    described_class.after_fork = nil
+  end
+
+  describe '#run_before_fork' do
+    context 'before_fork hook configured' do
+      before do
+        allow(canary).to receive(:trigger)
+
+        described_class.before_fork do |server, worker|
+          canary.trigger(server, worker)
+        end
+      end
+
+      it 'runs before fork' do
+        described_class.run_before_fork(server, worker)
+
+        expect(canary).to have_received(:trigger).with(server, worker)
+      end
+    end
+
+    context 'before_fork hook not configured' do
+      before do
+        allow(canary).to receive(:trigger)
+      end
+
+      it 'runs before fork' do
+        described_class.run_before_fork(server, worker)
+
+        expect(canary).not_to have_received(:trigger).with(server, worker)
+      end
+    end
+  end
+
+  describe '#run_after_fork' do
+    context 'after_fork hook configured' do
+      before do
+        allow(canary).to receive(:trigger)
+
+        described_class.after_fork do |server, worker|
+          canary.trigger(server, worker)
+        end
+      end
+
+      it 'runs before fork' do
+        described_class.run_after_fork(server, worker)
+
+        expect(canary).to have_received(:trigger).with(server, worker)
+      end
+    end
+
+    context 'after_fork hook not configured' do
+      before do
+        allow(canary).to receive(:trigger)
+      end
+
+      it 'runs before fork' do
+        described_class.run_after_fork(server, worker)
+
+        expect(canary).not_to have_received(:trigger).with(server, worker)
+      end
+    end
+  end
+end