wip

selfdrive/athena/athenad.py

@@ -19,7 +19,8 @@ from dataclasses import asdict, dataclass, replace
 from datetime import datetime
 from functools import partial
 from queue import Queue
-from typing import Callable, Dict, List, Optional, Set, Union, cast
+from typing import cast
+from collections.abc import Callable
 
 import requests
 from jsonrpc import JSONRPCResponseManager, dispatcher
@@ -55,17 +56,17 @@ WS_FRAME_SIZE = 4096
 
 NetworkType = log.DeviceState.NetworkType
 
-UploadFileDict = Dict[str, Union[str, int, float, bool]]
-UploadItemDict = Dict[str, Union[str, bool, int, float, Dict[str, str]]]
+UploadFileDict = dict[str, str | int | float | bool]
+UploadItemDict = dict[str, str | bool | int | float | dict[str, str]]
 
-UploadFilesToUrlResponse = Dict[str, Union[int, List[UploadItemDict], List[str]]]
+UploadFilesToUrlResponse = dict[str, int | list[UploadItemDict] | list[str]]
 
 
 @dataclass
 class UploadFile:
   fn: str
   url: str
-  headers: Dict[str, str]
+  headers: dict[str, str]
   allow_cellular: bool
 
   @classmethod
@@ -77,9 +78,9 @@ class UploadFile:
 class UploadItem:
   path: str
   url: str
-  headers: Dict[str, str]
+  headers: dict[str, str]
   created_at: int
-  id: Optional[str]
+  id: str | None
   retry_count: int = 0
   current: bool = False
   progress: float = 0
@@ -97,9 +98,9 @@ send_queue: Queue[str] = queue.Queue()
 upload_queue: Queue[UploadItem] = queue.Queue()
 low_priority_send_queue: Queue[str] = queue.Queue()
 log_recv_queue: Queue[str] = queue.Queue()
-cancelled_uploads: Set[str] = set()
+cancelled_uploads: set[str] = set()
 
-cur_upload_items: Dict[int, Optional[UploadItem]] = {}
+cur_upload_items: dict[int, UploadItem | None] = {}
 
 
 def strip_bz2_extension(fn: str) -> str:
@@ -127,14 +128,14 @@ class UploadQueueCache:
   @staticmethod
   def cache(upload_queue: Queue[UploadItem]) -> None:
     try:
-      queue: List[Optional[UploadItem]] = list(upload_queue.queue)
+      queue: list[UploadItem | None] = list(upload_queue.queue)
       items = [asdict(i) for i in queue if i is not None and (i.id not in cancelled_uploads)]
       Params().put("AthenadUploadQueue", json.dumps(items))
     except Exception:
       cloudlog.exception("athena.UploadQueueCache.cache.exception")
 
 
