Skip to content

Commit

Permalink
refactor(ota_proxy): refine request handling flow, allow external_cac…
Browse files Browse the repository at this point in the history 
…he mount without cache_enabled=True, fix disk usage checker (#447)

This PR simplifies the request handling flow, instead of implementing the guard conditions checks in retrieve_file API, now the guard conditions checks are implemented within each _retrieve_file handlers.

Minor fix:
1. ota_proxy: fixing the background disk space usage checker thread exits unexpectedly.

Behavior changes:
1. ota_proxy external cache feature now is enabled without checking cache_enabled flag.
  • Loading branch information
@Bodong-Yang
Bodong-Yang authored Dec 2, 2024
1 parent f72f22f commit e6a93db
Showing 1 changed file with 108 additions and 89 deletions.
197 changes: 108 additions & 89 deletions src/ota_proxy/ota_cache.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import AsyncIterator, Dict, List, Mapping, Optional, Tuple
from typing import AsyncIterator, Mapping, Optional
from urllib.parse import SplitResult, quote, urlsplit

import aiohttp
Expand Down Expand Up @@ -139,11 +139,7 @@ def __init__(

self._external_cache_data_dir = None
self._external_cache_mp = None
if (
cache_enabled
and external_cache_mnt_point
and mount_external_cache(external_cache_mnt_point)
):
if external_cache_mnt_point and mount_external_cache(external_cache_mnt_point):
logger.info(
f"external cache source is enabled at: {external_cache_mnt_point}"
)
Expand Down Expand Up @@ -263,7 +259,7 @@ def _background_check_free_space(self):
self._storage_below_hard_limit_event.clear()
break

current_used_p = disk_usage.used / disk_usage.total * 100
current_used_p = int(disk_usage.used / disk_usage.total * 100)

_previous_below_hard = self._storage_below_hard_limit_event.is_set()
_current_below_hard = True
Expand All @@ -281,17 +277,17 @@ def _background_check_free_space(self):
# logging space monitoring result
if _previous_below_hard and not _current_below_hard:
logger.warning(
f"disk usage reached hard limit({current_used_p=:.1}%,"
f"{cfg.DISK_USE_LIMIT_HARD_P:.1}%), cache disabled"
f"disk usage reached hard limit({current_used_p=}%,"
f"{cfg.DISK_USE_LIMIT_HARD_P=}%), cache disabled"
)
elif not _previous_below_hard and _current_below_hard:
logger.info(
f"disk usage is below hard limit({current_used_p=:.1}%),"
f"{cfg.DISK_USE_LIMIT_SOFT_P:.1}%), cache enabled again"
f"disk usage is below hard limit({current_used_p=}%),"
f"{cfg.DISK_USE_LIMIT_SOFT_P=}%), cache enabled again"
)
time.sleep(cfg.DISK_USE_PULL_INTERVAL)

def _cache_entries_cleanup(self, entry_hashes: List[str]):
def _cache_entries_cleanup(self, entry_hashes: list[str]) -> None:
"""Cleanup entries indicated by entry_hashes list."""
for entry_hash in entry_hashes:
# remove cache entry
Expand Down Expand Up @@ -383,7 +379,7 @@ async def _retrieve_file_by_downloading(
raw_url: str,
*,
headers: Mapping[str, str],
) -> Tuple[AsyncIterator[bytes], CIMultiDictProxy[str]]:
) -> tuple[AsyncIterator[bytes], CIMultiDictProxy[str]]:
async def _do_request() -> AsyncIterator[bytes]:
async with self._session.get(
self._process_raw_url(raw_url),
Expand All @@ -407,28 +403,33 @@ async def _do_request() -> AsyncIterator[bytes]:
resp_headers: CIMultiDictProxy[str] = await (_remote_fd := _do_request()).__anext__() # type: ignore
return _remote_fd, resp_headers

async def _retrieve_file_by_cache(
self, cache_identifier: str, *, retry_cache: bool
async def _retrieve_file_by_cache_lookup(
self, *, raw_url: str, cache_policy: OTAFileCacheControl
) -> tuple[AsyncIterator[bytes], CIMultiDict[str]] | None:
"""
Returns:
A tuple of bytes iterator and headers dict for back to client.
"""
# cache file available, lookup the db for metadata
if (
not self._cache_enabled
or cache_policy.no_cache
or cache_policy.retry_caching
):
return

cache_identifier = cache_policy.file_sha256
if not cache_identifier:
# fallback to use URL based hash, and clear compression_alg for such case
cache_identifier = url_based_hash(raw_url)

meta_db_entry = await self._lru_helper.lookup_entry(cache_identifier)
if not meta_db_entry:
return

# NOTE: db_entry.file_sha256 can be either
# 1. valid sha256 value for corresponding plain uncompressed OTA file
# 2. URL based sha256 value for corresponding requested URL
# otaclient indicates that this cache entry is invalid, cleanup and exit
cache_file = self._base_dir / cache_identifier
if retry_cache:
logger.debug(f"requested with retry_cache: {meta_db_entry=}..")
await self._lru_helper.remove_entry(cache_identifier)
cache_file.unlink(missing_ok=True)
return

# check if cache file exists
# NOTE(20240729): there is an edge condition that the finished cached file is not yet renamed,
Expand All @@ -438,7 +439,6 @@ async def _retrieve_file_by_cache(
for _retry_count in range(_retry_count_max):
if cache_file.is_file():
break

await asyncio.sleep(get_backoff(_retry_count, _factor, _backoff_max))

if not cache_file.is_file():
Expand All @@ -460,7 +460,12 @@ async def _retrieve_file_by_external_cache(
self, client_cache_policy: OTAFileCacheControl
) -> tuple[AsyncIterator[bytes], CIMultiDict[str]] | None:
# skip if not external cache or otaclient doesn't sent valid file_sha256
if not self._external_cache_data_dir or not client_cache_policy.file_sha256:
if (
not self._external_cache_data_dir
or client_cache_policy.no_cache
or client_cache_policy.retry_caching
or not client_cache_policy.file_sha256
):
return

cache_identifier = client_cache_policy.file_sha256
Expand Down Expand Up @@ -488,83 +493,34 @@ async def _retrieve_file_by_external_cache(
)
return read_file(cache_file, executor=self._executor), _header

# exposed API

async def retrieve_file(
async def _retrieve_file_by_new_caching(
self,
*,
raw_url: str,
headers_from_client: Dict[str, str],
) -> tuple[AsyncIterator[bytes], CIMultiDict[str] | CIMultiDictProxy[str]] | None:
"""Retrieve a file descriptor for the requested <raw_url>.
This method retrieves a file descriptor for incoming client request.
Upper uvicorn app can use this file descriptor to yield chunks of data,
and send chunks to the on-calling ota_client.
NOTE: use raw_url in all operations, except opening remote file.
Args:
raw_url: unquoted raw url received from uvicorn
headers_from_client: headers come from client's request, which will be
passthrough to upper otaproxy and/or remote OTA image server.
Returns:
A tuple contains an asyncio generator for upper server app to yield data chunks from
and headers dict that should be sent back to client in resp.
"""
if self._closed:
raise BaseOTACacheError("ota cache pool is closed")

cache_policy = OTAFileCacheControl.parse_header(
headers_from_client.get(HEADER_OTA_FILE_CACHE_CONTROL, "")
)
if cache_policy.no_cache:
logger.info(f"client indicates that do not cache for {raw_url=}")

if not self._upper_proxy:
headers_from_client.pop(HEADER_OTA_FILE_CACHE_CONTROL, None)

# --- case 1: not using cache, directly download file --- #
if (
not self._cache_enabled # ota_proxy is configured to not cache anything
or cache_policy.no_cache # ota_client send request with no_cache policy
or not self._storage_below_hard_limit_event.is_set() # disable cache if space hardlimit is reached
):
logger.debug(
f"not use cache({self._cache_enabled=}, {cache_policy=}, "
f"{self._storage_below_hard_limit_event.is_set()=}): {raw_url=}"
)
return await self._retrieve_file_by_downloading(
raw_url, headers=headers_from_client
)

# --- case 2: if externel cache source available, try to use it --- #
# NOTE: if client requsts with retry_caching directive, it may indicate cache corrupted
# in external cache storage, in such case we should skip the use of external cache.
cache_policy: OTAFileCacheControl,
headers_from_client: dict[str, str],
) -> tuple[AsyncIterator[bytes], CIMultiDictProxy[str] | CIMultiDict[str]] | None:
# NOTE(20241202): no new cache on hard limit being reached
if (
self._external_cache_data_dir
and not cache_policy.retry_caching
and (_res := await self._retrieve_file_by_external_cache(cache_policy))
not self._cache_enabled
or cache_policy.no_cache
or not self._storage_below_hard_limit_event.is_set()
):
return _res
return

# pre-calculated cache_identifier and corresponding compression_alg
cache_identifier = cache_policy.file_sha256
compression_alg = cache_policy.file_compression_alg

# fallback to use URL based hash, and clear compression_alg
if not cache_identifier:
# fallback to use URL based hash, and clear compression_alg for such case
cache_identifier = url_based_hash(raw_url)
compression_alg = ""

# --- case 3: try to use local cache --- #
if _res := await self._retrieve_file_by_cache(
cache_identifier, retry_cache=cache_policy.retry_caching
):
return _res
# if set, cleanup any previous cache file before starting new cache
if cache_policy.retry_caching:
logger.debug(f"requested with retry_cache for {raw_url=} ...")
await self._lru_helper.remove_entry(cache_identifier)
(self._base_dir / cache_identifier).unlink(missing_ok=True)

# --- case 4: no cache available, streaming remote file and cache --- #
# a online tracker is available for this requrest
if (tracker := self._on_going_caching.get_tracker(cache_identifier)) and (
subscription := await tracker.subscribe_tracker()
):
Expand Down Expand Up @@ -603,3 +559,66 @@ async def retrieve_file(
raise
finally:
tracker = None # remove ref

# exposed API

async def retrieve_file(
self, raw_url: str, headers_from_client: dict[str, str]
) -> tuple[AsyncIterator[bytes], CIMultiDict[str] | CIMultiDictProxy[str]] | None:
"""Retrieve a file descriptor for the requested <raw_url>.
This method retrieves a file descriptor for incoming client request.
Upper uvicorn app can use this file descriptor to yield chunks of data,
and send chunks to the on-calling ota_client.
NOTE: use raw_url in all operations, except opening remote file.
Args:
raw_url: unquoted raw url received from uvicorn
headers_from_client: headers come from client's request, which will be
passthrough to upper otaproxy and/or remote OTA image server.
Returns:
A tuple contains an asyncio generator for upper server app to yield data chunks from
and headers dict that should be sent back to client in resp.
"""
if self._closed:
raise BaseOTACacheError("ota cache pool is closed")

cache_policy = OTAFileCacheControl.parse_header(
headers_from_client.get(HEADER_OTA_FILE_CACHE_CONTROL, "")
)
if cache_policy.no_cache:
logger.info(f"client indicates that do not cache for {raw_url=}")

# when there is no upper_proxy, do not passthrough the OTA_FILE_CACHE_CONTROL header.
if not self._upper_proxy:
headers_from_client.pop(HEADER_OTA_FILE_CACHE_CONTROL, None)

# a fastpath when cache is not enabled or client requires so
if not self._cache_enabled or cache_policy.no_cache:
return await self._retrieve_file_by_downloading(
raw_url, headers=headers_from_client
)

# NOTE(20241202): behavior changed: even if _cache_enabled is False, if external_cache is configured
# and loaded, still try to use external cache source.
if _res := await self._retrieve_file_by_external_cache(cache_policy):
return _res

if _res := await self._retrieve_file_by_cache_lookup(
raw_url=raw_url, cache_policy=cache_policy
):
return _res

if _res := await self._retrieve_file_by_new_caching(
raw_url=raw_url,
cache_policy=cache_policy,
headers_from_client=headers_from_client,
):
return _res

# as last resort, finally try to handle the request by directly downloading
return await self._retrieve_file_by_downloading(
raw_url, headers=headers_from_client
)

1 comment on commit e6a93db

@github-actions
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Coverage

Coverage Report
FileStmtsMissCoverMissing
src/ota_metadata/legacy
   __init__.py10100% 
   parser.py3354885%106, 170, 175, 211–212, 222–223, 226, 238, 289–291, 295–298, 324–327, 396, 399, 407–409, 422, 431–432, 435–436, 601–603, 653–654, 657, 685–686, 689–690, 692, 696, 698–699, 753, 756–758
   types.py841384%37, 40–42, 112–116, 122–125
src/ota_metadata/utils
   cert_store.py86890%58–59, 73, 87, 91, 102, 123, 127
src/ota_proxy
   __init__.py15660%48, 50, 52, 61, 71–72
   __main__.py880%16, 18–20, 22, 24–25, 27
   _consts.py170100% 
   cache_control_header.py68494%71, 91, 113, 121
   cache_streaming.py1442284%154–156, 184–186, 211, 225, 229–230, 265–266, 268, 280, 349, 355–356, 359, 367–370
   config.py200100% 
   db.py741875%110, 116, 154, 160–161, 164, 170, 172, 193–200, 202–203
   errors.py50100% 
   external_cache.py282028%31, 35, 40–42, 44–45, 48–49, 51–53, 60, 63–65, 69–72
   lru_cache_helper.py47295%84–85
   ota_cache.py2346472%71–72, 143, 146–147, 159–160, 192–193, 210, 231, 250–254, 258–260, 262, 264–271, 273–275, 278–279, 283–284, 288, 335, 343–345, 418, 445, 448–449, 471–473, 477–479, 485, 487–489, 494, 520–522, 557–559, 586, 592, 607
   server_app.py1413972%79, 82, 88, 107, 111, 170, 179, 221–222, 224–226, 229, 234–235, 238, 241–242, 245, 248, 251, 254, 267–268, 271–272, 274, 277, 303–306, 309, 323–325, 331–333
   utils.py140100% 
src/otaclient
   __init__.py5260%17, 19
   __main__.py110%16
   _logging.py513335%43–44, 46–47, 49–54, 56–57, 59–60, 62–65, 67, 77, 80–82, 84–86, 89–90, 92–96
   _status_monitor.py1611193%46, 48–49, 157, 160, 177, 180, 196–197, 204, 207
   _types.py750100% 
   errors.py120199%97
   main.py492940%30, 32–39, 41–42, 44–45, 47–48, 52, 54, 59, 61, 67–69, 72–73, 77–80, 82
   ota_core.py31610566%101, 128, 130–131, 135–136, 138–140, 144–145, 150–151, 157, 159, 198–201, 207, 211, 217, 356, 388–389, 391, 400, 403, 408–409, 412, 418, 420–424, 460–463, 466–477, 480–483, 519–522, 538, 540–542, 605–612, 617, 620–627, 660–661, 667, 671–672, 678, 703–705, 707, 782, 810–811, 813–814, 822, 824–830
   utils.py37294%73–74
src/otaclient/boot_control
   __init__.py40100% 
   _firmware_package.py932276%82, 86, 136, 180, 186, 209–210, 213–218, 220–221, 224–229, 231
   _grub.py41812769%214, 262–265, 271–275, 312–313, 320–325, 328–334, 337, 340–341, 346, 348–350, 359–365, 367–368, 370–372, 381–383, 385–387, 466–467, 471–472, 524, 530, 556, 578, 582–583, 598–600, 624–627, 639, 643–645, 647–649, 708–711, 736–739, 762–765, 777–778, 781–782, 817, 823, 843–844, 846, 871–873, 891–894, 919–922, 929–932, 937–945, 950–957
   _jetson_cboot.py2612610%20, 22–25, 27–29, 35–40, 42, 58–60, 62, 64–65, 71, 75, 134, 137, 139–140, 143, 150–151, 159–160, 163, 169–170, 178, 187–191, 193, 199, 202–203, 209, 212–213, 218–219, 221, 227–228, 231–232, 235–237, 239, 245, 250–252, 254–256, 261, 263–266, 268–269, 278–279, 282–283, 288–289, 292–296, 299–300, 305–306, 309, 312–316, 321–324, 327, 330–331, 334, 337–338, 341, 345–350, 354–355, 359, 362–363, 366, 369–372, 374, 377–378, 382, 385, 388–391, 393, 400, 404–405, 408–409, 415–416, 422, 424–425, 429, 431, 433–435, 438, 442, 445, 448–449, 451, 454, 462–463, 470, 480, 483, 491–492, 497–500, 502, 509, 511–513, 519–520, 524–525, 528, 532, 535, 537, 544–548, 550, 562–565, 568, 571, 573, 580, 590–592, 594, 596, 599, 602, 605, 607–608, 611–615, 619–621, 623, 631–635, 637, 640, 644, 647, 658–659, 664, 674, 677–683, 687–693, 697–706, 710–718, 722, 724, 726–728
   _jetson_common.py1724573%132, 140, 288–291, 294, 311, 319, 354, 359–364, 382, 408–409, 411–413, 417–420, 422–423, 425–429, 431, 438–439, 442–443, 453, 456–457, 460, 462, 506–507
   _jetson_uefi.py40427432%124–126, 131–132, 151–153, 158–161, 328, 446, 448–451, 455, 459–460, 462–470, 472, 484–485, 488–489, 492–493, 496–498, 502–503, 508–510, 514, 518–519, 522–523, 526–527, 531, 534–535, 537, 542–543, 547, 550–551, 556, 560–561, 564, 568–570, 572, 576–579, 581–582, 604–605, 609–610, 612, 616, 620–621, 624–625, 632, 635–637, 640, 642–643, 648–649, 652–655, 657–658, 663, 665–666, 674, 677–680, 682–683, 685, 689–690, 694, 702–706, 709–710, 712, 715–719, 722, 725–729, 733–734, 737–742, 745–746, 749–752, 754–755, 762–763, 773–776, 779, 782–785, 788–792, 795–796, 799, 802–805, 808, 810, 815–816, 819, 822–825, 827, 833, 838–839, 858–859, 862, 870–871, 878, 888, 891, 898–899, 904–907, 915–918, 926–927, 939–942, 944, 947, 950, 958, 969–971, 973–975, 977–981, 986–987, 989, 1002, 1006, 1009, 1019, 1024, 1032–1033, 1036, 1040, 1042–1044, 1050–1051, 1056, 1064–1069, 1074–1079, 1084–1092, 1097–1104, 1112–1114
   _ota_status_control.py1021189%117, 122, 127, 240, 244–245, 248, 255, 257–258, 273
   _rpi_boot.py28713353%53, 56, 120–121, 125, 133–136, 150–153, 158–159, 161–162, 167–168, 171–172, 181–182, 222, 228–232, 235, 253–255, 259–261, 266–268, 272–274, 284–285, 288, 291, 293–294, 296–297, 299–301, 307, 310–311, 321–324, 332–336, 338, 340–341, 346–347, 354, 357–362, 393, 395–398, 408–411, 415–416, 418–422, 450–453, 472–475, 501–504, 509–517, 522–529, 544–547, 554–557, 565–567
   _slot_mnt_helper.py100100% 
   configs.py510100% 
   protocol.py60100% 
   selecter.py412929%44–46, 49–50, 54–55, 58–60, 63, 65, 69, 77–79, 81–82, 84–85, 89, 91, 93–94, 96, 98–99, 101, 103
src/otaclient/configs
   __init__.py170100% 
   _cfg_configurable.py470100% 
   _cfg_consts.py47197%97
   _common.py80100% 
   _ecu_info.py56394%59, 64–65
   _proxy_info.py51590%85, 87–88, 90, 101
   cfg.py190100% 
src/otaclient/create_standby
   __init__.py13192%36
   common.py2264480%59, 62–63, 67–69, 71, 75–76, 78, 126, 174–176, 178–180, 182, 185–188, 192, 203, 279–280, 282–287, 299, 339, 367, 370–372, 388–389, 403, 407, 429–430
   interface.py70100% 
   rebuild_mode.py1151091%98–100, 119, 150–155
src/otaclient/grpc
   _otaproxy_ctx.py644135%38, 40–41, 43–45, 47, 52–56, 58, 77–78, 80, 83, 87, 100–102, 106–107, 109, 111–112, 118–120, 124, 133–134, 136–141, 143–145
src/otaclient/grpc/api_v2
   ecu_status.py1531093%75, 78, 81, 147, 171, 173, 300, 370–371, 410
   ecu_tracker.py341944%40–42, 48, 52–54, 61–63, 66, 70–74, 77–79
   servicer.py1273671%82, 171–173, 180, 191–192, 233–234, 237–241, 250–259, 266, 272, 275–278, 284–285, 292, 295, 301, 304
   types.py46295%78–79
src/otaclient_api/v2
   api_caller.py39684%45–47, 83–85
   types.py2562391%86, 89–92, 131, 209–210, 212, 259, 262–263, 506–508, 512–513, 515, 518–519, 522–523, 586
src/otaclient_common
   __init__.py341555%42–44, 61, 63, 68–77
   _io.py64198%41
   cmdhelper.py130100% 
   common.py1061189%44, 148, 151–153, 168, 175–177, 271, 275
   downloader.py1991094%107–108, 126, 153, 369, 424, 428, 516–517, 526
   linux.py611575%51–53, 59, 69, 74, 76, 108–109, 133–134, 190, 195–196, 198
   logging.py29196%55
   persist_file_handling.py1181884%113, 118, 150–152, 163, 192–193, 228–232, 242–244, 246–247
   proto_streamer.py42880%33, 48, 66–67, 72, 81–82, 100
   proto_wrapper.py3984688%87, 165, 172, 184–186, 205, 210, 221, 257, 263, 268, 299, 303, 307, 402, 462, 469, 472, 492, 499, 501, 526, 532, 535, 537, 562, 568, 571, 573, 605, 609, 611, 625, 642, 669, 672, 676, 707, 713, 762–763, 765, 803–805
   retry_task_map.py105595%158–159, 161, 181–182
   shm_status.py952177%79–80, 83–84, 105, 120–122, 134, 139, 156–160, 169–170, 172, 179, 192, 204
   typing.py31487%48, 97–98, 100
TOTAL6605169474% 

Tests Skipped Failures Errors Time
236 0 💤 0 ❌ 0 🔥 12m 35s ⏱️

Please sign in to comment.