Harden query memory reservation deregistration ordering

[Caideyipi]

Description

A table standalone integration test exposed this failure at the query-memory boundary:

SharedTsBlockQueue.remove()
  -> MemoryPool.free()
  -> the fragment reservation map has already been removed

The failure proves that a free reached MemoryPool after deregistration. It does not by itself prove that a new reservation raced with FI deregistration.

After auditing the production FI terminal path, the normal ordering is:

1. clearShuffleSinkHandle();
2. driver.close();
3. releaseResourceWhenAllDriversAreClosed();
4. deregister the FI from MemoryPool.

Handle cleanup also cancels or completes pending reservation futures under the same future lock used when MemoryPool.free() grants them. I did not find a concrete normal-path caller that starts a new reservation after FI cleanup. The change below is therefore defensive hardening of the MemoryPool API boundary; it does not claim that the failure log demonstrates a production reserve/deregister race.

This is independent of the Pipe changes in #18266. That PR's head was based directly on b8ff9eeaea0d2e5afd79446e8bb4cf654c8c8038 and touched only PipeMemoryManager and its test.

Fix

• Serialize register, reserve, rollback, free, and deregister operations on the affected fragment reservation map.
• Revalidate the query/fragment map identities after acquiring the fragment lock so an operation never updates a detached map.
• If deregistration wins and removes the mapping, fail explicitly with IllegalArgumentException before allocating or releasing memory instead of retrying into a null dereference.
• Keep reservations for different fragments concurrent.

Tests

Added deterministic coverage for both lock orderings:

• reservation enters the fragment map first, so deregistration waits and observes the reservation;
• deregistration removes the map first, then the waiting reservation resumes, fails explicitly, and leaves the global memory usage unchanged.

Executed:

• mvn spotless:apply -pl iotdb-core/datanode
• mvn -o -nsu '-Ddevelocity.off=true' test -pl iotdb-core/datanode '-Dtest=MemoryPoolTest'
  • 24 tests, 0 failures, 0 errors
  • Checkstyle: 0 violations
  • Spotless check: passed
• mvn -o -nsu '-Ddevelocity.off=true' test-compile -pl iotdb-core/datanode -P with-zh-locale '-DskipTests'
  • BUILD SUCCESS

---

This PR has:

• been self-reviewed.
• added deterministic tests for both concurrency orderings.

[shuwenwei]

Could you clarify which production interleaving this test is intended to reproduce? The failure log only proves that SharedTsBlockQueue.remove() called MemoryPool.free() after the FI/query reservation map had already been removed; it does not show that reservation was concurrent with deregistration. In the normal FI terminal path, clearShuffleSinkHandle(), driver.close(), and releaseResourceWhenAllDriversAreClosed() all happen before deregistration, while handle cleanup also cancels pending reservation futures using the same future lock used by MemoryPool.free(). The new test directly races tryReserveForTest() with deregistration and therefore bypasses these lifecycle guarantees. There is also an uncovered reverse interleaving in the proposed fix: if deregistration acquires the fragment lock first and removes the map, the waiting operation fails identity validation and continues, but the next iteration dereferences a null queryRelatedMemory, resulting in another NPE. Could we add a deterministic test where deregistration wins first, and either document a concrete production caller that can still reserve after FI cleanup or handle the missing mapping explicitly?

[Caideyipi]

Good catch. I audited the production lifecycle and agree that the original description overstated what the failure log proves. In the normal FI terminal path, sink/driver/source cleanup happens before deregistration, and pending reservations are cancelled or completed under the future lock. I did not find a concrete normal-path caller that starts a new reservation after FI cleanup.

I addressed the actionable reverse ordering in 7af40af:

• added explicit missing query/fragment-map handling in free, tryReserve, and rollback, so a deregistration-first interleaving fails with a descriptive IllegalArgumentException instead of retrying into an NPE;
• added a deterministic test that lets the reservation read the old fragment map, lets deregistration remove it first, and then resumes the reservation; it verifies that the operation fails explicitly and global reserved bytes remain zero;
• kept the existing deterministic reservation-first test as defensive coverage of the MemoryPool API invariant.

I also revised the PR title/description to distinguish this defensive hardening from the still-unproven production interleaving.

Verification passed: MemoryPoolTest 24/24, Checkstyle, Spotless, and DataNode test-compile with the zh locale.