[ENH]: Propagate cancellations to NAC queue

[tanujnay112]

Description of changes

Summarize the changes made by this PR.

• Improvements & Bug fixes
  • When a future running a NAC request gets dropped, it "removes" its priority from the pending request if it hasn't been sent out to object storage yet.
  • The priority data structure for in-flight requests has been updated from a binary value of priorities to a set of counters, denoting how many tasks requested each priority level.
• New functionality
  • Ideally, the above changes should make it such that high priority NAC requests get lowered in priority when the calling requests are cancelled/timed out.

Test plan

How are these changes tested?

• Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Migration plan

N/A

Observability plan

It would be a good idea to add a metric that tracks the number of total waiters on the NAC layer.

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs section?

[tanujnay112]

• [ENH]: Propagate cancellations to NAC queue #5129 👈 (View in Graphite)
• [ENH]: Propagate cancellation to all spawned Tasks #5123
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[propel-code-bot]

Propagate Cancellations to NAC Queue and Improve Priority Management

This PR introduces substantial changes to the admission control (NAC) layer for S3-backed storage. Cancellations of in-flight NAC requests are now properly propagated, so when a future is dropped (i.e., caller cancels or times out), its associated priority is removed from the pending request's priority counters. The handling of priorities has been overhauled: rather than a single binary priority flag, a thread-safe set of per-level counters (PriorityHolder) now tracks how many outstanding tasks require each priority level. This allows the effective priority of a request to be dynamically lowered if all high-priority waiters are cancelled. The refactor introduces several new types (e.g., PriorityHolder, RollbackPriorityOnDrop), modifies all relevant internal APIs to use them, and consistently updates the codebase to maintain correct lock and notification semantics.

Key Changes

• Replaced atomic priority flag with a PriorityHolder struct using atomic counters for each priority level.
• Added logic to remove (rollback) a request's priority upon cancellation or future drop, ensuring correct request de-prioritization.
• Refactored in-flight request map to store Arc to facilitate shared state for cancellation management.
• Updated APIs, including parallel reads, plain reads, writes, and semaphore-based rate limiting, to utilize the new priority management constructs.
• Implemented RollbackPriorityOnDrop RAII pattern for automatic cleanup of priority claims and notification of waiting tasks.
• Updated all usages of rate limiter admission and related APIs to operate on the new priority structures.

Affected Areas

• rust/storage/src/admissioncontrolleds3.rs
• Request deduplication and coalescing logic
• Priority and rate limiting for S3 operations
• Internal APIs and data structures for priority propagation and cleanup

This summary was automatically generated by @propel-code-bot

[propel-code-bot]

[BestPractice]

Inconsistent error handling: The unwrap() call on line 717 could cause a panic if get_priority() returns None. While the comment suggests this shouldn't happen, defensive programming would handle this case gracefully.

Suggested change

	let current_priority = priority.get_priority().unwrap();
	let current_priority = match priority.get_priority() {
	Some(pri) => pri,
	None => {
	tracing::error!("Priority holder has no active priorities during acquire");
	return self.remaining_tokens[StorageRequestPriority::lowest().as_usize()]
	.acquire()
	.await
	.expect("Failed to acquire semaphore token");
	}
	};

⚡ Committable suggestion

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.