test: fix flaky QueryVirtualStorageTest.testQueryTooMuchData

[ykisana]

Fixes #20550.

Description

QueryVirtualStorageTest.testQueryTooMuchData runs a count(*) over ~3.7MB of segments against a 1MiB virtual storage cache and expects the query to fail for lack of space. The query does fail as intended, but the test only accepted the bundle-reservation message (Unable to reserve bundle[...]). Which reservation runs out of space first depends on load thread timing, so the test occasionally saw a different, equally correct CAPACITY_EXCEEDED error and failed at QueryVirtualStorageTest.java:166.

Why the error message varies

Native queries acquire segments with AcquireMode.FULL in a loop in ServerManager#getOrLoadSegmentReferences. Each acquireSegment call synchronously reserves the segment's partial metadata (SegmentLocalCacheManager#reservePartial, sized by virtualStorageMetadataReservationEstimate, 2KiB in this test) and then hands the bundle downloads to the load threads, which reserve bundles concurrently.

• Usually the space left over when a bundle fails to fit is larger than the metadata estimate. Every metadata reservation succeeds and the reported error is Unable to reserve bundle[...] for segment[...]; ensure enough disk space has been allocated to load all segments involved in the query.
• Occasionally, depending on which bundles the load threads win, the leftover is smaller than the estimate. The next metadata reservation fails with Unable to reserve partial metadata for segment[...]; ensure enough disk space has been allocated. That error is thrown from the acquire loop itself, so it is the one reported. It fails both of the old assertions, since it is not the bundle message and it does not end with "...to load all segments involved in the query".

I confirmed this locally. Raising the test's metadata estimate to 120KiB makes the partial-metadata outcome happen on every run. Under that setting the unchanged test fails with expected: <true> but was: <false> at line 166 (the same signature as the CI failure), and passes with this change. With the real 2KiB config the full class passes (4/4).

Relaxed the assertion in QueryVirtualStorageTest

testQueryTooMuchData now asserts on ensure enough disk space has been allocated, the text shared by every virtual storage capacity error (partial metadata, bundle, and full segment load). It still fails if the query errors for an unrelated reason, such as a timeout or a download failure. The assertion now passes t.getMessage() as its failure message. The CI failure only showed expected: <true> but was: <false>, which is why the actual error had to be reproduced locally.

Alternatives considered:

• Assert on the CAPACITY_EXCEEDED category. Not viable: by the time the error reaches the broker's HTTP response it is reported as "category":"UNCATEGORIZED", so only the message text is available to the test.
• Make the ordering deterministic, e.g. virtualStorageLoadThreads=1. Metadata is reserved on the query thread while bundles are reserved by load tasks, so this doesn't remove the race, and it would change the shared cluster that testQueryPartials relies on for its load-count expectations.
• Change the partial metadata message to end with "...to load all segments involved in the query" like the other two. That would be a user-facing message change made only to satisfy a test, so I left production code alone.

Added a unit test for the partial metadata capacity failure

SegmentLocalCacheManagerPartialAcquireTest#testFullAcquireFailsWhenPartialMetadataCannotBeReserved builds a manager whose only location is smaller than the metadata reservation estimate and checks that acquireSegment(..., AcquireMode.FULL):

• throws a CAPACITY_EXCEEDED DruidException;
• uses the Unable to reserve partial metadata for segment[...] message, including the shared "ensure enough disk space has been allocated" text the embedded test now relies on;
• leaves no cache entry behind.

No existing unit test covered this path. It is deterministic and runs in milliseconds, unlike the embedded race.

Release note

No user-facing changes. This PR only changes tests.

---

Key changed/added classes in this PR

• QueryVirtualStorageTest
• SegmentLocalCacheManagerPartialAcquireTest

---

This PR has:

• been self-reviewed.
• added comments explaining the "why" and the intent of the code wherever would not be obvious for an unfamiliar reader.
• added unit tests or modified existing tests to cover new code paths, ensuring the threshold for code coverage is met.

Generative AI usage

Generated by: Claude Opus 5.5