Skip to content

Commit

Permalink
refactor(backport v3.8.x): ota_proxy: refine request handling flow (#443
Browse files Browse the repository at this point in the history 
)

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.
  • Loading branch information
@Bodong-Yang
Bodong-Yang authored Dec 2, 2024
1 parent 23e98b4 commit a644fc4
Showing 1 changed file with 102 additions and 79 deletions.
181 changes: 102 additions & 79 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 @@ -276,7 +276,7 @@ def _background_check_free_space(self):
)
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 @@ -368,7 +368,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 @@ -392,28 +392,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 @@ -423,7 +428,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 @@ -445,7 +449,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 or not client_cache_policy.file_sha256:
if (
not self._external_cache
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 @@ -473,83 +482,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
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 @@ -588,3 +548,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 a644fc4

@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__.py110100% 
   parser.py3414088%103, 162, 167, 203–204, 214–215, 218, 230, 288, 298–301, 340–343, 423, 426, 434–436, 449, 458–459, 462–463, 675–676, 686, 688, 691, 718–720, 770, 773–775
   types.py841384%37, 40–42, 112–116, 122–125
src/ota_proxy
   __init__.py361072%59, 61, 63, 72, 81–82, 102, 104–106
   __main__.py770%16–18, 20, 22–23, 25
   _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.py170100% 
   db.py741875%110, 116, 154, 160–161, 164, 170, 172, 193–200, 202–203
   errors.py50100% 
   lru_cache_helper.py47295%84–85
   ota_cache.py2286173%70–71, 140, 151–152, 184–185, 202, 239–243, 247–249, 251, 253–260, 262–264, 267–268, 272–273, 277, 324, 332–334, 407, 434, 437–438, 460–462, 466–468, 474, 476–478, 483, 509–511, 546–548, 575, 581, 596
   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
   log_setting.py52590%53, 55, 64–66
src/otaclient/app
   __main__.py110%16
   configs.py760100% 
   errors.py1200100% 
   interface.py30100% 
   main.py46589%52–53, 75–77
   ota_client.py38111569%80, 88, 109, 136, 138–139, 142–143, 145–146, 150, 154–155, 160–161, 167, 169, 207–210, 216, 220, 226, 345, 357–358, 360, 369, 372, 377–378, 381, 387, 389–393, 412–415, 418–429, 457–460, 506–507, 511, 513–514, 544–545, 554–561, 568, 571–577, 622–625, 633, 669–671, 676–678, 681–682, 684–685, 687, 745–746, 749, 757–758, 761, 772–773, 776, 784–785, 788, 799, 818, 845, 864, 882
   ota_client_stub.py39310972%75–77, 79–80, 88–91, 94–96, 100, 105–106, 108–109, 112, 114–115, 118–120, 123–124, 127–129, 134–139, 143, 146–150, 152–153, 161–163, 166, 203–205, 210, 246, 271, 274, 277, 381, 405, 407, 431, 477, 534, 604–605, 644, 663–665, 671–674, 678–680, 687–689, 692, 696–699, 752, 841–843, 850, 880–881, 884–888, 897–906, 913, 919, 922–923, 927, 930
   update_stats.py104991%57, 103, 105, 114, 116, 125, 127, 148, 179
src/otaclient/boot_control
   __init__.py40100% 
   _common.py24811254%74–75, 96–98, 114–115, 135–136, 155–156, 175–176, 195–196, 218–220, 235–236, 260–266, 287, 295, 313, 321, 340–341, 344–345, 368, 370–379, 381–390, 392–394, 413, 416, 424, 432, 448–450, 452–457, 550, 555, 560, 673, 677–678, 681, 689, 691–692, 718–719, 721–724, 729, 735–736, 739–740, 742, 749–750, 761–767, 777–779, 783–784, 787–788, 791, 797
   _firmware_package.py942276%83, 87, 137, 181, 187, 210–211, 214–219, 221–222, 225–230, 232
   _grub.py41712869%217, 265–268, 274–278, 315–316, 323–328, 331–337, 340, 343–344, 349, 351–353, 362–368, 370–371, 373–375, 384–386, 388–390, 469–470, 474–475, 527, 533, 559, 581, 585–586, 601–603, 627–630, 642, 646–648, 650–652, 711–714, 739–742, 765–768, 780–781, 784–785, 820, 826, 846–847, 849, 861, 864, 867, 870, 874–876, 894–897, 925–928, 933–941, 946–954
   _jetson_cboot.py2622620%20, 22–25, 27–29, 35–38, 40–41, 57–58, 60, 62–63, 69, 73, 132, 135, 137–138, 141, 148–149, 157–158, 161, 167–168, 176, 185–189, 191, 197, 200–201, 207, 210–211, 216–217, 219, 225–226, 229–230, 233–235, 237, 243, 248–250, 252–254, 259, 261–264, 266–267, 276–277, 280–281, 286–287, 290–294, 297–298, 303–304, 307, 310–314, 319–322, 325, 328–329, 332, 335–336, 339, 343–348, 352–353, 357, 360–361, 364, 367–370, 372, 375–376, 380, 383, 386–389, 391, 398, 402–403, 406–407, 413–414, 420, 422–423, 427, 429, 431–433, 436, 440, 443, 446–447, 449, 452, 460–461, 468, 478, 481, 489–490, 495–498, 500, 507, 509–511, 517–518, 522–523, 526, 530, 533, 535, 542–546, 548, 560–563, 566, 569, 571, 578, 582–583, 585–586, 588–590, 592, 594, 597, 600, 603, 605–606, 609–613, 617–619, 621, 629–633, 635, 638, 642, 645, 656–657, 662, 672, 675–683, 687–696, 700–709, 713, 715–717, 719–720, 722–723
   _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.py39727131%127–129, 134–135, 154–156, 161–164, 331, 449, 451–454, 458, 462–463, 465–473, 475, 487–488, 491–492, 495–496, 499–501, 505–506, 511–513, 517, 521–522, 525–526, 529–530, 534, 537–538, 540, 545–546, 550, 553–554, 559, 563–565, 569–571, 573, 577–580, 582–583, 605–606, 610–611, 613, 617, 621–622, 625–626, 633, 636–638, 641, 643–644, 649–650, 653–656, 658–659, 664, 666–667, 675, 678–681, 683–684, 686, 690–691, 695, 703–707, 710–711, 713, 716–720, 723, 726–730, 734–735, 738–743, 746–747, 750–753, 755–756, 763–764, 774–777, 780, 783–786, 789–793, 796–797, 800, 803–806, 809, 811, 816–817, 820, 823–826, 828, 834, 839–840, 859–860, 863, 871–872, 879, 889, 892, 899–900, 905–908, 916–919, 927–928, 940–943, 945, 948, 951, 959, 970–972, 974–976, 978–982, 987–988, 990, 1003, 1007, 1010, 1020, 1025, 1033–1034, 1037, 1041, 1043–1045, 1051–1052, 1057, 1065–1072, 1077–1085, 1090–1098, 1104–1106
   _rpi_boot.py28713453%55, 58, 122–123, 127, 135–138, 152–155, 162–163, 165–166, 171–172, 175–176, 185–186, 224, 230–234, 237, 255–257, 261–263, 268–270, 274–276, 286–287, 290, 293, 295–296, 298–299, 301–303, 309, 312–313, 323–326, 334–338, 340, 342–343, 348–349, 356–362, 393, 395–398, 408–411, 415–416, 418–422, 450–453, 472–475, 480, 483, 501–504, 509–517, 522–530, 547–550, 556–558, 561, 564
   configs.py550100% 
   protocol.py40100% 
   selecter.py412929%45–47, 50–51, 55–56, 59–61, 64, 66, 70, 78–80, 82–83, 85–86, 90, 92, 94–95, 97, 99–100, 102, 104
src/otaclient/configs
   _common.py80100% 
   ecu_info.py58198%108
   proxy_info.py52296%88, 90
src/otaclient/create_standby
   __init__.py12558%29–31, 33, 35
   common.py2244480%62, 65–66, 70–72, 74, 78–79, 81, 127, 175–177, 179–181, 183, 186–189, 193, 204, 278–279, 281–286, 298, 335, 363, 366–368, 384–385, 399, 403, 425–426
   interface.py50100% 
   rebuild_mode.py97990%93–95, 107–112
src/otaclient_api/v2
   api_caller.py39684%45–47, 83–85
   api_stub.py170100% 
   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__.py34876%42–44, 61, 63, 69, 76–77
   common.py1561888%47, 202, 205–207, 222, 229–231, 297–299, 309, 318–320, 366, 370
   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.py3984887%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, 760–763, 765, 803–805
   retry_task_map.py105595%158–159, 161, 181–182
   typing.py25388%69–70, 72
TOTAL6332169073% 

Tests Skipped Failures Errors Time
217 0 💤 0 ❌ 0 🔥 13m 7s ⏱️

Please sign in to comment.