Configure puma to gracefully shutdown karafka producer on stop

karafka_puma_graceful_shutdown.rb

require 'puma/plugin'

module GracefulShutdown
  # Configures puma to gracefully shutdown karafka producer
  # Related documentation: https://karafka.io/docs/Producing-messages/#producer-shutdown
  # Usage:
  #
  # (at the end of config/puma.rb)
  #
  # GracefulShutdown::SetupPuma.new.call(self)
  class SetupPuma
    def call(puma)
      setup_master(puma)
      
      # detect cluster mode (not required, but yields warning if executed in single mode)
      # https://github.com/puma/puma/blob/55e4c0bc5cc901b45f503037a84fe30bcd871848/lib/puma/dsl.rb#L1142
      opts = puma.instance_variable_get(:@options)
      setup_workers(puma) if (opts[:workers] || 0) != 0
    end

    private

    def setup_master(puma)
      plugin_name = :graceful_karafka_shutdown
      plugin_class = generate_plugin

      Puma::Plugins.register plugin_name.to_s, plugin_class
      puma.plugin plugin_name
    end

    def generate_plugin
      Class.new(Puma::Plugin) do
        def start(launcher)
          launcher.events.on_stopped do
            Karafka.producer.close
          end
        end
      end
    end

    def setup_workers(puma)
      puma.on_worker_shutdown do
        Karafka.producer.close
      end
    end
  end
end