sunchao commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4433048625
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -532,6 +900,487 @@ mod tests {
assert_eq!(fake.released(), vec![100]);
}
+ #[test]
+ fn short_grant_is_released_and_reported_as_error() {
+ let stub = Arc::new(StubTaskMemory::new(GIB).short_every(1));
+ let pool = pool_with(&stub, 1_000);
+ let res = MemoryConsumer::new("consumer").register(&pool);
+
+ let err = res.try_grow(100).unwrap_err();
+ assert!(err.to_string().contains("only got"), "{err}");
+ assert_eq!(pool.reserved(), 0);
+ assert_eq!(stub.outstanding(), 0, "partial grant must be handed back");
+ }
+
+ #[test]
+ fn acquire_failure_leaves_accounting_unchanged() {
+ let stub = Arc::new(StubTaskMemory::new(GIB));
+ let pool = pool_with(&stub, 1_000);
+ let res = MemoryConsumer::new("consumer").register(&pool);
+
+ stub.fail_acquire.store(true, SeqCst);
+ assert!(res.try_grow(100).is_err());
+ assert_eq!(pool.reserved(), 0);
+ assert_eq!(stub.outstanding(), 0);
+ }
+
+ #[test]
+ fn panicking_try_grow_records_nothing_and_frees_the_pool() {
+ let stub = Arc::new(StubTaskMemory::new(GIB));
+ let pool = pool_with(&stub, 1_000);
+ let res = MemoryConsumer::new("consumer").register(&pool);
+ res.try_grow(100).unwrap();
+
+ stub.panic_acquire.store(true, SeqCst);
+ let panic = catch_unwind(AssertUnwindSafe(|| res.try_grow(600)));
+ assert!(panic.is_err(), "the bridge panic must propagate");
+ assert_eq!(pool.reserved(), 100, "the panicked grow left bytes
behind");
+
+ stub.panic_acquire.store(false, SeqCst);
+ let (res, result) = try_grow_on_another_thread(res, 600);
+ result.unwrap();
+ assert_eq!(pool.reserved(), 700);
+ res.free();
+ assert_eq!(stub.outstanding(), 0);
+ }
+
+ #[test]
+ fn panicking_grow_records_nothing_and_frees_the_pool() {
+ let stub = Arc::new(StubTaskMemory::new(GIB));
+ let pool = pool_with(&stub, 1_000);
+ let res = spillable("consumer", &pool);
+ res.grow(100);
+
+ stub.panic_acquire.store(true, SeqCst);
+ let panic = catch_unwind(AssertUnwindSafe(|| res.grow(600)));
+ assert!(panic.is_err(), "the bridge panic must propagate");
+ assert_eq!(pool.reserved(), 100, "the panicked grow left bytes
behind");
+ assert_eq!(res.size(), 100);
+
+ stub.panic_acquire.store(false, SeqCst);
+ let (res, result) = try_grow_on_another_thread(res, 900);
+ result.expect("the consumer's share still holds the panicked grow");
+ assert_eq!(pool.reserved(), 1_000);
+ res.free();
+ assert_eq!(stub.outstanding(), 0);
+ }
+
+ #[test]
+ fn concurrent_try_grow_and_shrink_keep_accounting_consistent() {
+ const THREADS: usize = 8;
+ const ITERS: usize = 500;
+ const POOL_SIZE: usize = 80_000;
+
+ let stub = Arc::new(StubTaskMemory::new(GIB).short_every(7));
+ let pool = pool_with(&stub, POOL_SIZE);
+
+ // Register every consumer up front so the fair limit stays fixed
while threads run.
+ let reservations: Vec<_> = (0..THREADS)
+ .map(|t| spillable(&format!("consumer-{t}"), &pool))
+ .collect();
+ let barrier = Arc::new(Barrier::new(THREADS));
+
+ let handles: Vec<_> = reservations
+ .into_iter()
+ .enumerate()
+ .map(|(t, res)| {
+ let pool = Arc::clone(&pool);
+ let barrier = Arc::clone(&barrier);
+ thread::spawn(move || {
+ for i in 0..ITERS {
+ let size = 1 + (i * 37 + t * 101) % 509;
+ // Refusals are expected; only the invariants below
must hold.
+ let _ = res.try_grow(size);
+ if i % 3 == 0 && res.size() > 0 {
+ res.shrink(res.size() / 2 + 1);
+ }
+ // Drains the task's balance to zero under acquires in
flight.
+ if (i + t) % 41 == 0 {
+ res.free();
+ }
+ // One reservation per consumer, so its size is the
consumer's share.
+ assert!(
+ res.size() <= POOL_SIZE / THREADS,
+ "consumer exceeded its fair limit"
+ );
+ assert!(pool.reserved() <= POOL_SIZE, "pool exceeded
its size");
+ }
+ // Stay registered until every thread stops so the fair
limit is fixed.
+ barrier.wait();
Review Comment:
If a worker fails one of the invariants above, it panics before reaching
this barrier. The other workers then wait forever, and the coordinator can hang
joining one of them instead of reporting the failed assertion. Could the
workers return their still-registered reservations to the coordinator, which
can keep them alive until all workers have finished? That preserves the fixed
consumer count without putting a barrier after a fallible assertion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
