laxmanmamididd commented on code in PR #7363:
URL: https://github.com/apache/datafusion-comet/pull/7363#discussion_r4475643683
##########
native/core/src/execution/operators/iceberg_common.rs:
##########
@@ -149,11 +149,30 @@ 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, or a commit that moves the metadata file, gives a table
new properties or a new
+ /// reference path, so the table is compared by its root: the reference
path up to `/metadata/`.
+ 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 != other
+ }
+}
+
+/// 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:
This moved to #7464 with the rest of the lifecycle change. It's fixed there
in 40963cfe6, and the reply is on the matching thread:
https://github.com/apache/datafusion-comet/pull/7464#discussion_r4475641962
##########
spark/src/main/java/org/apache/comet/cloud/s3/CometS3CredentialProvider.java:
##########
@@ -50,8 +50,10 @@ default void initialize(Map<String, String>
catalogProperties) {}
CometS3Credentials getCredentialsForPath(CometS3CredentialContext context)
throws Exception;
/**
- * Invoked from the dispatcher's best-effort JVM shutdown hook. Default
no-op suits stateless
- * providers; override to release HTTP clients, refresh executors, etc.
+ * Invoked once nothing reads through this instance any more, such as after
an Iceberg catalog
+ * re-vends credentials for a table, and from the dispatcher's best-effort
JVM shutdown hook for
+ * instances still in use then. Default no-op suits stateless providers;
override to release HTTP
+ * clients, refresh executors, etc.
*/
Review Comment:
This moved to #7464 with the rest of the lifecycle change. It's fixed there
in 40963cfe6, and the reply is on the matching thread:
https://github.com/apache/datafusion-comet/pull/7464#discussion_r4475642244
##########
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() {
+ long[] handles = new long[3];
+ for (int i = 0; i < handles.length; i++) {
+ Map<String, String> props = new HashMap<>();
+ props.put("s3.session-token", "token-" + i);
+ handles[i] =
CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, DK, props);
+ }
+ assertEquals(3, CometS3CredentialDispatcher.registeredCount());
+
+ CometS3CredentialDispatcher.release(handles[0]);
+ CometS3CredentialDispatcher.release(handles[1]);
+
+ assertEquals(2, TestCometS3CredentialProvider.closeCount.get());
+ assertEquals(1, CometS3CredentialDispatcher.registeredCount());
+ }
Review Comment:
This moved to #7464 with the rest of the lifecycle change. It's fixed there
in 40963cfe6, and the reply is on the matching thread:
https://github.com/apache/datafusion-comet/pull/7464#discussion_r4475642516
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala:
##########
@@ -608,6 +609,50 @@ object CometIcebergNativeScan extends
CometOperatorSerde[CometBatchScanExec] wit
}
}
+ /** Activation key for the `CometS3CredentialProvider` SPI on the native
Iceberg path. */
+ private[comet] val IcebergCredentialProviderClassKey =
"s3.comet.credential.provider.class"
+
+ /** The provider that refreshes REST-vended credentials (Spark 4.x builds
only). */
+ private[comet] val VendedCredentialProviderClass =
+ "org.apache.comet.cloud.s3.IcebergVendedS3CredentialProvider"
+
+ // Iceberg's AwsClientProperties.REFRESH_CREDENTIALS_ENDPOINT and
REFRESH_CREDENTIALS_ENABLED.
+ private val RefreshCredentialsEndpointKey =
"client.refresh-credentials-endpoint"
+ private val RefreshCredentialsEnabledKey =
"client.refresh-credentials-enabled"
+
+ // The provider ships in Comet's own jar, so look it up in Comet's loader,
not the user's.
+ private lazy val vendedCredentialProviderAvailable: Boolean =
+ // scalastyle:off classforname
+ Try(Class.forName(VendedCredentialProviderClass, false,
getClass.getClassLoader)).isSuccess
Review Comment:
Thanks for tracing this. Addressed in two commits:
- **a6c57053c** (pushed after your review). Selection checks on the driver
that Comet's loader can load `AwsCredentialsProvider`, and that
`AwsClientProperties` from the context loader links against that same class. In
your layout, with the only SDK on the user-jar loader, the first probe fails,
so Comet doesn't select the provider and the scan keeps the vended keys. That
commit also took selection off the native write path.
`requireSupportedS3FileIOProperties` rejected the refresh endpoint before it
anyway, so the write falls back to Spark rather than failing closed.
- **047118f30.** The check is now
`vendedCredentialProviderUsable(cometLoader, contextLoader)`, and
`CometIcebergNativeScanSuite` covers three layouts separately. One shared SDK
is usable. A Comet loader that can't see the SDK while Iceberg can is not
usable; that's your case. An Iceberg loader with its own copy of iceberg-aws
and the SDK is not usable either; that's the duplicate-SDK case. A partly
visible SDK raises a `LinkageError`, which the check now treats as not usable
instead of failing planning.
The check runs on the driver, so it assumes executors see the same loader
layout, as `--packages` and `--jars` give.
##########
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() {
+ long[] handles = new long[3];
+ for (int i = 0; i < handles.length; i++) {
+ Map<String, String> props = new HashMap<>();
+ props.put("s3.session-token", "token-" + i);
+ handles[i] =
CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, DK, props);
+ }
+ assertEquals(3, CometS3CredentialDispatcher.registeredCount());
+
+ CometS3CredentialDispatcher.release(handles[0]);
+ CometS3CredentialDispatcher.release(handles[1]);
+
+ assertEquals(2, TestCometS3CredentialProvider.closeCount.get());
+ assertEquals(1, CometS3CredentialDispatcher.registeredCount());
+ }
+
+ @Test
+ public void concurrentAcquireAndReleaseNeverUseAClosedProvider() throws
Exception {
+ int threads = 8;
+ int rounds = 2000;
+ ExecutorService pool = Executors.newFixedThreadPool(threads);
+ try {
+ List<Future<?>> futures = new ArrayList<>();
+ for (int t = 0; t < threads; t++) {
+ futures.add(
+ pool.submit(
+ () -> {
+ for (int i = 0; i < rounds; i++) {
+ long handle = init();
+ CometS3CredentialDispatcher.getCredentialsForPath(handle,
"b", "k", READ);
Review Comment:
Good catch. The lifecycle change now lives in #7464. There,
`TestCometS3CredentialProvider` sets a per-instance closed flag in `close()`
and throws from `getCredentialsForPath` afterwards, so
`concurrentAcquireAndReleaseNeverUseAClosedProvider` fails on use after close.
Your mutation, closing right after `initialize()`, now fails on the first
lookup. Fixed in 40963cfe6.
--
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]
