timsaucer opened a new issue, #24722:
URL: https://github.com/apache/datafusion/issues/24722
## Describe the bug
Several `datafusion-ffi` constructors unwrap an already-foreign input and
return its original handle. Three of them drop the arguments passed alongside,
without an error or a warning:
| Constructor | Unwraps | Silently discards |
|---|---|---|
| `FFI_LogicalExtensionCodec::new` (`proto/logical_extension_codec.rs:297`)
| `ForeignLogicalExtensionCodec` | `runtime`, `task_ctx_provider` |
| `FFI_PhysicalExtensionCodec::new`
(`proto/physical_extension_codec.rs:283`) | `ForeignPhysicalExtensionCodec` |
`runtime`, `task_ctx_provider` |
| `FFI_TableProvider::new_with_ffi_codec` (`table_provider.rs:405`) |
`ForeignTableProvider` | `runtime`, `logical_codec` |
```rust
pub fn new(
codec: Arc<dyn LogicalExtensionCodec>,
runtime: Option<Handle>,
task_ctx_provider: impl Into<FFI_TaskContextProvider>,
) -> Self {
if let Some(codec) = (Arc::clone(&codec) as Arc<dyn Any>)
.downcast_ref::<ForeignLogicalExtensionCodec>()
{
return codec.0.clone(); // runtime and task_ctx_provider are never
read
}
...
```
Two sibling constructors hit the same case and do the opposite — they adopt
the supplied values:
`FFI_QueryPlanner::new_with_ffi_codecs` (`query_planner.rs:252`), whose doc
comment makes the guarantee explicit:
> If `planner` is already foreign, this re-exports its original FFI handle
rather than adding another wrapper layer. The handle still adopts the codecs
supplied here, so they are never silently discarded.
```rust
if let Some(planner) = any_ref.downcast_ref::<ForeignQueryPlanner>() {
let mut planner = planner.0.clone();
planner.logical_codec = logical_codec;
planner.physical_codec = physical_codec;
return planner;
}
```
`FFI_SessionRef::new_with_ffi_codecs` (`session/mod.rs:481`) does the same
for `logical_codec` and `physical_codec`.
Given those two, the other three look like oversights rather than intent.
## To Reproduce
A consumer that imports a foreign codec can never rebind it afterwards.
Re-wrapping with a different `task_ctx_provider` compiles, runs, and has no
effect:
```rust
// P1 and P2 are distinct sessions.
let ffi = FFI_LogicalExtensionCodec::new(my_codec, None, &p1_provider);
let imported: Arc<dyn LogicalExtensionCodec> = (&ffi).into(); //
ForeignLogicalExtensionCodec
// Intent: rebind the imported codec to P2. Actual: no-op, still resolves P1.
let rebound = FFI_LogicalExtensionCodec::new(imported, None, &p2_provider);
```
(Sketch — I read this from the source rather than running it. Happy to add a
failing unit test to the PR if that is useful.)
The observable consequence, from `datafusion-python`:
1. A user installs a foreign `LogicalExtensionCodec`. Its FFI handle holds a
`Weak` to session **A**.
2. The user then installs a foreign `QueryPlanner`. `datafusion-python`
forks the session to **B** so the receiver is not mutated, and rebuilds its own
outer codec wrapper against **B**.
3. The inner foreign codec cannot be rebound, so it still points at **A**.
4. Decode callbacks in the extension library resolve names against **A**'s
registry. A UDF registered after the fork is invisible to them, and the config
they see is the pre-fork snapshot.
Worth noting where the correct context is lost: the host *does* pass the
right `TaskContext` down —
```rust
fn try_decode_table_provider(&self, buf, table_ref, schema, ctx:
&TaskContext) -> Result<...> {
self.inner.try_decode_table_provider(buf, table_ref, schema, ctx) //
ctx belongs to B
}
```
— but `try_decode_table_provider_fn_wrapper` takes no context parameter and
calls `codec.task_ctx()` instead, so the argument is dropped at the boundary
and the stale stored provider is substituted.
There is a second failure mode with the same root cause. Because the
discarded provider is held as a `Weak`, a consumer that cannot rebind must
instead keep the original session alive artificially, or the capsule starts
failing with `TaskContextProvider went out of scope over FFI boundary`.
## Expected behavior
The three constructors adopt the supplied values on the unwrap path,
matching `FFI_QueryPlanner::new_with_ffi_codecs`:
```rust
if let Some(codec) = (Arc::clone(&codec) as Arc<dyn Any>)
.downcast_ref::<ForeignLogicalExtensionCodec>()
{
let mut codec = codec.0.clone();
codec.task_ctx_provider = task_ctx_provider.into();
return codec;
}
```
If discarding is deliberate for any of them, then the arguments should not
be accepted silently — either document the behaviour on the constructor, or
change the signature so a caller cannot pass a value that will be ignored.
## Additional context
Found while adding FFI query planner support to `datafusion-python`:
https://github.com/apache/datafusion-python/pull/1677
The workaround there is to retain the pre-fork `SessionContext` so the
`Weak` stays valid. That prevents the crash but not the staleness, and it
retains memory that could otherwise be released. With the change above, the
fork can rebind the inner codecs and the workaround is deleted.
Code references are against the 55.x line (`branch-55`); line numbers should
be confirmed against `main`.
Separately, the `runtime` argument is discarded on the same paths. That is
the same bug shape, and the same fix applies, though it has been less visible
in practice.
Related design question, if it is worth a separate issue: `try_decode` and
`try_decode_table_provider` pull a `TaskContext` from the codec's stored
provider even though the calling side usually has one in hand and already
passes it to the trait method. Threading it through the FFI signature would
make the stored provider unnecessary for those paths and remove the rebinding
problem entirely, rather than making rebinding possible.
--
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]
