Add experimental support for Puma

This allows us (and others) to test drive Puma without it affecting all
users. Puma can be enabled by setting the environment variable
"EXPERIMENTAL_PUMA" to a non empty value.

2 files changed, 131 insertions, 0 deletions

diff --git a/lib/gitlab/cluster/lifecycle_events.rb b/lib/gitlab/cluster/lifecycle_events.rb
new file mode 100644
index 00000000000..b05dca409d1
--- /dev/null
+++ b/lib/gitlab/cluster/lifecycle_events.rb
@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Cluster
+    #
+    # LifecycleEvents lets Rails initializers register application startup hooks
+    # that are sensitive to forking. For example, to defer the creation of
+    # watchdog threads. This lets us abstract away the Unix process
+    # lifecycles of Unicorn, Sidekiq, Puma, Puma Cluster, etc.
+    #
+    # We have three lifecycle events.
+    #
+    # - before_fork (only in forking processes)
+    # - worker_start
+    # - before_master_restart (only in forking processes)
+    #
+    # Blocks will be executed in the order in which they are registered.
+    #
+    class LifecycleEvents
+      class << self
+        #
+        # Hook registration methods (called from initializers)
+        #
+        def on_worker_start(&block)
+          if in_clustered_environment?
+            # Defer block execution
+            (@worker_start_hooks ||= []) << block
+          else
+            yield
+          end
+        end
+
+        def on_before_fork(&block)
+          return unless in_clustered_environment?
+
+          # Defer block execution
+          (@before_fork_hooks ||= []) << block
+        end
+
+        def on_master_restart(&block)
+          return unless in_clustered_environment?
+
+          # Defer block execution
+          (@master_restart_hooks ||= []) << block
+        end
+
+        #
+        # Lifecycle integration methods (called from unicorn.rb, puma.rb, etc.)
+        #
+        def do_worker_start
+          @worker_start_hooks&.each do |block|
+            block.call
+          end
+        end
+
+        def do_before_fork
+          @before_fork_hooks&.each do |block|
+            block.call
+          end
+        end
+
+        def do_master_restart
+          @master_restart_hooks && @master_restart_hooks.each do |block|
+            block.call
+          end
+        end
+
+        # Puma doesn't use singletons (which is good) but
+        # this means we need to pass through whether the
+        # puma server is running in single mode or cluster mode
+        def set_puma_options(options)
+          @puma_options = options
+        end
+
+        private
+
+        def in_clustered_environment?
+          # Sidekiq doesn't fork
+          return false if Sidekiq.server?
+
+          # Unicorn always forks
+          return true if defined?(::Unicorn)
+
+          # Puma sometimes forks
+          return true if in_clustered_puma?
+
+          # Default assumption is that we don't fork
+          false
+        end
+
+        def in_clustered_puma?
+          return false unless defined?(::Puma)
+
+          @puma_options && @puma_options[:workers] && @puma_options[:workers] > 0
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/cluster/puma_worker_killer_initializer.rb b/lib/gitlab/cluster/puma_worker_killer_initializer.rb
new file mode 100644
index 00000000000..331c39f7d6b
--- /dev/null
+++ b/lib/gitlab/cluster/puma_worker_killer_initializer.rb
@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Cluster
+    class PumaWorkerKillerInitializer
+      def self.start(puma_options, puma_per_worker_max_memory_mb: 650)
+        require 'puma_worker_killer'
+
+        PumaWorkerKiller.config do |config|
+          # Note! ram is expressed in megabytes (whereas GITLAB_UNICORN_MEMORY_MAX is in bytes)
+          # Importantly RAM is for _all_workers (ie, the cluster),
+          # not each worker as is the case with GITLAB_UNICORN_MEMORY_MAX
+          worker_count = puma_options[:workers] || 1
+          config.ram = worker_count * puma_per_worker_max_memory_mb
+
+          config.frequency = 20 # seconds
+
+          # We just want to limit to a fixed maximum, unrelated to the total amount
+          # of available RAM.
+          config.percent_usage = 0.98
+
+          # Ideally we'll never hit the maximum amount of memory. If so the worker
+          # is restarted already, thus periodically restarting workers shouldn't be
+          # needed.
+          config.rolling_restart_frequency = false
+        end
+
+        PumaWorkerKiller.start
+      end
+    end
+  end
+end