dwsmith1983 commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4434172065
##########
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:
Done in a71355db2. The barrier is gone. Each worker runs its loop under
`catch_unwind` and hands its reservation back with the outcome, so all eight
consumers stay registered until every thread has been joined, even when one
fails. The coordinator then re-raises the first failure. With a failing
assertion injected into one worker, the old test hung until it was killed, and
the new one fails at once with that worker's message and no follow-on
fair-limit panics from the others.
--
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]
