From 8c057350529ca697e8c163f8bc664a9a83b30d40 Mon Sep 17 00:00:00 2001
From: albertkun <albertk@gmx.com>
Date: Mon, 20 Nov 2023 17:35:17 -0800
Subject: [PATCH] feat: WebSocket ping functionality feat: dockerfile uvicorn
 only, gunicorn removed

---
 fastapi/Dockerfile  |  6 +-----
 fastapi/app/main.py | 51 ++++++++++++++++++++++++++++++---------------
 2 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/fastapi/Dockerfile b/fastapi/Dockerfile
index a9bc39a..23b0b24 100644
--- a/fastapi/Dockerfile
+++ b/fastapi/Dockerfile
@@ -25,9 +25,6 @@ COPY . .
 # Create a new user and switch to that user
 RUN useradd -ms /bin/bash pgbouncer && \
     chown -R pgbouncer:pgbouncer /etc/pgbouncer
-
-RUN pip install gevent
-
 CMD python -c "\
 import os;\
 from urllib.parse import urlparse;\
@@ -48,6 +45,5 @@ print(f'export DB_NAME={result.path[1:]}')" > /tmp/env.sh && \
     echo "auth_file = /etc/pgbouncer/userlist.txt" >> /etc/pgbouncer/pgbouncer.ini && \
     echo "client_login_timeout = 120" >> /etc/pgbouncer/pgbouncer.ini && \
     gosu pgbouncer pgbouncer /etc/pgbouncer/pgbouncer.ini & \
-    gunicorn app.main:app --workers 3 --worker-class gevent --bind 0.0.0.0:80 --timeout 120 --keep-alive 120
-
+    exec uvicorn app.main:app --host 0.0.0.0 --port 80 --workers 2
 EXPOSE 80
\ No newline at end of file
diff --git a/fastapi/app/main.py b/fastapi/app/main.py
index dcad8c3..3dcd6cf 100644
--- a/fastapi/app/main.py
+++ b/fastapi/app/main.py
@@ -18,7 +18,7 @@
 
 from datetime import timedelta, date, datetime
 
-from fastapi import FastAPI, Request, Response, Depends, HTTPException, status, Query, Websocket
+from fastapi import FastAPI, Request, Response, Depends, HTTPException, status, Query, Websocket,WebSocketDisconnect
 from fastapi import Path as FastAPIPath
 # from fastapi import FastAPI, Request, Response, Depends, HTTPException, status
 from fastapi.encoders import jsonable_encoder
@@ -398,31 +398,48 @@ async def get_list_of_field_values(agency_id: AgencyIdEnum, field: VehiclePositi
         raise HTTPException(status_code=404, detail="Data not found")
     return data
 
-
 @app.websocket("/ws/{agency_id}/vehicle_positions")
 async def websocket_endpoint(websocket: WebSocket, agency_id: str, async_db: AsyncSession = Depends(get_async_db)):
     await websocket.accept()
-    while True:
-        data = await crud.get_all_data_async(async_db, models.VehiclePositions, agency_id)
-        if data is not None:
-            await websocket.send_json(data)
-        await asyncio.sleep(5)
+    try:
+        while True:
+            try:
+                data = await asyncio.wait_for(crud.get_all_data_async(async_db, models.VehiclePositions, agency_id), timeout=120)
+                if data is not None:
+                    await websocket.send_json(data)
+                await asyncio.sleep(10)
+                # Send a ping every 10 seconds
+                await websocket.send_json({"type": "ping"})
+            except asyncio.TimeoutError:
+                raise HTTPException(status_code=408, detail="Request timed out")
+    except WebSocketDisconnect:
+        # Handle the WebSocket disconnect event
+        print("WebSocket disconnected")
+
 
 @app.websocket("/ws/{agency_id}/vehicle_positions/{field}/{ids}")
 async def websocket_vehicle_positions_by_ids(websocket: WebSocket, agency_id: AgencyIdEnum, field: VehiclePositionsFieldsEnum, ids: str, async_db: AsyncSession = Depends(get_async_db)):
     await websocket.accept()
     model = models.VehiclePositions
     ids = ids.split(',')
-    while True:
-        data = {}
-        for id in ids:
-            result = await crud.get_data_async(async_db, model, agency_id.value, field.value, id)
-            if result is not None:
-                data[id] = result
-        if data:
-            await websocket.send_json(data)
-        await asyncio.sleep(5)
-
+    try:
+        while True:
+            data = {}
+            for id in ids:
+                try:
+                    result = await asyncio.wait_for(crud.get_data_async(async_db, model, agency_id.value, field.value, id), timeout=120)
+                    if result is not None:
+                        data[id] = result
+                except asyncio.TimeoutError:
+                    raise HTTPException(status_code=408, detail="Request timed out")
+            if data:
+                await websocket.send_json(data)
+            await asyncio.sleep(5)
+            # Send a ping every 5 seconds
+            await websocket.send_json({"type": "ping"})
+    except WebSocketDisconnect:
+        # Handle the WebSocket disconnect event
+        print("WebSocket disconnected")
 
 ##### todo: Needs to be tested