athenad: check end_event while uploading files (#31541)

* check end_event while uploading, throw abort exception if we need to shut down/restart

* fix

* draft test

stash

* Revert - there's no easy way to know if it breaks early in upload loop
or not yet

This reverts commit ad893687e196ebb31d276a114c19e9af963ed02a.

* todo for now

selfdrive/athena/athenad.py

@@ -206,13 +206,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 +256,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)

selfdrive/athena/tests/test_athenad.py

@@ -233,6 +233,7 @@ class TestAthenadMethods(unittest.TestCase):
     time.sleep(0.1)
 
     # TODO: verify that upload actually succeeded
+    # TODO: also check that end_event and metered network raises AbortTransferException
     self.assertEqual(athenad.upload_queue.qsize(), 0)
 
   @parameterized.expand([(500, True), (412, False)])