sunchao commented on code in PR #7363:
URL: https://github.com/apache/datafusion-comet/pull/7363#discussion_r4458627527
##########
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:
This checks that the adapter class can be found without initializing it,
which does not establish that its defining loader can resolve the AWS SDK. With
Comet on the parent/application classpath and Iceberg 1.11 plus the sole AWS
SDK copy on Spark's user-jar loader (`--packages`), this check succeeds.
Iceberg resolves the supplied fresh vended session credentials in that layout,
but initializing this head's Comet adapter throws `NoClassDefFoundError:
software/amazon/awssdk/auth/credentials/AwsCredentialsProvider`. The repository
documents this parent/child loader arrangement.
This becomes a regression through the new automatic selection:
`CometIcebergNativeWrite` now injects the provider even when the user
configured none. The prior native path used the still-valid static vended keys;
provider initialization errors make the native WRITE path fail closed in
`iceberg_common.rs` (reads instead warn and fall back). Could selection verify
that the adapter can actually resolve its SDK, or keep all Iceberg/SDK
interaction in the user loader? Please cover this one-SDK/default-delegation
layout separately from the duplicate-SDK case. I reproduced provider
initialization and credential resolution locally and traced the write failure
path; no live S3 write was attempted.
##########
spark/src/main/java/org/apache/comet/cloud/s3/CometS3CredentialDispatcher.java:
##########
@@ -198,20 +231,30 @@ private static CometS3CredentialProvider
instantiate(String providerClassName) {
/** Visible for tests; otherwise invoked from the JVM shutdown hook. */
static void closeAll() {
for (RegisteredProvider registered : INSTANCES.values()) {
- try {
- registered.provider.close();
- } catch (Throwable t) {
- LOG.warn(
- "Exception from {} (dispatchKey={}).close() during shutdown",
- registered.key.providerClassName,
- registered.key.dispatchKey,
- t);
- }
+ close(registered, "during shutdown");
Review Comment:
There is a shutdown-only race with the new `release()` path: `closeAll()`
can start closing an entry while its last native reference concurrently removes
and closes that same entry. With a test provider whose first close waits on a
latch, releasing the final handle during that wait produced two concurrent
close calls. Could both paths claim an entry for closing through the same
atomic operation? The shipped adapter forwards close to its delegate without a
guard. Normal rotation passed the stress probe; I have not observed a failure
in Iceberg's delegate from this shutdown overlap.
##########
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:
Could the fake provider reject `getCredentialsForPath` after its own
`close()`? This test's fake continues returning credentials after close, while
the final counters only count closures. In a scratch mutation I closed every
provider immediately after `initialize()` and removed the last-release close;
this test still passed all 16000 lookups. A per-instance closed flag, set by
`close()` and checked by credential lookup, would make the test detect the
use-after-close behavior its name promises to guard.
--
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]
