/
live.rb
390 lines (346 loc) · 12.8 KB
/
live.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
# frozen_string_literal: true
# :markup: markdown
require "action_dispatch/http/response"
require "delegate"
require "active_support/json"
module ActionController
# # Action Controller Live
#
# Mix this module into your controller, and all actions in that controller will
# be able to stream data to the client as it's written.
#
# class MyController < ActionController::Base
# include ActionController::Live
#
# def stream
# response.headers['Content-Type'] = 'text/event-stream'
# 100.times {
# response.stream.write "hello world\n"
# sleep 1
# }
# ensure
# response.stream.close
# end
# end
#
# There are a few caveats with this module. You **cannot** write headers after
# the response has been committed (Response#committed? will return truthy).
# Calling `write` or `close` on the response stream will cause the response
# object to be committed. Make sure all headers are set before calling write or
# close on your stream.
#
# You **must** call close on your stream when you're finished, otherwise the
# socket may be left open forever.
#
# The final caveat is that your actions are executed in a separate thread than
# the main thread. Make sure your actions are thread safe, and this shouldn't be
# a problem (don't share state across threads, etc).
#
# Note that Rails includes `Rack::ETag` by default, which will buffer your
# response. As a result, streaming responses may not work properly with Rack
# 2.2.x, and you may need to implement workarounds in your application. You can
# either set the `ETag` or `Last-Modified` response headers or remove
# `Rack::ETag` from the middleware stack to address this issue.
#
# Here's an example of how you can set the `Last-Modified` header if your Rack
# version is 2.2.x:
#
# def stream
# response.headers["Content-Type"] = "text/event-stream"
# response.headers["Last-Modified"] = Time.now.httpdate # Add this line if your Rack version is 2.2.x
# ...
# end
module Live
extend ActiveSupport::Concern
module ClassMethods
def make_response!(request)
if request.get_header("HTTP_VERSION") == "HTTP/1.0"
super
else
Live::Response.new.tap do |res|
res.request = request
end
end
end
end
# # Action Controller Live Server Sent Events
#
# This class provides the ability to write an SSE (Server Sent Event) to an IO
# stream. The class is initialized with a stream and can be used to either write
# a JSON string or an object which can be converted to JSON.
#
# Writing an object will convert it into standard SSE format with whatever
# options you have configured. You may choose to set the following options:
#
# 1) Event. If specified, an event with this name will be dispatched on
# the browser.
# 2) Retry. The reconnection time in milliseconds used when attempting
# to send the event.
# 3) Id. If the connection dies while sending an SSE to the browser, then
# the server will receive a +Last-Event-ID+ header with value equal to +id+.
#
# After setting an option in the constructor of the SSE object, all future SSEs
# sent across the stream will use those options unless overridden.
#
# Example Usage:
#
# class MyController < ActionController::Base
# include ActionController::Live
#
# def index
# response.headers['Content-Type'] = 'text/event-stream'
# sse = SSE.new(response.stream, retry: 300, event: "event-name")
# sse.write({ name: 'John'})
# sse.write({ name: 'John'}, id: 10)
# sse.write({ name: 'John'}, id: 10, event: "other-event")
# sse.write({ name: 'John'}, id: 10, event: "other-event", retry: 500)
# ensure
# sse.close
# end
# end
#
# Note: SSEs are not currently supported by IE. However, they are supported by
# Chrome, Firefox, Opera, and Safari.
class SSE
PERMITTED_OPTIONS = %w( retry event id )
def initialize(stream, options = {})
@stream = stream
@options = options
end
def close
@stream.close
end
def write(object, options = {})
case object
when String
perform_write(object, options)
else
perform_write(ActiveSupport::JSON.encode(object), options)
end
end
private
def perform_write(json, options)
current_options = @options.merge(options).stringify_keys
PERMITTED_OPTIONS.each do |option_name|
if (option_value = current_options[option_name])
@stream.write "#{option_name}: #{option_value}\n"
end
end
message = json.gsub("\n", "\ndata: ")
@stream.write "data: #{message}\n\n"
end
end
class ClientDisconnected < RuntimeError
end
class Buffer < ActionDispatch::Response::Buffer # :nodoc:
include MonitorMixin
class << self
attr_accessor :queue_size
end
@queue_size = 10
# Ignore that the client has disconnected.
#
# If this value is `true`, calling `write` after the client disconnects will
# result in the written content being silently discarded. If this value is
# `false` (the default), a ClientDisconnected exception will be raised.
attr_accessor :ignore_disconnect
def initialize(response)
super(response, build_queue(self.class.queue_size))
@error_callback = lambda { true }
@cv = new_cond
@aborted = false
@ignore_disconnect = false
end
# ActionDispatch::Response delegates #to_ary to the internal
# ActionDispatch::Response::Buffer, defining #to_ary is an indicator that the
# response body can be buffered and/or cached by Rack middlewares, this is not
# the case for Live responses so we undefine it for this Buffer subclass.
undef_method :to_ary
def write(string)
unless @response.committed?
@response.headers["Cache-Control"] ||= "no-cache"
@response.delete_header "Content-Length"
end
super
unless connected?
@buf.clear
unless @ignore_disconnect
# Raise ClientDisconnected, which is a RuntimeError (not an IOError), because
# that's more appropriate for something beyond the developer's control.
raise ClientDisconnected, "client disconnected"
end
end
end
# Same as `write` but automatically include a newline at the end of the string.
def writeln(string)
write string.end_with?("\n") ? string : "#{string}\n"
end
# Write a 'close' event to the buffer; the producer/writing thread uses this to
# notify us that it's finished supplying content.
#
# See also #abort.
def close
synchronize do
super
@buf.push nil
@cv.broadcast
end
end
# Inform the producer/writing thread that the client has disconnected; the
# reading thread is no longer interested in anything that's being written.
#
# See also #close.
def abort
synchronize do
@aborted = true
@buf.clear
end
end
# Is the client still connected and waiting for content?
#
# The result of calling `write` when this is `false` is determined by
# `ignore_disconnect`.
def connected?
!@aborted
end
def on_error(&block)
@error_callback = block
end
def call_on_error
@error_callback.call
end
private
def each_chunk(&block)
loop do
str = nil
ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
str = @buf.pop
end
break unless str
yield str
end
end
def build_queue(queue_size)
queue_size ? SizedQueue.new(queue_size) : Queue.new
end
end
class Response < ActionDispatch::Response # :nodoc: all
private
def before_committed
super
jar = request.cookie_jar
# The response can be committed multiple times
jar.write self unless committed?
end
def build_buffer(response, body)
buf = Live::Buffer.new response
body.each { |part| buf.write part }
buf
end
end
def process(name)
t1 = Thread.current
locals = t1.keys.map { |key| [key, t1[key]] }
error = nil
# This processes the action in a child thread. It lets us return the response
# code and headers back up the Rack stack, and still process the body in
# parallel with sending data to the client.
new_controller_thread {
ActiveSupport::Dependencies.interlock.running do
t2 = Thread.current
# Since we're processing the view in a different thread, copy the thread locals
# from the main thread to the child thread. :'(
locals.each { |k, v| t2[k] = v }
ActiveSupport::IsolatedExecutionState.share_with(t1)
begin
super(name)
rescue => e
if @_response.committed?
begin
@_response.stream.write(ActionView::Base.streaming_completion_on_exception) if request.format == :html
@_response.stream.call_on_error
rescue => exception
log_error(exception)
ensure
log_error(e)
@_response.stream.close
end
else
error = e
end
ensure
@_response.commit!
end
end
}
ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
@_response.await_commit
end
raise error if error
end
def response_body=(body)
super
response.close if response
end
# Sends a stream to the browser, which is helpful when you're generating exports
# or other running data where you don't want the entire file buffered in memory
# first. Similar to send_data, but where the data is generated live.
#
# Options:
# * `:filename` - suggests a filename for the browser to use.
# * `:type` - specifies an HTTP content type. You can specify either a string
# or a symbol for a registered type with `Mime::Type.register`, for example
# :json. If omitted, type will be inferred from the file extension specified
# in `:filename`. If no content type is registered for the extension, the
# default type 'application/octet-stream' will be used.
# * `:disposition` - specifies whether the file will be shown inline or
# downloaded. Valid values are 'inline' and 'attachment' (default).
#
#
# Example of generating a csv export:
#
# send_stream(filename: "subscribers.csv") do |stream|
# stream.write "email_address,updated_at\n"
#
# @subscribers.find_each do |subscriber|
# stream.write "#{subscriber.email_address},#{subscriber.updated_at}\n"
# end
# end
def send_stream(filename:, disposition: "attachment", type: nil)
payload = { filename: filename, disposition: disposition, type: type }
ActiveSupport::Notifications.instrument("send_stream.action_controller", payload) do
response.headers["Content-Type"] =
(type.is_a?(Symbol) ? Mime[type].to_s : type) ||
Mime::Type.lookup_by_extension(File.extname(filename).downcase.delete("."))&.to_s ||
"application/octet-stream"
response.headers["Content-Disposition"] =
ActionDispatch::Http::ContentDisposition.format(disposition: disposition, filename: filename)
yield response.stream
end
ensure
response.stream.close
end
private
# Spawn a new thread to serve up the controller in. This is to get around the
# fact that Rack isn't based around IOs and we need to use a thread to stream
# data from the response bodies. Nobody should call this method except in Rails
# internals. Seriously!
def new_controller_thread # :nodoc:
Thread.new {
t2 = Thread.current
t2.abort_on_exception = true
yield
}
end
def log_error(exception)
logger = ActionController::Base.logger
return unless logger
logger.fatal do
message = +"\n#{exception.class} (#{exception.message}):\n"
message << exception.annotated_source_code.to_s if exception.respond_to?(:annotated_source_code)
message << " " << exception.backtrace.join("\n ")
"#{message}\n\n"
end
end
end
end