Feature: Store client as thread variable, so any handler code can ask if the socket/connection is closed, and avoid wasted computation

Is your feature request related to a problem? Please describe. Puma currently checks often (when supported) if the current connection is closed, and it skips handling a request or writing a response when that's the case. This is extremely helpful.

However if a connection is closed after Puma has passed the request to the app, the app will continue to respond to the request, perhaps doing some expensive or time-consuming work, and then hands the response back to Puma. That response is then thrown away because there's no IO to write it. It would be great to handle these cases and skip the work, to allow Puma to sooner handle more open requests.

I've found out it's possible to store a reference to the current request's Client instance, and then call Puma::Server.current.closed_socket?(client.io) to find out if the connection has been dropped in my app, and skip a bunch of stuff if it's closed.

I'm currently doing this by prepending Puma::Server#handle_request and storing the given client on a globally-accessible object, calling super, then unsetting it.

I've gotten it to work in my app by raising an exception if Puma::Server.current.closed_socket?(client.io) is true, then handling that with an empty HTTP 204 response.

Describe the solution you'd like Something like Puma::Client.current.closed? maybe? Which references the correct io object for the current thread/request, and uses the same method Puma uses internally to ask if it's closed.

Describe alternatives you've considered

1. The ideas in #1349 and #867 are to terminate the request thread or entire worker. Whilst that would be a simpler feature, it can cause issues with e.g. database connections not being cleaned up properly, so I understand why Puma hasn't done that.

2. I wasn't sure how one might send a message into whatever app you're having Puma serve. If that's possible, I don't think it would be preferable because it would likely be a lot of app-specific code in Puma. Unless an app could register a proc to be called? Even then, it still requires a lot of Puma work to constantly check if the connection is closed, and send a message as soon as that occurs. I think it'd be better to give responsibility of checking if the connection is closed to the app, if Puma can support a simple one-liner to make the check easy.

Additional context I'm mostly happy with supporting my own patch, but I just don't know if it's a bad thing in terms of Puma or performance to keep a reference to the client during handling a request. So making an official feature request to Puma I think will find the best solution for it.

Here is my example, in a Rails app:

# this file is required at bootup
require 'puma'

# Wrap internal methods so we can gain access to internal objects for each request.
module PumaServerOverload
  # This is when a request is passed over to Rails, during which we want to be able to stop execution at an opportune
  # moment when the connection is closed. We intercept here so we can store the client which has access to the socket
  # that can tell us if the connection has been closed.
  def handle_request(processor, client, requests)
    return super unless defined?(::PumaGlobal)

    wrap_with_saved_client(client) do
      super
    end
  end

  private

    def wrap_with_saved_client(client)
      ::PumaGlobal.reset!
      ::PumaGlobal.client = client
      yield
    ensure
      ::PumaGlobal.reset!
    end
end

warn "PATCHING PUMA SERVER"
Puma::Server.prepend(PumaServerOverload) unless Puma::Server < PumaServerOverload

# Track Puma objects so we can check on the connection.
class PumaGlobal < ActiveSupport::CurrentAttributes
  attribute :client

  # Don't reset automatically when CurrentAttributes gets reset at the start of a request, because we need to store
  # objects just before that happens (`handle_request` goes on to call into Rails, which then triggers the reset).
  alias_method :reset!, :reset
  def reset; end

  def connection_closed?
    client&.io && Puma::Server.current&.closed_socket?(client.io)
  end
end

class ApplicationController < ActionController::Base

  # Action can be slow if it makes database queries or otherwise does hard work. This starts before render.
  before_action :check_connection_open!

  # Render can be slow if it ends up making lots of individual database queries.
  def render(...)
    check_connection_open!
    super
  end

  class ConnectionClosedError < StandardError; end

  rescue_from ConnectionClosedError, with: :no_content

  private

    def no_content
      head :no_content
    end

    def check_connection_open!
      raise ConnectionClosedError, "connection has been closed: ok to skip rest of action" if ::PumaGlobal.connection_closed?
    end
end

I test it by firing off a bunch of requests in parallel, then terminating them all after a small duration, long enough for my app to have started handling the requests (with as many workers*threads Puma is running). I then see 204 No Content responses written by my app.

PIDS=""
for i in {1..25}; do
  echo $i
  curl https://myapp.localhost/page -s -o /dev/null & PIDS="$PIDS $!"
done
sleep 0.2
kill $PIDS

I'd be happy to start a pull-request. Unless there's a reason not to store the current client/io for the lifetime of a request?

Thanks in advance! Henry