Add support to keep open streaming connections with Puma

sinatra · Jan 2, 2023 · cf969dc · cf969dc
1 parent 9b6ba81
commit cf969dc
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 11 deletions.
diff --git a/examples/chat.rb b/examples/chat.rb
@@ -1,11 +1,13 @@
 #!/usr/bin/env ruby -I ../lib -I lib
 # frozen_string_literal: true
 
-require_relative 'rainbows'
+# This example does *not* work properly with WEBrick or other
+# servers that buffer output. To shut down the server, close any
+# open browser tabs that are connected to the chat server.
 
 require 'sinatra'
-set :server, :rainbows
-connections = []
+set :server, :puma
+connections = Set.new
 
 get '/' do
   halt erb(:login) unless params[:user]
@@ -14,13 +16,22 @@
 
 get '/stream', provides: 'text/event-stream' do
   stream :keep_open do |out|
-    connections << out
-    out.callback { connections.delete(out) }
+    if connections.add?(out)
+      out.callback { connections.delete(out) }
+    end
+    out << "heartbeat:\n"
+    sleep 1
+  rescue
+    out.close
   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
 
@@ -37,10 +48,10 @@
 </html>
 
 @@ login
-<form action='/'>
+<form action="/">
   <label for='user'>User Name:</label>
-  <input name='user' value='' />
-  <input type='submit' value="GO!" />
+  <input name="user" value="" />
+  <input type="submit" value="GO!" />
 </form>
 
 @@ chat

diff --git a/lib/sinatra/base.rb b/lib/sinatra/base.rb
@@ -474,8 +474,9 @@ def each(&front)
             @back.call(self)
           rescue Exception => e
             @scheduler.schedule { raise e }
+          ensure
+            close unless @keep_open
           end
-          close unless @keep_open
         end
       end
 
@@ -506,7 +507,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) } }
+      stream = if scheduler == Stream  && keep_open
+        Stream.new(scheduler, false) do |out|
+          until out.closed?
+            with_params(current) { yield(out) }
+          end
+        end
+      else
+        Stream.new(scheduler, keep_open) { |out| with_params(current) { yield(out) } }
+      end
+      body stream
     end
 
     # Specify response freshness policy for HTTP caches (Cache-Control header).

diff --git a/test/streaming_test.rb b/test/streaming_test.rb
@@ -133,6 +133,7 @@ def close.errback; end
         env['async.close'] = close
         stream(:keep_open) do |out|
           out.callback { ran = true }
+          out.close
         end
       end
     end