FrankChen021 commented on code in PR #19671:
URL: https://github.com/apache/druid/pull/19671#discussion_r3558725842
##########
server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java:
##########
@@ -872,6 +937,290 @@ private record ReservedPartial(
{
}
+ /**
+ * Apply (or reconcile) a partial-load rule for a segment coming through
{@link #load}. The rule state lives on the
+ * {@link PartialSegmentMetadataCacheEntry} itself; this method just drives
it:
+ * <ol>
+ * <li>Materialize the wrapper's {@link PartialLoadSpec} and open a range
reader over the backend. If the backend
+ * can't do range reads, clear any prior rule on the segment and
return; the segment falls back to the
+ * weak-full-load path at query time.</li>
+ * <li>Acquire (or reserve) the metadata entry via the shared on-demand
path ({@link #findOrReservePartial}),
+ * take a transient eviction-protective hold across the rest of the
flow.</li>
+ * <li>Mount the metadata entry (idempotent).</li>
+ * <li>Compute the rule's selected bundle names from the wrapper + parsed
segment metadata, and call
+ * {@link PartialSegmentMetadataCacheEntry#applyRule} which installs
the metadata self-hold and takes
+ * holds on any already-registered selected bundles.</li>
+ * <li>Kick off async eager downloads on the loading pool for any selected
bundle not yet linked. Each async task
+ * drives {@link
PartialSegmentMetadataCacheEntry#ensureBundleResidentForRule}: its bundle mount
will call
+ * {@link PartialSegmentMetadataCacheEntry#registerBundle}, which
acquires the rule-hold on that bundle.</li>
+ * <li>Release the transient hold; the metadata's self-hold keeps the
entry resident under the applied rule.</li>
+ * </ol>
+ */
+ private void loadPartial(DataSegment dataSegment) throws
SegmentLoadingException
+ {
+ final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment);
+ final SegmentRangeReader rangeReader = openPartialRangeReader(dataSegment,
wrapper);
+ final ReferenceCountingLock lock = lock(dataSegment);
+ synchronized (lock) {
+ try {
+ // If a stale non-partial cache entry sits at this segment id (a
CompleteSegmentCacheEntry created by a prior
+ // acquireSegment while virtualStoragePartialDownloadsEnabled=false,
for example), evict it before any
+ // partial-entry lookup or reservation; otherwise
findExistingPartialWithHold's defensive type-check would
+ // throw, and reservePartial's addWeakReservationHold would land on
the incompatible entry. If the stale
+ // entry is currently held (in-flight query), this throws a retryable
SegmentLoadingException; the
+ // coordinator's load queue retries on next sync, and by then the
query should have released.
+ evictStaleNonPartialWeakEntry(dataSegment.getId());
+
+ if (rangeReader == null) {
+ // Backend doesn't support range reads (e.g. zipped deep storage).
The rule can't be honored as a partial
+ // load; clear any prior rule so the segment falls through to the
ordinary weak-full-load path at query
+ // time.
+ final ReservedPartial existing =
findExistingPartialWithHold(dataSegment.getId());
+ if (existing != null) {
+ try {
+ existing.metadata().clearRule();
+ log.warn(
+ "Backend for segment[%s] does not support range reads;
released rule[fingerprint=%s], segment "
+ + "will fall back to weak full-load at query time",
+ dataSegment.getId(),
+ wrapper.getFingerprint()
+ );
+ }
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(
+ existing.hold(),
+ t -> log.warn(t, "Failed to release transient hold on
partial metadata for segment[%s]",
+ dataSegment.getId())
+ );
+ }
+ } else {
+ // No prior entry AND range-reader is null on this fresh load: the
coordinator asked for a partial rule
+ // this historical can't honor. Log so operators can debug "why
isn't my rule applied on historical X".
+ log.warn(
+ "Backend for segment[%s] does not support range reads and no
prior partial entry exists; rule"
+ + "[fingerprint=%s] cannot be applied on this historical",
+ dataSegment.getId(),
+ wrapper.getFingerprint()
+ );
+ }
+ return;
+ }
+
+ final ReservedPartial reserved = findOrReservePartial(dataSegment,
rangeReader);
+ try {
+ final PartialSegmentMetadataCacheEntry metadata =
reserved.metadata();
+ try {
+ metadata.mount(reserved.location());
+ }
+ catch (IOException e) {
+ throw new SegmentLoadingException(
+ e,
+ "Failed to mount partial metadata for segment[%s]",
+ dataSegment.getId()
+ );
+ }
+ final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper();
+ if (mapper == null) {
+ throw DruidException.defensive(
+ "Partial metadata for segment[%s] mounted without a file
mapper",
+ dataSegment.getId()
+ );
+ }
+ // Register the info-file cleanup hook BEFORE anything that could
throw. Any throw between mount and
+ // awaitEagerDownloadsOrClearRule leaves the metadata entry
weak-reserved with the hook attached; when
+ // cache eventually reclaims, the info file gets deleted along with
the entry.
+ metadata.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
+ final Set<String> selected = Set.copyOf(
+ wrapper.getSelectedBundleNames(dataSegment,
mapper.getSegmentFileMetadata())
+ );
+ final String priorFingerprint = metadata.getRuleFingerprint();
+ metadata.applyRule(wrapper.getFingerprint(), selected);
+ if (priorFingerprint != null &&
!priorFingerprint.equals(wrapper.getFingerprint())) {
+ log.info(
+ "Reconciled partial-load rule for segment[%s]: fingerprint
transitioned [%s] → [%s]",
+ dataSegment.getId(),
+ priorFingerprint,
+ wrapper.getFingerprint()
+ );
+ }
+ // Block until every eager download completes so the announcement
fingerprint reflects reality: any failure
+ // clears the rule state (releasing self-hold + all bundle
rule-holds) and propagates as a load failure so
+ // the coordinator's load queue can retry on its next sync. The
announced fingerprint == "rule fully
+ // realized" contract stays intact.
+ awaitEagerDownloadsOrClearRule(dataSegment, metadata, selected);
+ }
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(
+ reserved.hold(),
+ t -> log.warn(t, "Failed to release transient hold on partial
metadata for segment[%s]",
+ dataSegment.getId())
+ );
+ }
+ }
+ finally {
+ unlock(dataSegment, lock);
+ }
+ }
+ }
+
+ /**
+ * Submit eager-download tasks to the loading pool for every rule-selected
bundle not yet registered with
+ * {@code metadata}, then block until every task completes. On any failure
the rule state is cleared
+ * and a {@link SegmentLoadingException} is thrown so the caller treats the
load as failed and retries.
+ */
+ private void awaitEagerDownloadsOrClearRule(
+ DataSegment dataSegment,
+ PartialSegmentMetadataCacheEntry metadata,
+ Set<String> selected
+ ) throws SegmentLoadingException
+ {
+ final List<Future<?>> pending = new ArrayList<>();
+ Throwable firstFailure = null;
+ try {
+ for (String bundleName : selected) {
+ // Skip only when the bundle is BOTH rule-held AND fully downloaded
(every container's files present on disk).
+ if (metadata.isBundleRuleHeld(bundleName) &&
metadata.isBundleFullyDownloaded(bundleName)) {
+ continue;
+ }
+
pending.add(virtualStorageLoadingThreadPool.getExecutorService().submit(() -> {
+ metadata.ensureBundleResidentForRule(bundleName);
Review Comment:
[P1] Download transitive bundle dependencies
`ensureBundleResidentForRule` downloads only the selected bundle, although
mounting a projection may infer `__base` as a required dependency. The
dependency is sparse-mounted and held, but its files are never eagerly
downloaded; the otherwise-unused `bundlesInMountOrder` confirms this missing
expansion. Queries can therefore perform on-demand deep-storage reads after the
rule was reported fully realized, and fail if deep storage is unavailable.
##########
server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java:
##########
@@ -927,6 +1276,15 @@ public void load(final DataSegment dataSegment) throws
SegmentLoadingException
"load() should not be called when virtualStorageIsEphemeral is
true"
);
}
+ // Partial-load-rule routing: a wrapped load spec carrying a fingerprint
+ per-spec selection (cluster groups,
+ // projections, ...) means the coordinator wants this segment loaded
with rule holds pinning the selected
+ // bundles + metadata. Install the holds via loadPartial; other segments
(no wrapper, or partials disabled)
+ // take the existing weak/no-op path below.
+ if (config.isVirtualStoragePartialDownloadsEnabled()
+ && PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec()))
{
+ loadPartial(dataSegment);
Review Comment:
[P1] Announce the realized partial footprint
After this new partial path succeeds, `SegmentLoadDropHandler` still
announces the original `DataSegment`, and unchanged
`SegmentChangeRequestLoad.forAnnouncement` stamps `loadedBytes =
segment.getSize()`. `HttpServerInventoryView` consequently classifies every
successful partial load as a full fallback, overstates server disk usage, and
can make coordinator capacity accounting reject otherwise-valid placements. The
successful path needs to expose its realized loaded bytes to announcement
generation.
##########
server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java:
##########
@@ -199,6 +230,322 @@ File getLocalCacheDir()
return localCacheDir;
}
+ /**
+ * The fingerprint of the partial-load rule currently applied to this entry,
or {@code null} if no rule has been
+ * applied. Set by {@link #applyRule(String, Set)} and cleared by {@link
#clearRule()}. {@link #doActualUnmount()}
+ * asserts that this field is {@code null} at unmount time, the {@code
metadataSelfHold} taken by applyRule pins
+ * the entry's phaser and prevents eviction reclaim / doActualUnmount from
firing while a rule is applied. A fresh
+ * mount therefore always starts with no rule applied.
+ */
+ @Nullable
+ public String getRuleFingerprint()
+ {
+ entryLock.lock();
+ try {
+ return ruleFingerprint;
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * The set of bundle names currently pinned by an applied partial-load rule,
or an empty set if no rule is applied.
+ * Return value is immutable; safe to inspect without further
synchronization.
+ */
+ public Set<String> getRuleSelectedBundleNames()
+ {
+ entryLock.lock();
+ try {
+ return ruleSelectedBundleNames;
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * Whether a partial-load rule is currently applied to this entry.
Equivalent to {@code getRuleFingerprint() != null}.
+ * While {@code true}, this entry holds a self-referential {@link
StorageLocation.ReservationHold} that prevents cache
+ * from evicting it, so the metadata's V10 header stays resident until
{@link #clearRule()} runs.
+ */
+ public boolean isRuleHeld()
+ {
+ entryLock.lock();
+ try {
+ return ruleFingerprint != null;
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * Apply (or replace) a partial-load rule on this entry. Must be called
after {@link #mount(StorageLocation)}, the
+ * entry must be registered with its {@link StorageLocation} so it can
acquire holds on itself and on its bundles.
+ * <p>
+ * <b>Concurrency contract.</b> The caller MUST serialize {@code applyRule}
and {@link #clearRule} for a given
+ * segment id through an external per-segment lock ({@code
SegmentLocalCacheManager} uses its
+ * {@code ReferenceCountingLock segmentLocks}). This method releases {@link
#entryLock} between Phase 2 (acquire
+ * holds outside the lock) and Phase 4 (release holds outside the lock); a
concurrent invocation of
+ * {@code applyRule}/{@code clearRule} on the same entry can observe
intermediate state during those windows and
+ * silently miss rule-hold installations or leak them. {@link #entryLock}
alone is insufficient because it is
+ * released across the acquire/release phases by design (holding it across
{@link StorageLocation} lock acquisition
+ * would invert the writeLock → entryLock ordering the storage layer
already uses for eviction).
+ * <p>
+ * On first application ({@code ruleFingerprint} transitions from {@code
null}), a self-referential
+ * {@link StorageLocation.ReservationHold} is taken to pin the metadata
entry in-cache. On repeat calls with the same
+ * {@code fingerprint} and matching {@code selectedBundleNames}, this is a
no-op. On a genuine rule swap, only the
+ * delta between the previous and new selection is applied: bundle holds for
names dropped from the selection are
+ * closed, and holds for names newly added are acquired for any bundle
already registered with this entry (later
+ * arrivals get their hold via {@link #registerBundle}).
+ * <p>
+ * Bundles whose {@link StorageLocation.ReservationHold} is currently
zero-refcount (never mounted, or evicted but
+ * unregistered before this call) will not appear in {@link #linkedBundles}
yet, so no hold is taken for them here;
+ * the caller is expected to drive an on-demand acquire per selected bundle
name to trigger a fresh mount, which will
+ * call {@link #registerBundle} and pick up the rule-hold at that point.
+ *
+ * @throws DruidException if the entry is not currently mounted
+ */
+ public void applyRule(String fingerprint, Set<String> selectedBundleNames)
+ {
+ Objects.requireNonNull(fingerprint, "fingerprint");
+ Objects.requireNonNull(selectedBundleNames, "selectedBundleNames");
+ final Set<String> newSelection = Set.copyOf(selectedBundleNames);
+
+ // Phase 1: snapshot under entryLock and compute the diff. Do NOT call
StorageLocation methods here — that would
+ // acquire the location's readLock while holding entryLock and could
deadlock with a concurrent writeLock holder
+ // that needs entryLock
+ final StorageLocation loc;
+ final boolean needsSelfHold;
+ final List<String> namesToRelease;
+ final List<String> namesToAcquire;
+ entryLock.lock();
+ try {
+ if (location == null) {
+ throw DruidException.defensive(
+ "applyRule on partial metadata entry[%s] requires the entry to be
mounted",
+ id
+ );
+ }
+ if (fingerprint.equals(ruleFingerprint) &&
newSelection.equals(ruleSelectedBundleNames)) {
+ return;
+ }
+ loc = location;
+ needsSelfHold = (metadataSelfHold == null);
+ namesToRelease = new ArrayList<>();
+ for (String name : ruleBundleHolds.keySet()) {
+ if (!newSelection.contains(name)) {
+ namesToRelease.add(name);
+ }
+ }
+ namesToAcquire = new ArrayList<>();
+ for (String name : newSelection) {
+ if (!ruleBundleHolds.containsKey(name) && findLinkedBundleByName(name)
!= null) {
Review Comment:
[P2] Close the concurrent registration gap
A query can insert a newly mounted bundle into `linkedBundles` after this
snapshot but before `ruleSelectedBundleNames` is committed. `registerBundle`
observes the old selection and skips its hold, while the eager reconciliation
later acquires the already-mounted bundle without invoking `registerBundle`
again. The selected bundle remains permanently absent from `ruleBundleHolds`
and can be SIEVE-evicted despite the applied rule.
##########
server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java:
##########
@@ -872,6 +937,290 @@ private record ReservedPartial(
{
}
+ /**
+ * Apply (or reconcile) a partial-load rule for a segment coming through
{@link #load}. The rule state lives on the
+ * {@link PartialSegmentMetadataCacheEntry} itself; this method just drives
it:
+ * <ol>
+ * <li>Materialize the wrapper's {@link PartialLoadSpec} and open a range
reader over the backend. If the backend
+ * can't do range reads, clear any prior rule on the segment and
return; the segment falls back to the
+ * weak-full-load path at query time.</li>
+ * <li>Acquire (or reserve) the metadata entry via the shared on-demand
path ({@link #findOrReservePartial}),
+ * take a transient eviction-protective hold across the rest of the
flow.</li>
+ * <li>Mount the metadata entry (idempotent).</li>
+ * <li>Compute the rule's selected bundle names from the wrapper + parsed
segment metadata, and call
+ * {@link PartialSegmentMetadataCacheEntry#applyRule} which installs
the metadata self-hold and takes
+ * holds on any already-registered selected bundles.</li>
+ * <li>Kick off async eager downloads on the loading pool for any selected
bundle not yet linked. Each async task
+ * drives {@link
PartialSegmentMetadataCacheEntry#ensureBundleResidentForRule}: its bundle mount
will call
+ * {@link PartialSegmentMetadataCacheEntry#registerBundle}, which
acquires the rule-hold on that bundle.</li>
+ * <li>Release the transient hold; the metadata's self-hold keeps the
entry resident under the applied rule.</li>
+ * </ol>
+ */
+ private void loadPartial(DataSegment dataSegment) throws
SegmentLoadingException
+ {
+ final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment);
+ final SegmentRangeReader rangeReader = openPartialRangeReader(dataSegment,
wrapper);
+ final ReferenceCountingLock lock = lock(dataSegment);
+ synchronized (lock) {
+ try {
+ // If a stale non-partial cache entry sits at this segment id (a
CompleteSegmentCacheEntry created by a prior
+ // acquireSegment while virtualStoragePartialDownloadsEnabled=false,
for example), evict it before any
+ // partial-entry lookup or reservation; otherwise
findExistingPartialWithHold's defensive type-check would
+ // throw, and reservePartial's addWeakReservationHold would land on
the incompatible entry. If the stale
+ // entry is currently held (in-flight query), this throws a retryable
SegmentLoadingException; the
+ // coordinator's load queue retries on next sync, and by then the
query should have released.
+ evictStaleNonPartialWeakEntry(dataSegment.getId());
+
+ if (rangeReader == null) {
+ // Backend doesn't support range reads (e.g. zipped deep storage).
The rule can't be honored as a partial
+ // load; clear any prior rule so the segment falls through to the
ordinary weak-full-load path at query
+ // time.
+ final ReservedPartial existing =
findExistingPartialWithHold(dataSegment.getId());
+ if (existing != null) {
+ try {
+ existing.metadata().clearRule();
+ log.warn(
+ "Backend for segment[%s] does not support range reads;
released rule[fingerprint=%s], segment "
+ + "will fall back to weak full-load at query time",
+ dataSegment.getId(),
+ wrapper.getFingerprint()
+ );
+ }
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(
+ existing.hold(),
+ t -> log.warn(t, "Failed to release transient hold on
partial metadata for segment[%s]",
+ dataSegment.getId())
+ );
+ }
+ } else {
+ // No prior entry AND range-reader is null on this fresh load: the
coordinator asked for a partial rule
+ // this historical can't honor. Log so operators can debug "why
isn't my rule applied on historical X".
+ log.warn(
+ "Backend for segment[%s] does not support range reads and no
prior partial entry exists; rule"
+ + "[fingerprint=%s] cannot be applied on this historical",
+ dataSegment.getId(),
+ wrapper.getFingerprint()
+ );
+ }
+ return;
+ }
+
+ final ReservedPartial reserved = findOrReservePartial(dataSegment,
rangeReader);
Review Comment:
[P2] Persist rule changes when reusing metadata
`findOrReservePartial` returns an existing entry without passing through
`reservePartial`, the only partial path that rewrites the info file. A
successful fingerprint/selection change is therefore applied in memory while
disk retains the prior wrapped load spec; after restart, bootstrap re-applies
and announces the old rule until the coordinator corrects it.
--
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]
