-
Notifications
You must be signed in to change notification settings - Fork 772
perf(file-system): Fix quadratic file rescans in the FS request queue #2011
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
vdusek
wants to merge
1
commit into
master
Choose a base branch
from
worktree-fix-fs-rq
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+151
−70
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,7 +1,6 @@ | ||
| from __future__ import annotations | ||
|
|
||
| import asyncio | ||
| import functools | ||
| import json | ||
| import shutil | ||
| from collections import deque | ||
|
|
@@ -333,7 +332,8 @@ async def add_batch_of_requests( | |
| unprocessed_requests = list[UnprocessedRequest]() | ||
| state = self._state.current_value | ||
|
|
||
| all_requests = state.forefront_requests | state.regular_requests | ||
| def is_in_queue(unique_key: str) -> bool: | ||
| return unique_key in state.forefront_requests or unique_key in state.regular_requests | ||
|
|
||
| requests_to_enqueue = {} | ||
|
|
||
|
|
@@ -352,7 +352,7 @@ async def add_batch_of_requests( | |
| # Or if the request is already in the queue and the `forefront` flag is not used, we do not change the | ||
| # position of the request. | ||
| elif (request.unique_key in state.in_progress_requests) or ( | ||
| request.unique_key in all_requests and not forefront | ||
| is_in_queue(request.unique_key) and not forefront | ||
| ): | ||
| processed_requests.append( | ||
| ProcessedRequest( | ||
|
|
@@ -368,7 +368,7 @@ async def add_batch_of_requests( | |
| # Process each request in the batch. | ||
| for request in requests_to_enqueue.values(): | ||
| # If the request is not already in the RQ, this is a new request. | ||
| if request.unique_key not in all_requests: | ||
| if not is_in_queue(request.unique_key): | ||
| request_path = self._get_request_path(request.unique_key) | ||
| # Add sequence number to ensure FIFO ordering using state. | ||
| if forefront: | ||
|
|
@@ -384,6 +384,16 @@ async def add_batch_of_requests( | |
| request_data = await json_dumps(request.model_dump()) | ||
| await atomic_write(request_path, request_data) | ||
|
|
||
| # A new forefront request belongs to the very front of the queue, so it can go straight | ||
| # to the cache without a full refresh, unless the cache is at capacity. New regular | ||
| # requests must not be appended to the cache: it may be truncated, so its tail is not | ||
| # necessarily the end of the queue; they are picked up by a refresh once the cache drains. | ||
| if forefront: | ||
| if len(self._request_cache) < self._MAX_REQUESTS_IN_CACHE: | ||
| self._request_cache.appendleft(request) | ||
| else: | ||
| self._request_cache_needs_refresh = True | ||
|
|
||
| # Update the metadata counts. | ||
| new_total_request_count += 1 | ||
| new_pending_request_count += 1 | ||
|
|
@@ -406,6 +416,9 @@ async def add_batch_of_requests( | |
| state.forefront_requests[request.unique_key] = state.forefront_sequence_counter | ||
| state.forefront_sequence_counter += 1 | ||
|
|
||
| # The request may already sit elsewhere in the cache, so its position must be recomputed. | ||
| self._request_cache_needs_refresh = True | ||
|
|
||
| processed_requests.append( | ||
| ProcessedRequest( | ||
| unique_key=request.unique_key, | ||
|
|
@@ -431,10 +444,6 @@ async def add_batch_of_requests( | |
| new_pending_request_count=new_pending_request_count, | ||
| ) | ||
|
|
||
| # Invalidate the cache if we added forefront requests. | ||
| if forefront: | ||
| self._request_cache_needs_refresh = True | ||
|
|
||
| # Invalidate is_empty cache. | ||
| self._is_empty_cache = None | ||
|
|
||
|
|
@@ -506,8 +515,12 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | | |
| request_data = await json_dumps(request.model_dump()) | ||
| await atomic_write(request_path, request_data) | ||
|
|
||
| # Update state: remove from in-progress and add to handled. | ||
| # Update state: remove from in-progress and pending, and add to handled. Dropping the key from | ||
| # the pending mappings keeps them (and the persisted state) sized by the backlog rather than by | ||
| # every request ever processed. | ||
| state.in_progress_requests.discard(request.unique_key) | ||
| state.forefront_requests.pop(request.unique_key, None) | ||
| state.regular_requests.pop(request.unique_key, None) | ||
| state.handled_requests.add(request.unique_key) | ||
|
|
||
| # Update RQ metadata. | ||
|
|
@@ -696,63 +709,36 @@ async def _update_metadata( | |
| await atomic_write(self.path_to_metadata, data) | ||
|
|
||
| async def _refresh_cache(self) -> None: | ||
| """Refresh the request cache from filesystem. | ||
| """Refresh the request cache from the filesystem. | ||
|
|
||
| This method loads up to _MAX_REQUESTS_IN_CACHE requests from the filesystem, | ||
| prioritizing forefront requests and maintaining proper ordering. | ||
| This method loads up to `_MAX_REQUESTS_IN_CACHE` requests from the filesystem, prioritizing forefront | ||
| requests and maintaining proper ordering. Only requests that are pending according to the state are | ||
| read, so files of already handled or in-progress requests are not touched at all. | ||
| """ | ||
| self._request_cache.clear() | ||
| state = self._state.current_value | ||
|
|
||
| forefront_requests = list[tuple[Request, int]]() # (request, sequence) | ||
| regular_requests = list[tuple[Request, int]]() # (request, sequence) | ||
|
|
||
| request_files = await self._get_request_files(self.path_to_rq) | ||
|
|
||
| for request_file in request_files: | ||
| request = await self._parse_request_file(request_file) | ||
|
|
||
| if request is None: | ||
| continue | ||
|
|
||
| # Skip handled requests | ||
| if request.unique_key in state.handled_requests: | ||
| continue | ||
|
|
||
| # Skip in-progress requests | ||
| if request.unique_key in state.in_progress_requests: | ||
| continue | ||
|
|
||
| # Determine if request is forefront or regular based on state | ||
| if request.unique_key in state.forefront_requests: | ||
| sequence = state.forefront_requests[request.unique_key] | ||
| forefront_requests.append((request, sequence)) | ||
| elif request.unique_key in state.regular_requests: | ||
| sequence = state.regular_requests[request.unique_key] | ||
| regular_requests.append((request, sequence)) | ||
| else: | ||
| # Request not in state, skip it (might be orphaned) | ||
| logger.warning(f'Request {request.unique_key} not found in state, skipping.') | ||
| continue | ||
|
|
||
| # Sort forefront requests by sequence (newest first for LIFO behavior). | ||
| forefront_requests.sort(key=lambda item: item[1], reverse=True) | ||
| def is_pending(unique_key: str) -> bool: | ||
| return unique_key not in state.handled_requests and unique_key not in state.in_progress_requests | ||
|
|
||
| # Sort regular requests by sequence (oldest first for FIFO behavior). | ||
| regular_requests.sort(key=lambda item: item[1], reverse=False) | ||
| # Forefront requests are fetched newest first (LIFO), regular requests oldest first (FIFO). | ||
| forefront_keys = sorted( | ||
| filter(is_pending, state.forefront_requests), | ||
| key=lambda unique_key: state.forefront_requests[unique_key], | ||
| reverse=True, | ||
| ) | ||
| regular_keys = sorted( | ||
| filter(is_pending, state.regular_requests), | ||
| key=lambda unique_key: state.regular_requests[unique_key], | ||
| ) | ||
|
|
||
| # Add forefront requests to the beginning of the cache (left side). Since forefront_requests are sorted | ||
| # by sequence (newest first), we need to add them in reverse order to maintain correct priority. | ||
| for request, _ in reversed(forefront_requests): | ||
| for unique_key in forefront_keys + regular_keys: | ||
| if len(self._request_cache) >= self._MAX_REQUESTS_IN_CACHE: | ||
| break | ||
| self._request_cache.appendleft(request) | ||
|
|
||
| # Add regular requests to the end of the cache (right side). | ||
| for request, _ in regular_requests: | ||
| if len(self._request_cache) >= self._MAX_REQUESTS_IN_CACHE: | ||
| break | ||
| self._request_cache.append(request) | ||
| request = await self._parse_request_file(self._get_request_path(unique_key)) | ||
| if request is not None: | ||
| self._request_cache.append(request) | ||
|
Comment on lines
+740
to
+741
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's add a log warning when there is no file for the |
||
|
|
||
| self._request_cache_needs_refresh = False | ||
|
|
||
|
|
@@ -787,25 +773,23 @@ async def _parse_request_file(cls, file_path: Path) -> Request | None: | |
| Returns: | ||
| The parsed `Request` object or `None` if the file could not be read or parsed. | ||
| """ | ||
| # Open the request file. | ||
| # Read the request file in a thread to avoid blocking the event loop. | ||
| try: | ||
| file = await asyncio.to_thread(functools.partial(file_path.open, mode='r', encoding='utf-8')) | ||
| file_content = await asyncio.to_thread(file_path.read_text, encoding='utf-8') | ||
| except FileNotFoundError: | ||
| logger.warning(f'Request file "{file_path}" not found.') | ||
| return None | ||
|
|
||
| # Read the file content and parse it as JSON. | ||
| # Parse the file content as JSON. | ||
| try: | ||
| file_content = json.load(file) | ||
| parsed = json.loads(file_content) | ||
| except json.JSONDecodeError as exc: | ||
| logger.warning(f'Failed to parse request file {file_path}: {exc!s}') | ||
| return None | ||
| finally: | ||
| await asyncio.to_thread(file.close) | ||
|
|
||
| # Validate the content against the Request model. | ||
| try: | ||
| return Request.model_validate(file_content) | ||
| return Request.model_validate(parsed) | ||
| except ValidationError as exc: | ||
| logger.warning(f'Failed to validate request file {file_path}: {exc!s}') | ||
| return None | ||
|
|
@@ -830,16 +814,17 @@ async def _discover_existing_requests(self) -> None: | |
| if request is None: | ||
| continue | ||
|
|
||
| # Add request to state as regular request (assign sequence numbers) | ||
| if request.unique_key not in state.regular_requests and request.unique_key not in state.forefront_requests: | ||
| # Assign as regular request with current sequence counter | ||
| # Already handled requests are tracked only for deduplication, not as pending work. | ||
| if request.handled_at is not None: | ||
| state.handled_requests.add(request.unique_key) | ||
|
|
||
| # Add pending request to state as regular request (assign sequence numbers) | ||
| elif ( | ||
| request.unique_key not in state.regular_requests and request.unique_key not in state.forefront_requests | ||
| ): | ||
| state.regular_requests[request.unique_key] = state.sequence_counter | ||
| state.sequence_counter += 1 | ||
|
|
||
| # Check if request was already handled | ||
| if request.handled_at is not None: | ||
| state.handled_requests.add(request.unique_key) | ||
|
|
||
| @staticmethod | ||
| def _get_file_base_name_from_unique_key(unique_key: str) -> str: | ||
| """Generate a deterministic file name for a unique_key. | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -259,3 +259,99 @@ async def test_get_request_does_not_mark_in_progress(rq_client: FileSystemReques | |
| next_request = await rq_client.fetch_next_request() | ||
| assert next_request is not None | ||
| assert next_request.unique_key == request.unique_key | ||
|
|
||
|
|
||
| async def test_forefront_add_fetch_handle_parses_linear_number_of_files( | ||
| rq_client: FileSystemRequestQueueClient, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| """Test that per-request forefront adds (the `RequestManagerTandem` pattern) do not rescan all request files.""" | ||
| parse_count = 0 | ||
| original_parse = type(rq_client)._parse_request_file | ||
|
|
||
| async def counting_parse(file_path: Path) -> Request | None: | ||
| nonlocal parse_count | ||
| parse_count += 1 | ||
| return await original_parse(file_path) | ||
|
|
||
| monkeypatch.setattr(type(rq_client), '_parse_request_file', staticmethod(counting_parse)) | ||
|
|
||
| n = 25 | ||
| for i in range(n): | ||
| await rq_client.add_batch_of_requests([Request.from_url(f'https://example.com/{i}')], forefront=True) | ||
| request = await rq_client.fetch_next_request() | ||
| assert request is not None | ||
| await rq_client.mark_request_as_handled(request) | ||
|
|
||
| # Before the fix, each forefront add invalidated the cache and each fetch re-parsed every request file | ||
| # ever written, giving n * (n + 1) / 2 parses in total. With the fix, new forefront requests go straight | ||
| # to the cache, so only the initial refresh parses anything. | ||
| assert parse_count <= 5 | ||
|
Comment on lines
+264
to
+288
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Maybe we could use For example from unittest.mock import patch
async def test_forefront_add_fetch_handle_parses_linear_number_of_files(
rq_client: FileSystemRequestQueueClient,
) -> None:
with patch.object(
rq_client,
'_parse_request_file',
wraps=rq_client._parse_request_file,
) as call_counter:
n = 25
for i in range(n):
await rq_client.add_batch_of_requests([Request.from_url(f'https://example.com/{i}')], forefront=True)
request = await rq_client.fetch_next_request()
assert request is not None
await rq_client.mark_request_as_handled(request)
assert call_counter.call_count <= 5 |
||
|
|
||
|
|
||
| async def test_cache_refresh_skips_handled_request_files( | ||
| rq_client: FileSystemRequestQueueClient, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| """Test that a cache refresh reads only pending request files, not files of already handled requests.""" | ||
| await rq_client.add_batch_of_requests([Request.from_url(f'https://example.com/{i}') for i in range(5)]) | ||
| for _ in range(5): | ||
| request = await rq_client.fetch_next_request() | ||
| assert request is not None | ||
| await rq_client.mark_request_as_handled(request) | ||
|
|
||
| # The 5 handled request files stay on disk. Add 5 new requests; fetching the next one refreshes | ||
| # the cache, which must not parse the handled files again. | ||
| await rq_client.add_batch_of_requests([Request.from_url(f'https://example.com/new/{i}') for i in range(5)]) | ||
|
|
||
| parse_count = 0 | ||
| original_parse = type(rq_client)._parse_request_file | ||
|
|
||
| async def counting_parse(file_path: Path) -> Request | None: | ||
| nonlocal parse_count | ||
| parse_count += 1 | ||
| return await original_parse(file_path) | ||
|
|
||
| monkeypatch.setattr(type(rq_client), '_parse_request_file', staticmethod(counting_parse)) | ||
|
|
||
| request = await rq_client.fetch_next_request() | ||
| assert request is not None | ||
| assert parse_count == 5 | ||
|
|
||
|
|
||
| async def test_handled_requests_pruned_from_pending_state(rq_client: FileSystemRequestQueueClient) -> None: | ||
| """Test that handling a request removes it from the pending state mappings so the state stays bounded.""" | ||
| await rq_client.add_batch_of_requests([Request.from_url(f'https://example.com/{i}') for i in range(3)]) | ||
|
|
||
| handled_keys = set() | ||
| for _ in range(3): | ||
| request = await rq_client.fetch_next_request() | ||
| assert request is not None | ||
| await rq_client.mark_request_as_handled(request) | ||
| handled_keys.add(request.unique_key) | ||
|
|
||
| state = rq_client._state.current_value | ||
| assert not state.regular_requests | ||
| assert not state.forefront_requests | ||
| assert not state.in_progress_requests | ||
| assert state.handled_requests == handled_keys | ||
|
|
||
|
|
||
| async def test_handled_requests_deduplicated_after_reopen() -> None: | ||
| """Test that requests handled before a reopen are not served again and still deduplicate re-adds.""" | ||
| storage_client = FileSystemStorageClient() | ||
| client = await storage_client.create_rq_client(name='handled-reopen-test') | ||
|
|
||
| request = Request.from_url('https://example.com/handled') | ||
| await client.add_batch_of_requests([request]) | ||
| fetched = await client.fetch_next_request() | ||
| assert fetched is not None | ||
| await client.mark_request_as_handled(fetched) | ||
| await client._state.persist_state() | ||
|
|
||
| rq_id = (await client.get_metadata()).id | ||
| reopened = await storage_client.create_rq_client(id=rq_id) | ||
|
|
||
| response = await reopened.add_batch_of_requests([request]) | ||
| assert response.processed_requests[0].was_already_handled is True | ||
| assert await reopened.fetch_next_request() is None | ||
|
|
||
| await reopened.drop() | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We can probably skip this if
self._request_cache_needs_refreshis alreadyTrue.