comphead commented on code in PR #7464:
URL: https://github.com/apache/datafusion-comet/pull/7464#discussion_r4470290017
##########
spark/src/main/java/org/apache/comet/cloud/s3/CometS3CredentialDispatcher.java:
##########
@@ -77,15 +86,39 @@ public static long ensureInitialized(
: Collections.unmodifiableMap(new HashMap<>(catalogProperties));
InstanceKey key =
new InstanceKey(providerClassName, dispatchKey == null ? "" :
dispatchKey, snapshot);
- return KEY_TO_HANDLE.computeIfAbsent(
- key,
- k -> {
- CometS3CredentialProvider provider =
instantiate(k.providerClassName);
- provider.initialize(k.catalogProperties);
- long handle = HANDLE_SEQ.getAndIncrement();
- INSTANCES.put(handle, new RegisteredProvider(provider, k));
- return handle;
- });
+ while (true) {
+ long handle =
+ KEY_TO_HANDLE.computeIfAbsent(
+ key,
+ k -> {
+ CometS3CredentialProvider provider =
instantiate(k.providerClassName);
+ provider.initialize(k.catalogProperties);
+ long h = HANDLE_SEQ.getAndIncrement();
+ INSTANCES.put(h, new RegisteredProvider(provider, k));
+ return h;
+ });
+ RegisteredProvider registered = INSTANCES.get(handle);
+ if (registered != null && registered.acquire()) {
+ return handle;
+ }
+ // The provider was retired after the lookup; drop its stale mapping and
register anew.
+ KEY_TO_HANDLE.remove(key, handle);
+ }
+ }
+
+ /**
+ * Invoked by native code when it drops the last reader built on {@code
handle}. Gives back one
+ * reference taken by {@link #ensureInitialized}, and closes the provider
when it was the last. An
+ * unknown handle, such as one already closed at shutdown, is ignored.
+ */
+ public static void release(long handle) {
+ RegisteredProvider registered = INSTANCES.get(handle);
+ if (registered == null || !registered.releaseLast()) {
+ return;
+ }
+ KEY_TO_HANDLE.remove(registered.key, handle);
+ INSTANCES.remove(handle, registered);
+ close(registered, "after its last reference was released");
Review Comment:
https://github.com/apache/datafusion-comet/pull/7363#discussion_r4459280742
still applies here: `close()` runs inline on whichever thread drops the last
bridge, and neither the `close()` Javadoc nor the Lifecycle section of
`s3-credential-providers.md` tells provider authors that it must not block. One
more case I did not mention there: when the cache holds the last reference, the
drop in `cached_file_io` runs on the task thread that just missed the cache,
inside `execute`. A slow `close()` of the superseded provider then delays the
read that already has the new credentials.
sunchao's shutdown race
(https://github.com/apache/datafusion-comet/pull/7363#discussion_r4458627539)
is also still open in this method and `closeAll`. The `refs` counter could be
the claim both paths share: `closeAll` closes an entry only when
`refs.getAndSet(RETIRED)` did not already return `RETIRED`, and a later
`releaseLast` then sees `RETIRED` and returns false. That is the same guard
Iceberg's `S3FileIO.close()` uses with `isResourceClosed.compareAndSet`. I have
not run it.
##########
native/core/src/execution/operators/iceberg_common.rs:
##########
@@ -149,11 +149,46 @@ impl FileIoCacheKey {
properties: sorted_properties(catalog_properties),
})
}
+
+ /// Whether `other` is a `FileIO` for the same table, access mode and
catalog built from
+ /// different properties, which this key's `FileIO` supersedes. A catalog
that re-vends
+ /// credentials gives a table new properties, and a commit moves its
metadata file, so the table
+ /// is compared by its root: the reference path up to `/metadata/`. A
commit that keeps the
+ /// properties supersedes nothing, so neighboring snapshots share one
provider.
+ ///
+ /// When both carry `s3.session-token-expires-at-ms`, only a later expiry
supersedes. Two
+ /// sessions reading one table with different vended tokens then keep both
entries instead of
+ /// evicting each other on every insert; the LRU bound retires the older
one.
+ fn supersedes(&self, other: &Self) -> bool {
+ self.access_mode == other.access_mode
+ && self.catalog_name == other.catalog_name
+ && table_root(&self.reference_path) ==
table_root(&other.reference_path)
+ && self.properties != other.properties
+ && match (self.token_expires_at_ms(), other.token_expires_at_ms())
{
+ (Some(new), Some(old)) => new > old,
+ _ => true,
+ }
+ }
+
+ fn token_expires_at_ms(&self) -> Option<u64> {
+ self.properties
+ .iter()
+ .find(|(k, _)| k == "s3.session-token-expires-at-ms")
+ .and_then(|(_, v)| v.parse().ok())
+ }
+}
+
+/// The table location of a metadata file path, or the path itself when it
isn't one.
+fn table_root(reference_path: &str) -> &str {
+ reference_path
+ .rfind("/metadata/")
+ .map_or(reference_path, |i| &reference_path[..i])
}
Review Comment:
The custom `write.metadata.path` case from
https://github.com/apache/datafusion-comet/pull/7363#discussion_r4459280729
still applies here, since `table_root` still cuts at the last `/metadata/`.
With `write.metadata.path=s3://b/meta/db/t1` each commit gets a new root, so
the old `FileIO` and its provider wait for the LRU bound. With
`s3://b/metadata/db/t1` every table laid out that way gets the root `s3://b`,
so two of them in one query evict each other. Comparing the metadata file's
parent directory would handle both, and a case for each layout in
`cache_evicts_the_file_io_a_table_supersedes` would pin it. The design doc
sentence about "the reference path up to `/metadata/`" would change to match.
##########
spark/src/test/java/org/apache/comet/cloud/s3/CometS3CredentialDispatcherTest.java:
##########
@@ -267,4 +272,105 @@ public void closeAllSwallowsProviderExceptions() {
// Both close() invocations ran even though one (the first to be visited)
threw.
assertEquals(2, TestCometS3CredentialProvider.closeCount.get());
}
+
+ @Test
+ public void releasingTheLastReferenceClosesTheProvider() throws Exception {
+ long h1 = init();
+ long h2 = init();
+ assertEquals(h1, h2);
+
+ CometS3CredentialDispatcher.release(h1);
+ assertEquals(0, TestCometS3CredentialProvider.closeCount.get());
+ CometS3CredentialDispatcher.getCredentialsForPath(h2, "b", "k", READ);
+
+ CometS3CredentialDispatcher.release(h2);
+ assertEquals(1, TestCometS3CredentialProvider.closeCount.get());
+ assertEquals(0, CometS3CredentialDispatcher.registeredCount());
+ Exception thrown =
+ assertThrows(
+ IllegalStateException.class,
+ () -> CometS3CredentialDispatcher.getCredentialsForPath(h2, "b",
"k", READ));
+ assertTrue(thrown.getMessage().contains("not initialized"));
+ }
+
+ @Test
+ public void aRetiredKeyGetsANewProvider() {
+ long retired = init();
+ CometS3CredentialDispatcher.release(retired);
+
+ long fresh = init();
+
+ assertNotEquals(retired, fresh);
+ assertEquals(2, TestCometS3CredentialProvider.initCount.get());
+ assertEquals(1, CometS3CredentialDispatcher.registeredCount());
+ }
+
+ /** An extra release, or one for a handle closed at shutdown, changes
nothing. */
+ @Test
+ public void unknownOrRepeatedReleasesAreIgnored() {
+ CometS3CredentialDispatcher.release(Long.MAX_VALUE);
+ long handle = init();
+ CometS3CredentialDispatcher.release(handle);
+ CometS3CredentialDispatcher.release(handle);
+
+ assertEquals(1, TestCometS3CredentialProvider.closeCount.get());
+ long other = init();
+ CometS3CredentialDispatcher.release(handle);
+ assertEquals(1, CometS3CredentialDispatcher.registeredCount());
+ CometS3CredentialDispatcher.release(other);
+ assertEquals(2, TestCometS3CredentialProvider.closeCount.get());
+ }
+
+ /**
+ * A REST catalog that re-vends credentials for a table changes its FileIO
properties, so each
+ * snapshot registers its own provider. Once native code releases a
snapshot's handle, that
+ * provider is closed, and only the snapshots still in use stay registered.
+ */
+ @Test
+ public void supersededSnapshotsAreClosedOnceReleased() {
Review Comment:
The overlap from
https://github.com/apache/datafusion-comet/pull/7363#discussion_r4459280752
still holds. Now that `CometS3CredentialBridgeSuite` checks `registeredCount()`
across two `FileIO`s for one table, this test covers nothing that
`distinctCatalogPropertiesIsolateInstances`,
`releasingTheLastReferenceClosesTheProvider` and that suite don't. Could it be
removed?
--
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]
