Attempt to make stream(:keep_open) work with Puma

Only running the async tests passes in Docker for me now

    TESTOPTS="--name='/streams_async/'" bundle exec rake

but I've had the tests (or a server) being out of memory killed, is it
webrick that's the issue? Going to see how this behave on GitHub Actions

Co-authored-by: Jordan Owens <jkowens@gmail.com>

examples/chat.rb

@@ -1,11 +1,9 @@
 #!/usr/bin/env ruby -I ../lib -I lib
 # frozen_string_literal: true
 
-require_relative 'rainbows'
-
 require 'sinatra'
-set :server, :rainbows
-connections = []
+set :server, :puma
+connections = Set.new
 
 get '/' do
   halt erb(:login) unless params[:user]
@@ -16,11 +14,16 @@
   stream :keep_open do |out|
     connections << out
     out.callback { connections.delete(out) }
+    sleep 1
   end
 end
 
 post '/' do
-  connections.each { |out| out << "data: #{params[:msg]}\n\n" }
+  connections.each do |out|
+    out << "data: #{params[:msg]}\n\n"
+  rescue
+    out.close
+  end
   204 # response without entity body
 end
 

lib/sinatra/base.rb

@@ -477,8 +477,9 @@ def each(&front)
             @back.call(self)
           rescue Exception => e
             @scheduler.schedule { raise e }
+          ensure
+            close
           end
-          close unless @keep_open
         end
       end
 
@@ -509,7 +510,16 @@ def closed?
     def stream(keep_open = false)
       scheduler = env['async.callback'] ? EventMachine : Stream
       current   = @params.dup
-      body Stream.new(scheduler, keep_open) { |out| with_params(current) { yield(out) } }
+      block = if scheduler == Stream && keep_open
+        proc do |out|
+          until out.closed?
+            with_params(current) { yield(out) }
+          end
+        end
+      else
+        proc { |out| with_params(current) { yield(out) } }
+      end
+      body Stream.new(scheduler, keep_open, &block)
     end
 
     # Specify response freshness policy for HTTP caches (Cache-Control header).

test/integration/app.rb

@@ -33,9 +33,17 @@
   end
 end
 
+# the "sleep 1" in async routes makes the async tests work with puma
+# more in https://github.com/sinatra/sinatra/pull/1853
+
 set :out, nil
 get '/async' do
-  stream(:keep_open) { |o| (settings.out = o) << "hi!" }
+  stream(:keep_open) do |out_stream|
+    (settings.out = out_stream) << "hi!"
+    sleep 1
+  rescue
+    out_stream.close
+  end
 end
 
 get '/send' do
@@ -64,7 +72,12 @@
 class Subclass < Sinatra::Base
   set :out, nil
   get '/subclass/async' do
-    stream(:keep_open) { |o| (settings.out = o) << "hi!" }
+    stream(:keep_open) do |out_stream|
+      (settings.out = out_stream) << "hi!"
+      sleep 1
+    rescue
+      out_stream.close
+    end
   end
 
   get '/subclass/send' do