Implement SSE endpoints in Tornado #3

server/handlers/api/spots.py

@@ -1,9 +1,11 @@
 import json
 import logging
 from datetime import datetime, timedelta
+from queue import Queue
 
 import pytz
 import tornado
+import tornado_eventsource.handler
 
 from core.prometheus_metrics_handler import api_requests_counter
 from core.utils import serialize_everything
@@ -44,28 +46,55 @@ class APISpotsHandler(tornado.web.RequestHandler):
 
 
 # API request handler for /api/v1/spots/stream
-class APISpotsStreamHandler(tornado.web.RequestHandler):
-    def get(self):
-        # todo
-        # try:
-        #     # Metrics
-        #     api_requests_counter.inc()
-        #
-        #     response.content_type  = 'text/event-stream'
-        #     response.cache_control = 'no-cache'
-        #     yield 'retry: 1000\n\n'
-        #
-        #     spot_queue = Queue(maxsize=100)
-        #     self.sse_spot_queues.append(spot_queue)
-        #     while True:
-        #         if spot_queue.empty():
-        #             gevent.sleep(1)
-        #         else:
-        #             spot = spot_queue.get()
-        #             yield 'data: ' + json.dumps(spot, default=serialize_everything) + '\n\n'
-        # except Exception as e:
-        #     logging.warn("Exception when serving SSE socket", e)
-        pass
+class APISpotsStreamHandler(tornado_eventsource.handler.EventSourceHandler):
+    def initialize(self, sse_spot_queues, web_server_metrics):
+        self.sse_spot_queues = sse_spot_queues
+        self.web_server_metrics = web_server_metrics
+
+    # Called once on the client opening a connection, set things up
+    def open(self):
+        try:
+            # Metrics
+            self.web_server_metrics["last_api_access_time"] = datetime.now(pytz.UTC)
+            self.web_server_metrics["api_access_counter"] += 1
+            self.web_server_metrics["status"] = "OK"
+            api_requests_counter.inc()
+
+            # Create a spot queue and add it to the web server's list. The web server will fill this when spots arrive
+            self.spot_queue = Queue(maxsize=1000)
+            self.sse_spot_queues.append(self.spot_queue)
+
+            # Set up a timed callback to check if anything is in the queue
+            self.heartbeat = tornado.ioloop.PeriodicCallback(self._callback, 1000)
+            self.heartbeat.start()
+
+        except Exception as e:
+            logging.warn("Exception when serving SSE socket", e)
+
+    # When the user closes the socket, empty our queue and remove it from the list so the server no longer fills it
+    def close(self):
+        try:
+            if self.spot_queue in self.sse_spot_queues:
+                self.sse_spot_queues.remove(self.spot_queue)
+            self.spot_queue.empty()
+        except:
+            pass
+        self.spot_queue = None
+        super().close()
+
+    # Callback to check if anything has arrived in the queue, and if so send it to the client
+    def _callback(self):
+        try:
+            if self.spot_queue:
+                while not self.spot_queue.empty():
+                    spot = self.spot_queue.get()
+                    self.write_message(msg=json.dumps(spot, default=serialize_everything))
+                if self.spot_queue not in self.sse_spot_queues:
+                    logging.error("Web server cleared up a queue of an active connection!")
+                    self.close()
+        except:
+            logging.warn("Exception in SSE callback, connection will be closed.")
+            self.close()