-def handle_long_poll(ws: WebSocket, exit_event: Optional[threading.Event]) -> None:
+def handle_long_poll(ws: WebSocket, exit_event: threading.Event | None) -> None:
   end_event = threading.Event()
 
   threads = [
@@ -206,13 +207,17 @@ def retry_upload(tid: int, end_event: threading.Event, increase_count: bool = Tr
         break
 
 
-def cb(sm, item, tid, sz: int, cur: int) -> None:
+def cb(sm, item, tid, end_event: threading.Event, sz: int, cur: int) -> None:
   # Abort transfer if connection changed to metered after starting upload
+  # or if athenad is shutting down to re-connect the websocket
   sm.update(0)
   metered = sm['deviceState'].networkMetered
   if metered and (not item.allow_cellular):
     raise AbortTransferException
 
+  if end_event.is_set():
+    raise AbortTransferException
+
   cur_upload_items[tid] = replace(item, progress=cur / sz if sz else 1)
 
 
@@ -252,7 +257,7 @@ def upload_handler(end_event: threading.Event) -> None:
           sz = -1
 
         cloudlog.event("athena.upload_handler.upload_start", fn=fn, sz=sz, network_type=network_type, metered=metered, retry_count=item.retry_count)
-        response = _do_upload(item, partial(cb, sm, item, tid))
+        response = _do_upload(item, partial(cb, sm, item, tid, end_event))
 
         if response.status_code not in (200, 201, 401, 403, 412):
           cloudlog.event("athena.upload_handler.retry", status_code=response.status_code, fn=fn, sz=sz, network_type=network_type, metered=metered)
@@ -274,7 +279,7 @@ def upload_handler(end_event: threading.Event) -> None:
       cloudlog.exception("athena.upload_handler.exception")
 
 
-def _do_upload(upload_item: UploadItem, callback: Optional[Callable] = None) -> requests.Response:
+def _do_upload(upload_item: UploadItem, callback: Callable = None) -> requests.Response:
   path = upload_item.path
   compress = False
 
@@ -313,7 +318,7 @@ def getMessage(service: str, timeout: int = 1000) -> dict:
 
 
 @dispatcher.add_method
-def getVersion() -> Dict[str, str]:
+def getVersion() -> dict[str, str]:
   return {
     "version": get_version(),
     "remote": get_normalized_origin(),
@@ -323,7 +328,7 @@ def getVersion() -> Dict[str, str]:
 
 
 @dispatcher.add_method
-def setNavDestination(latitude: int = 0, longitude: int = 0, place_name: Optional[str] = None, place_details: Optional[str] = None) -> Dict[str, int]:
+def setNavDestination(latitude: int = 0, longitude: int = 0, place_name: str = None, place_details: str = None) -> dict[str, int]:
   destination = {
     "latitude": latitude,
     "longitude": longitude,
@@ -335,7 +340,7 @@ def setNavDestination(latitude: int = 0, longitude: int = 0, place_name: Optiona
   return {"success": 1}
 
 
-def scan_dir(path: str, prefix: str) -> List[str]:
+def scan_dir(path: str, prefix: str) -> list[str]:
   files = []
   # only walk directories that match the prefix
   # (glob and friends traverse entire dir tree)
@@ -355,12 +360,12 @@ def scan_dir(path: str, prefix: str) -> List[str]:
   return files
 
 @dispatcher.add_method
-def listDataDirectory(prefix='') -> List[str]:
+def listDataDirectory(prefix='') -> list[str]:
   return scan_dir(Paths.log_root(), prefix)
 
 
 @dispatcher.add_method
-def uploadFileToUrl(fn: str, url: str, headers: Dict[str, str]) -> UploadFilesToUrlResponse:
+def uploadFileToUrl(fn: str, url: str, headers: dict[str, str]) -> UploadFilesToUrlResponse:
   # this is because mypy doesn't understand that the decorator doesn't change the return type
   response: UploadFilesToUrlResponse = uploadFilesToUrls([{
     "fn": fn,
@@ -371,11 +376,11 @@ def uploadFileToUrl(fn: str, url: str, headers: Dict[str, str]) -> UploadFilesTo
 
 
 @dispatcher.add_method
-def uploadFilesToUrls(files_data: List[UploadFileDict]) -> UploadFilesToUrlResponse:
+def uploadFilesToUrls(files_data: list[UploadFileDict]) -> UploadFilesToUrlResponse:
   files = map(UploadFile.from_dict, files_data)
 
-  items: List[UploadItemDict] = []
-  failed: List[str] = []
+  items: list[UploadItemDict] = []
+  failed: list[str] = []
   for file in files:
     if len(file.fn) == 0 or file.fn[0] == '/' or '..' in file.fn or len(file.url) == 0:
       failed.append(file.fn)
@@ -414,13 +419,13 @@ def uploadFilesToUrls(files_data: List[UploadFileDict]) -> UploadFilesToUrlRespo
 
 
 @dispatcher.add_method
-def listUploadQueue() -> List[UploadItemDict]:
+def listUploadQueue() -> list[UploadItemDict]:
   items = list(upload_queue.queue) + list(cur_upload_items.values())
   return [asdict(i) for i in items if (i is not None) and (i.id not in cancelled_uploads)]
 
 
 @dispatcher.add_method
-def cancelUpload(upload_id: Union[str, List[str]]) -> Dict[str, Union[int, str]]:
+def cancelUpload(upload_id: str | list[str]) -> dict[str, int | str]:
   if not isinstance(upload_id, list):
     upload_id = [upload_id]
 
@@ -433,7 +438,7 @@ def cancelUpload(upload_id: Union[str, List[str]]) -> Dict[str, Union[int, str]]
   return {"success": 1}
 
 @dispatcher.add_method
-def setRouteViewed(route: str) -> Dict[str, Union[int, str]]:
+def setRouteViewed(route: str) -> dict[str, int | str]:
   # maintain a list of the last 10 routes viewed in connect
   params = Params()
 
@@ -448,7 +453,7 @@ def setRouteViewed(route: str) -> Dict[str, Union[int, str]]:
   return {"success": 1}
 
 
-def startLocalProxy(global_end_event: threading.Event, remote_ws_uri: str, local_port: int) -> Dict[str, int]:
+def startLocalProxy(global_end_event: threading.Event, remote_ws_uri: str, local_port: int) -> dict[str, int]:
   try:
     if local_port not in LOCAL_PORT_WHITELIST:
       raise Exception("Requested local port not whitelisted")
@@ -482,7 +487,7 @@ def startLocalProxy(global_end_event: threading.Event, remote_ws_uri: str, local
 
 
 @dispatcher.add_method
-def getPublicKey() -> Optional[str]:
+def getPublicKey() -> str | None:
   if not os.path.isfile(Paths.persist_root() + '/comma/id_rsa.pub'):
     return None
 
@@ -522,7 +527,7 @@ def getNetworks():
 
 
 @dispatcher.add_method
-def takeSnapshot() -> Optional[Union[str, Dict[str, str]]]:
+def takeSnapshot() -> str | dict[str, str] | None:
   from openpilot.system.camerad.snapshot.snapshot import jpeg_write, snapshot
   ret = snapshot()
   if ret is not None:
@@ -539,7 +544,7 @@ def takeSnapshot() -> Optional[Union[str, Dict[str, str]]]:
     raise Exception("not available while camerad is started")
 
 
-def get_logs_to_send_sorted() -> List[str]:
+def get_logs_to_send_sorted() -> list[str]:
   # TODO: scan once then use inotify to detect file creation/deletion
   curr_time = int(time.time())
   logs = []
@@ -746,6 +751,9 @@ def ws_manage(ws: WebSocket, end_event: threading.Event) -> None:
       onroad_prev = onroad
 
       if sock is not None:
+        # While not sending data, onroad, we can expect to time out in 7 + (7 * 2) = 21s
+        #                         offroad, we can expect to time out in 30 + (10 * 3) = 60s
+        # FIXME: TCP_USER_TIMEOUT is effectively 2x for some reason (32s), so it's mostly unused
         sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT, 16000 if onroad else 0)
         sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 7 if onroad else 30)
         sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 7 if onroad else 10)
@@ -759,7 +767,7 @@ def backoff(retries: int) -> int:
   return random.randrange(0, min(128, int(2 ** retries)))
 
 
-def main(exit_event: Optional[threading.Event] = None):
+def main(exit_event: threading.Event = None):
   try:
     set_core_affinity([0, 1, 2, 3])
   except Exception: