Merge pull request #22 from zachasme/puma-plugin

Add Puma plugin

README.md

@@ -6,17 +6,23 @@
 
 Install the gem and add to the application's Gemfile by executing:
 
-    $ bundle add litestream
+```sh
+bundle add litestream
+```
 
 If bundler is not being used to manage dependencies, install the gem by executing:
 
-    $ gem install litestream
+```sh
+gem install litestream
+```
 
 After installing the gem, run the installer:
 
-    $ rails generate litestream:install
+```sh
+rails generate litestream:install
+```
 
-The installer will create a configuration file at `config/litestream.yml`, an initializer file for configuring the gem at `config/initializers/litestream.rb`, as well as a `Procfile` so that you can run the Litestream replication process alongside your Rails application in production.
+The installer will create a configuration file at `config/litestream.yml` and an initializer file for configuring the gem at `config/initializers/litestream.rb`.
 
 This gem wraps the standalone executable version of the [Litestream](https://litestream.io/install/source/) utility. These executables are platform specific, so there are actually separate underlying gems per platform, but the correct gem will automatically be picked for your platform. Litestream itself doesn't support Windows, so this gem doesn't either.
 
@@ -77,14 +83,28 @@ However, if you need manual control over the Litestream configuration, you can m
 
 ### Replication
 
-By default, the gem will create or append to a `Procfile` to start the Litestream process via the gem's provided `litestream:replicate` rake task. This rake task will automatically load the configuration file and set the environment variables before starting the Litestream process. You can also execute this rake task yourself:
+In order to stream changes to your configured replicas, you need to start the Litestream replication process.
 
-```shell
-bin/rails litestream:replicate
-# or
-bundle exec rake litestream:replicate
+The simplest way to run the Litestream replication process is use the Puma plugin provided by the gem. This allows you to run the Litestream replication process together with Puma and have Puma monitor and manage it. You just need to add
+
+```ruby
+plugin :litestream
 ```
 
+to your `puma.rb` configuration.
+
+If you would prefer to run the Litestream replication process separately from Puma, you can use the provided `litestream:replicate` rake task. This rake task will automatically load the configuration file and set the environment variables before starting the Litestream process.
+
+The simplest way to spin up a Litestream process separately from your Rails application is to use a `Procfile`:
+
+```yaml
+# Procfile
+rails: bundle exec rails server --port $PORT
+litestream: bin/rails litestream:replicate
+```
+
+Alternatively, you could setup a `systemd` service to manage the Litestream replication process, but setting this up is outside the scope of this README.
+
 If you need to pass arguments through the rake task to the underlying `litestream` command, that can be done with argument forwarding:
 
 ```shell
@@ -94,6 +114,7 @@ bin/rails litestream:replicate -- -exec "foreman start"
 This example utilizes the `-exec` option available on [the `replicate` command](https://litestream.io/reference/replicate/) which provides basic process management, since Litestream will exit when the child process exits. In this example, we only launch our collection of Rails application processes (like Rails and SolidQueue, for example) after the Litestream replication process is ready.
 
 The Litestream `replicate` command supports the following options, which can be passed through the rake task:
+
 ```shell
 -config PATH
     Specifies the configuration file.

lib/litestream/generators/litestream/install_generator.rb

@@ -15,19 +15,6 @@ def copy_initializer_file
         template "initializer.rb", "config/initializers/litestream.rb"
       end
 
-      def create_or_update_procfile
-        if File.exist?("Procfile")
-          append_to_file "Procfile", "litestream: bin/rails litestream:replicate"
-        else
-          create_file "Procfile" do
-            <<~PROCFILE
-              rails: bundle exec rails server --port $PORT
-              litestream: bin/rails litestream:replicate
-            PROCFILE
-          end
-        end
-      end
-
       private
 
       def production_sqlite_databases

lib/puma/plugin/litestream.rb

@@ -0,0 +1,69 @@
+require "puma/plugin"
+
+# Copied from https://github.com/rails/solid_queue/blob/15408647f1780033dad223d3198761ea2e1e983e/lib/puma/plugin/solid_queue.rb
+Puma::Plugin.create do
+  attr_reader :puma_pid, :litestream_pid, :log_writer
+
+  def start(launcher)
+    @log_writer = launcher.log_writer
+    @puma_pid = $$
+
+    launcher.events.on_booted do
+      @litestream_pid = fork do
+        Thread.new { monitor_puma }
+        Litestream::Commands.replicate(async: true)
+      end
+
+      in_background do
+        monitor_litestream
+      end
+    end
+
+    launcher.events.on_stopped { stop_litestream }
+    launcher.events.on_restart { stop_litestream }
+  end
+
+  private
+
+  def stop_litestream
+    Process.waitpid(litestream_pid, Process::WNOHANG)
+    log_writer.log "Stopping Litestream..."
+    Process.kill(:INT, litestream_pid) if litestream_pid
+    Process.wait(litestream_pid)
+  rescue Errno::ECHILD, Errno::ESRCH
+  end
+
+  def monitor_puma
+    monitor(:puma_dead?, "Detected Puma has gone away, stopping Litestream...")
+  end
+
+  def monitor_litestream
+    monitor(:litestream_dead?, "Detected Litestream has gone away, stopping Puma...")
+  end
+
+  def monitor(process_dead, message)
+    loop do
+      if send(process_dead)
+        log message
+        Process.kill(:INT, $$)
+        break
+      end
+      sleep 2
+    end
+  end
+
+  def litestream_dead?
+    Process.waitpid(litestream_pid, Process::WNOHANG)
+    false
+  rescue Errno::ECHILD, Errno::ESRCH
+    true
+  end
+
+  def puma_dead?
+    Process.ppid != puma_pid
+  end
+
+  def log(...)
+    log_writer.log(...)
+  end
+end

test/generators/test_install.rb

@@ -32,10 +32,5 @@ def after_teardown
       assert_match "config.replica_key_id = litestream_credentials.replica_key_id", content
       assert_match "config.replica_access_key = litestream_credentials.replica_access_key", content
     end
-
-    assert_file "Procfile" do |content|
-      assert_match "rails: bundle exec rails server --port $PORT", content
-      assert_match "litestream: bin/rails litestream:replicate", content
-    end
   end
 end