Skip to content

Fix query memory reservation deregistration race#18272

Open
Caideyipi wants to merge 1 commit into
apache:masterfrom
Caideyipi:fix/query-memory-reservation-deregister-race
Open

Fix query memory reservation deregistration race#18272
Caideyipi wants to merge 1 commit into
apache:masterfrom
Caideyipi:fix/query-memory-reservation-deregister-race

Conversation

@Caideyipi

Copy link
Copy Markdown
Collaborator

Description

A table standalone integration test exposed a race in MemoryPool:

SharedTsBlockQueue.remove()
  -> MemoryPool.free()
  -> NullPointerException because the fragment reservation map no longer exists

deRegisterFragmentInstanceFromQueryMemoryMap checked whether a fragment's reservations were zero before taking the map-removal lock. A concurrent reservation could update the fragment map after that zero snapshot but before deregistration removed it. The memory block was then charged while its reservation map was detached, so the downstream consumer failed when freeing the buffered TsBlock.

This is independent of the Pipe changes in #18266. That PR's head was based directly on b8ff9eeaea0d2e5afd79446e8bb4cf654c8c8038 (origin/master at the time), and its diff only touched 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 no operation updates a detached map.
  • Keep reservations for different fragments concurrent.

Tests

Added a deterministic concurrency test that pauses Map.merge while deregistration runs. With the origin/master implementation, it fails because the reservation map is removed (expected 256 but was 0). With this change:

  • MemoryPoolTest: 23 tests run, 0 failures
  • mvn spotless:apply -pl iotdb-core/datanode: success
  • Checkstyle reported 0 violations before compilation was blocked by unrelated local cross-module/generated-source state

The local full reactor could not complete because the Windows Thrift 0.14 generator emitted javax.annotation.Generated while this branch uses Jakarta annotations; the changed production and test sources were therefore compiled directly with JDK 17 against the Maven-resolved datanode test classpath for the targeted JUnit run.


This PR has:

  • been self-reviewed.
  • added tests that fail without the fix.

@shuwenwei

Copy link
Copy Markdown
Member

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants