Plugin hooks for WebSockets

docs/plugins.md

@@ -1,28 +1,52 @@
 ## Plugins
 
-Puma 3.0 added support for plugins that can augment configuration and service operations.
+Puma 3.0 added support for plugins that can augment configuration and service
+operations.
 
 2 canonical plugins to look to aid in development of further plugins:
 
-* [tmp\_restart](https://github.com/puma/puma/blob/master/lib/puma/plugin/tmp_restart.rb): Restarts the server if the file `tmp/restart.txt` is touched
-* [heroku](https://github.com/puma/puma-heroku/blob/master/lib/puma/plugin/heroku.rb): Packages up the default configuration used by puma on Heroku
+* [tmp\_restart](https://github.com/puma/puma/blob/master/lib/puma/plugin/tmp_restart.rb):
+  Restarts the server if the file `tmp/restart.txt` is touched
+* [heroku](https://github.com/puma/puma-heroku/blob/master/lib/puma/plugin/heroku.rb):
+  Packages up the default configuration used by puma on Heroku
 
-Plugins are activated in a puma configuration file (such as `config/puma.rb'`) by adding `plugin "name"`, such as `plugin "heroku"`.
+Plugins are activated in a puma configuration file (such as `config/puma.rb'`)
+by adding `plugin "name"`, such as `plugin "heroku"`.
 
-Plugins are activated based simply on path requirements so, activating the `heroku` plugin will simply be doing `require "puma/plugin/heroku"`. This allows gems to provide multiple plugins (as well as unrelated gems to provide puma plugins).
+Plugins are activated based simply on path requirements so, activating the
+`heroku` plugin will simply be doing `require "puma/plugin/heroku"`. This
+allows gems to provide multiple plugins (as well as unrelated gems to provide
+puma plugins).
 
 The `tmp_restart` plugin is bundled with puma, so it can always be used.
 
 To use the `heroku` plugin, add `puma-heroku` to your Gemfile or install it.
 
 ### API
 
-At present, there are 2 hooks that plugins can use: `start` and `config`.
+## Server-wide hooks
 
-`start` runs when the server has started and allows the plugin to start other functionality to augment puma.
+Plugins can use a couple of hooks at server level: `start` and `config`.
 
-`config` runs when the server is being configured and is passed a `Puma::DSL` object that can be used to add additional configuration.
+`start` runs when the server has started and allows the plugin to start other
+functionality to augment puma.
 
-Any public methods in `Puma::Plugin` are the public API that any plugin may use.
+`config` runs when the server is being configured and is passed a `Puma::DSL`
+object that can be used to add additional configuration.
 
-In the future, more hooks and APIs will be added.
+Any public methods in `Puma::Plugin` are the public API that any plugin may
+use.
+
+## Per request hooks
+
+`#on_before_rack(env)` will be called right before the Rack application is
+invoked. The called hook may modify `env` just like any Rack middleware.
+
+`#on_after_rack(env, headers, io)` will be called after the Rack application
+has completed its execution and before any response content is written to the
+client. A plugin may take over from here by returning an object that responds
+to `#stream?` with a truthy value. Check out `lib/puma/stream_client.rb` to
+know more about this interface.
+
+If more than one plugin arises interest in taking over, an exception will
+be happen and Puma will serve a 500.

lib/puma/client.rb

@@ -141,6 +141,14 @@ def close
       end
     end
 
+    def on_shutdown
+      close
+    end
+
+    def stream?
+      false
+    end
+
     # The object used for a request with no body. All requests with
     # no body share this one object since it has no state.
     EmptyBody = NullIO.new

lib/puma/plugin.rb

@@ -57,6 +57,10 @@ def find(name)
       raise UnknownPlugin, "file failed to register a plugin"
     end
 
+    def each
+      @plugins.each_value { |plugin| yield plugin }
+    end
+
     def add_background(blk)
       @background << blk
     end

lib/puma/reactor.rb

@@ -189,7 +189,7 @@ def run_internal
                     if submon.value == @ready
                       false
                     else
-                      submon.value.close
+                      submon.value.on_shutdown
                       begin
                         selector.deregister submon.value
                       rescue IOError
@@ -215,6 +215,18 @@ def run_internal
                 end
               end
 
+              if c.stream?
+                if c.on_read_ready
+                  @app_pool << c
+                end
+
+                if c.closed?
+                  clear_monitor mon
+                end
+
+                next
+              end
+
               begin
                 if c.try_to_finish
                   @app_pool << c

lib/puma/server.rb

@@ -85,6 +85,19 @@ def initialize(app, events=Events.stdio, options={})
       @mode = :http
 
       @precheck_closing = true
+
+      @on_before_rack = nil
+      @on_after_rack = nil
+      Plugins.each do |plugin|
+        if plugin.respond_to? :on_before_rack
+          @on_before_rack ||= []
+          @on_before_rack << plugin.method(:on_before_rack)
+        end
+        if plugin.respond_to? :on_after_rack
+          @on_after_rack ||= []
+          @on_after_rack << plugin.method(:on_after_rack)
+        end
+      end
     end
 
     attr_accessor :binder, :leak_stack_on_error, :early_hints
@@ -296,6 +309,13 @@ def run(background=true)
                                     @max_threads,
                                     IOBuffer) do |client, buffer|
 
+        if client.respond_to? :churn
+          more_to_churn = client.churn
+          @thread_pool << client if more_to_churn
+
+          next
+        end
+
         # Advertise this server into the thread
         Thread.current[ThreadLocalKey] = self
 
@@ -652,6 +672,10 @@ def handle_request(req, lines)
       #
       after_reply = env[RACK_AFTER_REPLY] = []
 
+      if @on_before_rack
+        @on_before_rack.each { |hook| hook.call(env) }
+      end
+
       begin
         begin
           status, headers, res_body = @app.call(env)
@@ -681,6 +705,30 @@ def handle_request(req, lines)
           status, headers, res_body = lowlevel_error(e, env)
         end
 
+        if @on_after_rack
+          is_async = false
+
+          @on_after_rack.reverse_each do |hook|
+            stream_client = hook.call(env, headers, req.io)
+
+            if stream_client && stream_client.stream?
+              if is_async
+                raise "Only one #on_after_rack hook should take over"
+              else
+                if stream_client.on_read_ready
+                  @thread_pool << stream_client
+                end
+
+                @reactor.add stream_client
+
+                is_async = true
+              end
+            end
+          end
+
+          return :async if is_async
+        end
+
         content_length = nil
         no_body = head
 

lib/puma/stream_client.rb

@@ -0,0 +1,72 @@
+module Puma
+  # This serves as a base class for any plugin that wants to take over the
+  # socket and wait on IO.
+  #
+  # The subclasses are expected to implement:
+  #
+  #  1. Methods to respond to changes in the underlying socket. These are
+  #     `#on_read_ready`, `#on_broken_pipe` & `#on_shutdown`.
+  #  2. A `#churn` method that runs within the thread pool, this is what can
+  #     be used to invoke app's logic.
+  #
+  # The underlying socket is available through `@io`.
+  #
+  # The other methods should never be overriden.
+  class StreamClient
+    def initialize(io)
+      @io = io
+    end
+
+    def to_io
+      @io
+    end
+
+    def timeout_at
+      false
+    end
+
+    def closed?
+      @io.closed?
+    end
+
+    def stream?
+      true
+    end
+
+    # This method will be invoked when the IO descriptor has new data pending
+    # to be read. You can read from `@io` at this time.
+    #
+    # You can return a truthy value to add this client to the thread pool.
+    def on_read_ready
+      raise NotImplementedError
+    end
+
+    # This method will be invoked when the underlying connection is broken
+    # for whatever reason (client timeout, abrupt client disconnection, etc.).
+    #
+    # You can return a truthy value to add this client to the thread pool.
+    def on_broken_pipe
+      raise NotImplementedError
+    end
+
+    # This method will be invoked when the server is being stopped.
+    #
+    # It's not possible to run more work on the thread pool at this stage.
+    def on_shutdown
+      raise NotImplementedError
+    end
+
+    # This is the method that the thread pool will be consuming. A "churn" is
+    # enqueued on the thread pool each time `#churn` or `#read_more` return a
+    # truthy value.
+    #
+    # You can return a truthy value to add this client to the thread pool
+    # again.
+    #
+    # This allows the plugin to process its work on the thread pool as any
+    # other regular HTTP request.
+    def churn
+      raise NotImplementedError
+    end
+  end
+end