apache / apache/datafusion-comet

Bound `CometS3CredentialDispatcher` cache via refcounted handle lifecycle

Open
#4,456 0 comments 0 reactions 0 assignees View on GitHub
area:scan priority:medium
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 4h
Merged PRs (30d)
198

Description

## Problem

`CometS3CredentialDispatcher` caches one provider instance per `(FQCN, dispatchKey, catalogProperties)` triple in `KEY_TO_HANDLE` and `INSTANCES` for the lifetime of the executor JVM. Today eviction only happens via the JVM shutdown hook.

Catalogs that vend per-table credentials inject those credentials into the table's FileIO property bag, which Comet captures into `catalogProperties` at plan time (see comment in `IcebergReflection.scala` `getFileIOProperties`). The Iceberg REST + STS AssumeRole pattern (e.g. [Apache Gravitino's `S3TokenCredential`](https://gravitino.apache.org/docs/0.9.0-incubating/security/credential-vending) with default 1h TTL) is the canonical example. Each `loadTable` returns a fresh session token, so each unique table → a new `InstanceKey` → a new cached provider instance. On long-running JVMs (Spark Connect, Thrift Server) the maps grow without bound.

Discussion thread: https://github.com/apache/datafusion-comet/pull/4309#discussion_r3290519239

## Why a plain LRU is not the fix

Native `CometS3CredentialBridge` instances hold the handle by value and are reused across scans. Time- or size-based eviction can invalidate a handle that a live bridge still references mid-job, surfacing as opaque credential failures from `object_store` or `opendal`.

## Proposed approach

Refcounted handle lifecycle driven by the native side:

- Java: add a JNI-callable `releaseHandle(long)` that decrements a refcount and, at zero, removes the entry from both maps and calls `provider.close()` (swallowing exceptions, matching the shutdown-hook precedent).
- Native: wrap the handle in a struct whose `Drop` calls `releaseHandle`. The bridge itself is already shared via `Arc`, so the wrapping struct's `Drop` fires when the last `Arc` goes away.
- Guard against re-entry during JVM shutdown (the existing shutdown hook already drains everything) and during panic unwind.

The `ensureInitialized` path stays as-is: `computeIfAbsent` returns the existing handle and increments the refcount; first-time insertion sets count to 1.

## Acceptance criteria

- `KEY_TO_HANDLE.size()` does not grow unboundedly under per-table-vended-credential workloads.
- Steady-state path (same triple reused across many scans) keeps a single cached instance.
- Concurrent `ensureInitialized` + `releaseHandle` against the same key does not double-free or resurrect a released entry.
- Provider `close()` exceptions on release are swallowed and logged.
- Existing `CometS3CredentialBridgeSuite` and dispatcher unit tests continue to pass.

## Out of scope

- Driver→executor refresh of `catalogProperties` for long-running scans. That side of the AssumeRole story is addressed by composing with Spark's `HadoopDelegationTokenProvider`, documented in `docs/source/user-guide/latest/s3-credential-providers.md` and `docs/source/contributor-guide/s3-credential-provider-design.md`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.