1use asupersync::types::CancelReason;
7use asupersync::{Budget, Cx, Outcome, RegionId, TaskId, Time};
8use std::sync::Arc;
9
10use crate::dependency::{CleanupStack, DependencyCache, DependencyOverrides, ResolutionStack};
11
12pub const DEFAULT_MAX_BODY_SIZE: usize = 1024 * 1024;
14
15#[derive(Debug, Clone, Copy)]
21pub struct BodyLimitConfig {
22 max_size: usize,
24}
25
26impl Default for BodyLimitConfig {
27 fn default() -> Self {
28 Self {
29 max_size: DEFAULT_MAX_BODY_SIZE,
30 }
31 }
32}
33
34impl BodyLimitConfig {
35 #[must_use]
37 pub fn new(max_size: usize) -> Self {
38 Self { max_size }
39 }
40
41 #[must_use]
43 pub fn max_size(self) -> usize {
44 self.max_size
45 }
46}
47
48#[derive(Debug, Clone)]
75pub struct RequestContext {
76 cx: Cx,
78 request_id: u64,
80 dependency_cache: Arc<DependencyCache>,
82 dependency_overrides: Arc<DependencyOverrides>,
84 resolution_stack: Arc<ResolutionStack>,
86 cleanup_stack: Arc<CleanupStack>,
88 body_limit: BodyLimitConfig,
90 deadline: Option<Time>,
99}
100
101impl RequestContext {
102 #[must_use]
108 pub fn new(cx: Cx, request_id: u64) -> Self {
109 Self {
110 cx,
111 request_id,
112 dependency_cache: Arc::new(DependencyCache::new()),
113 dependency_overrides: Arc::new(DependencyOverrides::new()),
114 resolution_stack: Arc::new(ResolutionStack::new()),
115 cleanup_stack: Arc::new(CleanupStack::new()),
116 body_limit: BodyLimitConfig::default(),
117 deadline: None,
118 }
119 }
120
121 #[must_use]
126 pub fn with_body_limit(cx: Cx, request_id: u64, max_body_size: usize) -> Self {
127 Self {
128 cx,
129 request_id,
130 dependency_cache: Arc::new(DependencyCache::new()),
131 dependency_overrides: Arc::new(DependencyOverrides::new()),
132 resolution_stack: Arc::new(ResolutionStack::new()),
133 cleanup_stack: Arc::new(CleanupStack::new()),
134 body_limit: BodyLimitConfig::new(max_body_size),
135 deadline: None,
136 }
137 }
138
139 #[must_use]
141 pub fn with_overrides(cx: Cx, request_id: u64, overrides: Arc<DependencyOverrides>) -> Self {
142 Self {
143 cx,
144 request_id,
145 dependency_cache: Arc::new(DependencyCache::new()),
146 dependency_overrides: overrides,
147 resolution_stack: Arc::new(ResolutionStack::new()),
148 cleanup_stack: Arc::new(CleanupStack::new()),
149 body_limit: BodyLimitConfig::default(),
150 deadline: None,
151 }
152 }
153
154 #[must_use]
156 pub fn with_overrides_and_body_limit(
157 cx: Cx,
158 request_id: u64,
159 overrides: Arc<DependencyOverrides>,
160 max_body_size: usize,
161 ) -> Self {
162 Self {
163 cx,
164 request_id,
165 dependency_cache: Arc::new(DependencyCache::new()),
166 dependency_overrides: overrides,
167 resolution_stack: Arc::new(ResolutionStack::new()),
168 cleanup_stack: Arc::new(CleanupStack::new()),
169 body_limit: BodyLimitConfig::new(max_body_size),
170 deadline: None,
171 }
172 }
173
174 #[must_use]
181 pub fn with_deadline(mut self, deadline: Time) -> Self {
182 self.deadline = Some(deadline);
183 self
184 }
185
186 #[must_use]
188 pub fn deadline(&self) -> Option<Time> {
189 self.deadline
190 }
191
192 #[must_use]
202 pub fn deadline_exceeded(&self) -> bool {
203 self.deadline
204 .is_some_and(|deadline| self.cx.now() >= deadline)
205 }
206
207 #[must_use]
211 pub fn request_id(&self) -> u64 {
212 self.request_id
213 }
214
215 #[must_use]
217 pub fn dependency_cache(&self) -> &DependencyCache {
218 &self.dependency_cache
219 }
220
221 #[must_use]
223 pub fn dependency_overrides(&self) -> &DependencyOverrides {
224 &self.dependency_overrides
225 }
226
227 #[must_use]
229 pub fn resolution_stack(&self) -> &ResolutionStack {
230 &self.resolution_stack
231 }
232
233 #[must_use]
237 pub fn cleanup_stack(&self) -> &CleanupStack {
238 &self.cleanup_stack
239 }
240
241 #[must_use]
246 pub fn body_limit(&self) -> &BodyLimitConfig {
247 &self.body_limit
248 }
249
250 #[must_use]
254 pub fn max_body_size(&self) -> usize {
255 self.body_limit.max_size()
256 }
257
258 #[must_use]
264 pub fn region_id(&self) -> RegionId {
265 self.cx.region_id()
266 }
267
268 #[must_use]
270 pub fn task_id(&self) -> TaskId {
271 self.cx.task_id()
272 }
273
274 #[must_use]
280 pub fn budget(&self) -> Budget {
281 self.cx.budget()
282 }
283
284 #[must_use]
289 pub fn is_cancelled(&self) -> bool {
290 self.cx.is_cancel_requested()
291 }
292
293 pub fn checkpoint(&self) -> Result<(), CancelledError> {
315 self.cx.checkpoint().map_err(|_| CancelledError)
316 }
317
318 pub fn masked<F, R>(&self, f: F) -> R
334 where
335 F: FnOnce() -> R,
336 {
337 self.cx.masked(f)
338 }
339
340 pub fn trace(&self, message: &str) {
345 self.cx.trace(message);
346 }
347
348 #[must_use]
353 pub fn cx(&self) -> &Cx {
354 &self.cx
355 }
356}
357
358#[derive(Debug, Clone, Copy)]
364pub struct CancelledError;
365
366impl std::fmt::Display for CancelledError {
367 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
368 write!(f, "request cancelled")
369 }
370}
371
372impl std::error::Error for CancelledError {}
373
374impl From<CancelledError> for crate::error::HttpError {
378 fn from(_: CancelledError) -> Self {
379 crate::error::HttpError::new(crate::response::StatusCode::CLIENT_CLOSED_REQUEST)
380 .with_detail("request cancelled")
381 }
382}
383
384impl crate::response::IntoResponse for CancelledError {
385 fn into_response(self) -> crate::response::Response {
386 crate::error::HttpError::from(self).into_response()
387 }
388}
389
390pub trait IntoOutcome<T, E> {
395 fn into_outcome(self) -> Outcome<T, E>;
397}
398
399impl<T, E> IntoOutcome<T, E> for Result<T, E> {
400 fn into_outcome(self) -> Outcome<T, E> {
401 match self {
402 Ok(v) => Outcome::Ok(v),
403 Err(e) => Outcome::Err(e),
404 }
405 }
406}
407
408impl<T, E> IntoOutcome<T, E> for Result<T, CancelledError>
409where
410 E: Default,
411{
412 fn into_outcome(self) -> Outcome<T, E> {
413 match self {
414 Ok(v) => Outcome::Ok(v),
415 Err(CancelledError) => Outcome::Cancelled(CancelReason::user("request cancelled")),
416 }
417 }
418}
419
420#[cfg(test)]
421mod tests {
422 use super::*;
423
424 #[test]
425 fn cancelled_error_display() {
426 let err = CancelledError;
427 assert_eq!(format!("{err}"), "request cancelled");
428 }
429
430 #[test]
431 fn checkpoint_returns_error_when_cancel_requested() {
432 let cx = Cx::for_testing();
433 let ctx = RequestContext::new(cx, 1);
434 ctx.cx().set_cancel_requested(true);
435 assert!(ctx.checkpoint().is_err());
436 }
437
438 #[test]
439 fn deadline_defaults_to_none_and_is_never_exceeded() {
440 let ctx = RequestContext::new(Cx::for_testing(), 1);
441 assert_eq!(ctx.deadline(), None);
442 assert!(!ctx.deadline_exceeded());
443 }
444
445 #[test]
446 fn with_deadline_exposes_the_server_deadline() {
447 let deadline = Time::from_secs(5);
448 let ctx = RequestContext::new(Cx::for_testing(), 1).with_deadline(deadline);
449 assert_eq!(ctx.deadline(), Some(deadline));
450 }
451
452 #[test]
453 fn deadline_exceeded_reflects_the_runtime_clock() {
454 let past = RequestContext::new(Cx::for_testing(), 1).with_deadline(Time::ZERO);
456 assert!(past.deadline_exceeded());
457
458 let future =
460 RequestContext::new(Cx::for_testing(), 1).with_deadline(Time::from_nanos(u64::MAX));
461 assert!(!future.deadline_exceeded());
462 }
463
464 #[test]
465 fn masked_defers_cancellation_at_checkpoint() {
466 let cx = Cx::for_testing();
467 let ctx = RequestContext::new(cx, 1);
468 ctx.cx().set_cancel_requested(true);
469
470 let result = ctx.masked(|| ctx.checkpoint());
471 assert!(result.is_ok());
472 assert!(ctx.checkpoint().is_err());
473 }
474
475 #[test]
480 fn body_limit_config_default() {
481 let config = BodyLimitConfig::default();
482 assert_eq!(config.max_size(), DEFAULT_MAX_BODY_SIZE);
483 assert_eq!(config.max_size(), 1024 * 1024); }
485
486 #[test]
487 fn body_limit_config_custom() {
488 let config = BodyLimitConfig::new(512 * 1024);
489 assert_eq!(config.max_size(), 512 * 1024); }
491
492 #[test]
493 fn request_context_default_body_limit() {
494 let cx = Cx::for_testing();
495 let ctx = RequestContext::new(cx, 1);
496 assert_eq!(ctx.max_body_size(), DEFAULT_MAX_BODY_SIZE);
497 assert_eq!(ctx.body_limit().max_size(), DEFAULT_MAX_BODY_SIZE);
498 }
499
500 #[test]
501 fn request_context_custom_body_limit() {
502 let cx = Cx::for_testing();
503 let ctx = RequestContext::with_body_limit(cx, 1, 2 * 1024 * 1024);
504 assert_eq!(ctx.max_body_size(), 2 * 1024 * 1024); }
506
507 #[test]
508 fn request_context_with_overrides_has_default_limit() {
509 let cx = Cx::for_testing();
510 let overrides = Arc::new(DependencyOverrides::new());
511 let ctx = RequestContext::with_overrides(cx, 1, overrides);
512 assert_eq!(ctx.max_body_size(), DEFAULT_MAX_BODY_SIZE);
513 }
514
515 #[test]
516 fn request_context_with_overrides_and_custom_limit() {
517 let cx = Cx::for_testing();
518 let overrides = Arc::new(DependencyOverrides::new());
519 let ctx = RequestContext::with_overrides_and_body_limit(cx, 1, overrides, 4 * 1024 * 1024);
520 assert_eq!(ctx.max_body_size(), 4 * 1024 * 1024); }
522
523 #[test]
528 #[allow(clippy::similar_names)]
529 fn request_id_isolation_unique_per_context() {
530 let cx1 = Cx::for_testing();
532 let cx2 = Cx::for_testing();
533 let cx3 = Cx::for_testing();
534
535 let ctx1 = RequestContext::new(cx1, 100);
536 let ctx2 = RequestContext::new(cx2, 200);
537 let ctx3 = RequestContext::new(cx3, 300);
538
539 assert_eq!(ctx1.request_id(), 100);
541 assert_eq!(ctx2.request_id(), 200);
542 assert_eq!(ctx3.request_id(), 300);
543
544 assert_ne!(ctx1.request_id(), ctx2.request_id());
546 assert_ne!(ctx2.request_id(), ctx3.request_id());
547 }
548
549 #[test]
550 #[allow(clippy::similar_names)]
551 fn dependency_cache_isolation_per_request() {
552 let cx1 = Cx::for_testing();
554 let cx2 = Cx::for_testing();
555
556 let ctx1 = RequestContext::new(cx1, 1);
557 let ctx2 = RequestContext::new(cx2, 2);
558
559 ctx1.dependency_cache().insert::<i32>(42);
561
562 let value1 = ctx1.dependency_cache().get::<i32>();
564 let value2 = ctx2.dependency_cache().get::<i32>();
565
566 assert!(value1.is_some(), "ctx1 should have cached value");
567 assert_eq!(value1.unwrap(), 42);
568 assert!(value2.is_none(), "ctx2 should NOT have ctx1's cached value");
569 }
570
571 #[test]
572 #[allow(clippy::similar_names)]
573 fn cleanup_stack_isolation_per_request() {
574 use std::sync::atomic::{AtomicUsize, Ordering};
576
577 let cleanup_counter1 = Arc::new(AtomicUsize::new(0));
578 let cleanup_counter2 = Arc::new(AtomicUsize::new(0));
579
580 let cx1 = Cx::for_testing();
581 let cx2 = Cx::for_testing();
582
583 let ctx1 = RequestContext::new(cx1, 1);
584 let ctx2 = RequestContext::new(cx2, 2);
585
586 {
588 let counter = cleanup_counter1.clone();
589 ctx1.cleanup_stack().push(Box::new(move || {
590 Box::pin(async move {
591 counter.fetch_add(1, Ordering::SeqCst);
592 })
593 }));
594 }
595
596 {
598 let counter = cleanup_counter2.clone();
599 ctx2.cleanup_stack().push(Box::new(move || {
600 Box::pin(async move {
601 counter.fetch_add(1, Ordering::SeqCst);
602 })
603 }));
604 }
605
606 futures_executor::block_on(ctx1.cleanup_stack().run_cleanups());
608
609 assert_eq!(
611 cleanup_counter1.load(Ordering::SeqCst),
612 1,
613 "ctx1 cleanup should have run"
614 );
615 assert_eq!(
616 cleanup_counter2.load(Ordering::SeqCst),
617 0,
618 "ctx2 cleanup should NOT have run"
619 );
620
621 futures_executor::block_on(ctx2.cleanup_stack().run_cleanups());
623 assert_eq!(
624 cleanup_counter2.load(Ordering::SeqCst),
625 1,
626 "ctx2 cleanup should have run"
627 );
628 }
629
630 #[test]
631 #[allow(clippy::similar_names)]
632 fn cx_cancellation_isolation_per_request() {
633 let cx1 = Cx::for_testing();
635 let cx2 = Cx::for_testing();
636 let cx3 = Cx::for_testing();
637
638 let ctx1 = RequestContext::new(cx1, 1);
639 let ctx2 = RequestContext::new(cx2, 2);
640 let ctx3 = RequestContext::new(cx3, 3);
641
642 assert!(ctx1.checkpoint().is_ok(), "ctx1 should not be cancelled");
644 assert!(ctx2.checkpoint().is_ok(), "ctx2 should not be cancelled");
645 assert!(ctx3.checkpoint().is_ok(), "ctx3 should not be cancelled");
646
647 ctx2.cx().set_cancel_requested(true);
649
650 assert!(
652 ctx1.checkpoint().is_ok(),
653 "ctx1 should still not be cancelled"
654 );
655 assert!(ctx2.checkpoint().is_err(), "ctx2 should be cancelled");
656 assert!(
657 ctx3.checkpoint().is_ok(),
658 "ctx3 should still not be cancelled"
659 );
660 }
661
662 #[test]
663 #[allow(clippy::similar_names)]
664 fn body_limit_isolation_per_request() {
665 let cx1 = Cx::for_testing();
667 let cx2 = Cx::for_testing();
668
669 let ctx1 = RequestContext::with_body_limit(cx1, 1, 1024); let ctx2 = RequestContext::with_body_limit(cx2, 2, 1024 * 1024); assert_eq!(ctx1.max_body_size(), 1024);
675 assert_eq!(ctx2.max_body_size(), 1024 * 1024);
676
677 assert_ne!(ctx1.max_body_size(), ctx2.max_body_size());
679 }
680
681 #[test]
682 fn concurrent_requests_fully_isolated() {
683 use std::thread;
685
686 const NUM_REQUESTS: usize = 100;
687 let results = Arc::new(parking_lot::Mutex::new(Vec::with_capacity(NUM_REQUESTS)));
688
689 let handles: Vec<_> = (0..NUM_REQUESTS)
690 .map(|i| {
691 let results = results.clone();
692 thread::spawn(move || {
693 let cx = Cx::for_testing();
694 let request_id = (i + 1) as u64 * 1000; let ctx = RequestContext::new(cx, request_id);
696
697 ctx.dependency_cache().insert::<u64>(request_id);
699
700 let cached = ctx.dependency_cache().get::<u64>();
702 let retrieved = cached.unwrap_or(0);
703
704 results.lock().push((request_id, retrieved));
705 })
706 })
707 .collect();
708
709 for handle in handles {
711 handle.join().expect("Thread panicked");
712 }
713
714 let results = results.lock();
716 assert_eq!(results.len(), NUM_REQUESTS);
717
718 for (request_id, retrieved) in results.iter() {
719 assert_eq!(
720 request_id, retrieved,
721 "Request {request_id} should retrieve its own cached value, not another request's"
722 );
723 }
724 }
725
726 #[test]
727 #[allow(clippy::similar_names)]
728 fn resolution_stack_isolation_per_request() {
729 use crate::dependency::DependencyScope;
731
732 let cx1 = Cx::for_testing();
733 let cx2 = Cx::for_testing();
734
735 let ctx1 = RequestContext::new(cx1, 1);
736 let ctx2 = RequestContext::new(cx2, 2);
737
738 ctx1.resolution_stack()
740 .push::<i32>("i32", DependencyScope::Request);
741
742 let cycle1 = ctx1.resolution_stack().check_cycle::<i32>("i32");
744 assert!(cycle1.is_some(), "ctx1 should detect cycle for i32");
745
746 let cycle2 = ctx2.resolution_stack().check_cycle::<i32>("i32");
748 assert!(
749 cycle2.is_none(),
750 "ctx2 should NOT see ctx1's resolution stack"
751 );
752
753 ctx2.resolution_stack()
755 .push::<i32>("i32", DependencyScope::Request);
756 assert_eq!(ctx2.resolution_stack().depth(), 1);
757
758 assert_eq!(ctx1.resolution_stack().depth(), 1);
760 assert_eq!(ctx2.resolution_stack().depth(), 1);
761
762 ctx1.resolution_stack().pop();
764 ctx2.resolution_stack().pop();
765 assert!(ctx1.resolution_stack().is_empty());
766 assert!(ctx2.resolution_stack().is_empty());
767 }
768}