@@ -21,24 +21,28 @@ String _sseHeaders(String origin) => 'HTTP/1.1 200 OK\r\n'
2121
2222/// A bi-directional SSE connection between server and browser.
2323class SseConnection extends StreamChannelMixin <String > {
24+ /// Incoming messages from the Browser client.
2425 final _incomingController = StreamController <String >();
26+
27+ /// Outgoing messages to the Browser client.
2528 final _outgoingController = StreamController <String >();
26- final _closeCompleter = Completer < Null >();
29+
2730 final Sink _sink;
28- final String _clientId;
2931
30- SseConnection (this ._sink, this ._clientId) {
32+ final _closedCompleter = Completer <void >();
33+
34+ SseConnection (this ._sink) {
3135 _outgoingController.stream.listen ((data) {
32- if (! _closeCompleter .isCompleted) {
36+ if (! _closedCompleter .isCompleted) {
3337 // JSON encode the message to escape new lines.
3438 _sink.add ('data: ${json .encode (data )}\n ' );
3539 _sink.add ('\n ' );
3640 }
3741 });
42+ _outgoingController.onCancel = _close;
43+ _incomingController.onCancel = _close;
3844 }
3945
40- Future get onClose => _closeCompleter.future;
41-
4246 /// The message added to the sink has to be JSON encodable.
4347 @override
4448 StreamSink <String > get sink => _outgoingController.sink;
@@ -50,8 +54,13 @@ class SseConnection extends StreamChannelMixin<String> {
5054 @override
5155 Stream <String > get stream => _incomingController.stream;
5256
53- void close () {
54- if (! _closeCompleter.isCompleted) _closeCompleter.complete ();
57+ void _close () {
58+ if (! _closedCompleter.isCompleted) {
59+ _closedCompleter.complete ();
60+ _sink.close ();
61+ if (! _outgoingController.isClosed) _outgoingController.close ();
62+ if (! _incomingController.isClosed) _incomingController.close ();
63+ }
5564 }
5665}
5766
@@ -63,15 +72,15 @@ class SseConnection extends StreamChannelMixin<String> {
6372class SseHandler {
6473 final _logger = Logger ('SseHandler' );
6574 final Uri _uri;
66-
67- final Set <SseConnection > _connections = Set <SseConnection >();
68-
75+ final _connections = < String , SseConnection > {};
6976 final _connectionController = StreamController <SseConnection >();
7077
78+ StreamQueue <SseConnection > _connectionsStream;
79+
7180 SseHandler (this ._uri);
7281
7382 StreamQueue <SseConnection > get connections =>
74- StreamQueue (_connectionController.stream);
83+ _connectionsStream ?? = StreamQueue (_connectionController.stream);
7584
7685 shelf.Handler get handler => _handle;
7786
@@ -82,19 +91,22 @@ class SseHandler {
8291 var sink = utf8.encoder.startChunkedConversion (channel.sink);
8392 sink.add (_sseHeaders (req.headers['origin' ]));
8493 var clientId = req.url.queryParameters['sseClientId' ];
85- var connection = SseConnection (sink, clientId );
86- _connections. add ( connection) ;
87- unawaited (connection.onClose .then ((_) {
88- _connections.remove (connection );
94+ var connection = SseConnection (sink);
95+ _connections[clientId] = connection;
96+ unawaited (connection._closedCompleter.future .then ((_) {
97+ _connections.remove (clientId );
8998 }));
99+ // Remove connection when it is remotely closed or the stream is
100+ // cancelled.
90101 channel.stream.listen ((_) {
91102 // SSE is unidirectional. Responses are handled through POST requests.
92103 }, onDone: () {
93- connection.close ();
104+ connection._close ();
94105 });
106+
95107 _connectionController.add (connection);
96108 });
97- return null ;
109+ return shelf. Response . notFound ( '' ) ;
98110 }
99111
100112 String _getOriginalPath (shelf.Request req) => req.requestedUri.path;
@@ -122,11 +134,7 @@ class SseHandler {
122134 var clientId = req.url.queryParameters['sseClientId' ];
123135 var message = await req.readAsString ();
124136 var jsonObject = json.decode (message) as String ;
125- for (var connection in _connections) {
126- if (connection._clientId == clientId) {
127- connection._incomingController.add (jsonObject);
128- }
129- }
137+ _connections[clientId]? ._incomingController? .add (jsonObject);
130138 } catch (e, st) {
131139 _logger.fine ('Failed to handle incoming message. $e $st ' );
132140 }
0 commit comments