FFI constructors silently discard arguments when the input is already foreign

[timsaucer]

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	task_ctx_provider
FFI_PhysicalExtensionCodec::new (proto/physical_extension_codec.rs:283)	ForeignPhysicalExtensionCodec	task_ctx_provider
FFI_TableProvider::new_with_ffi_codec (table_provider.rs:567)	ForeignTableProvider	logical_codec

The runtime argument is dropped on the same paths, but it is out of scope for this issue. Unlike the fields above, runtime lives in private_data, which belongs to the library that owns the handle, so the importing side cannot write it without an ABI change. FFI_SessionRef::new_with_ffi_codecs already takes that position explicitly ("retaining its original private data and runtime"). It should be documented on these three constructors rather than adopted.

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();   // task_ctx_provider is 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.

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.

The tests below run inside the datafusion-ffi crate (they touch pub(crate) state) against main at 1038d35. Each asserts the expected behaviour, so each fails today and passes once the constructors adopt their arguments.

One thing worth stating up front, because it is easy to write a repro that silently proves nothing: impl From<&FFI_LogicalExtensionCodec> for Arc<dyn LogicalExtensionCodec> (proto/logical_extension_codec.rs:348) compares library_marker_id first, and on a match returns Arc::clone(provider.inner()) — the original local Arc, not a ForeignLogicalExtensionCodec. Within one library the buggy branch is therefore never reached and the rebind appears to work. A repro must force the foreign path; the crate's crate::mock_foreign_marker_id exists for exactly this.

Shared helpers:

fn session() -> (Arc<SessionContext>, Arc<dyn TaskContextProvider>) {
    let ctx = Arc::new(SessionContext::new());
    let provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
    (ctx, provider)
}

/// Import an FFI handle as if it had crossed a library boundary, producing the
/// `Foreign*` wrapper rather than unwrapping back to the local `Arc`.
macro_rules! import_as_foreign {
    ($ffi:expr, $ty:ty) => {{
        let mut ffi = $ffi;
        ffi.library_marker_id = crate::mock_foreign_marker_id;
        let imported: Arc<$ty> = (&ffi).into();
        imported
    }};
}

1. FFI_LogicalExtensionCodec::new discards task_ctx_provider

#[test]
fn logical_codec_rebind_is_a_noop() {
    let (ctx_a, provider_a) = session();
    let (ctx_b, provider_b) = session();

    let ffi_a = FFI_LogicalExtensionCodec::new(
        Arc::new(DefaultLogicalExtensionCodec {}),
        None,
        &provider_a,
    );

    let imported = import_as_foreign!(ffi_a, dyn LogicalExtensionCodec);
    assert!(
        (Arc::clone(&imported) as Arc<dyn Any>)
            .downcast_ref::<ForeignLogicalExtensionCodec>()
            .is_some(),
        "expected the import to produce a ForeignLogicalExtensionCodec"
    );

    // Intent: rebind the imported codec to session B.
    let rebound = FFI_LogicalExtensionCodec::new(imported, None, &provider_b);

    let bound_to: Arc<TaskContext> = (&rebound.task_ctx_provider).try_into().unwrap();
    assert_eq!(
        bound_to.session_id(),
        ctx_b.task_ctx().session_id(),
        "rebound codec still resolves against session A ({}), not B",
        ctx_a.task_ctx().session_id()
    );
}

assertion `left == right` failed: rebound codec still resolves against session A (1cccd7f2-6086-4d9e-bcb1-db43d278ef75), not B
  left: "1cccd7f2-6086-4d9e-bcb1-db43d278ef75"
 right: "351dd7bb-0ae6-4adb-8fe8-c0b02f60b994"

2. The discarded provider is a Weak, so the handle also goes stale

#[test]
fn logical_codec_rebind_leaves_a_dangling_weak() {
    let (ctx_a, provider_a) = session();
    let (_ctx_b, provider_b) = session();

    let ffi_a = FFI_LogicalExtensionCodec::new(
        Arc::new(DefaultLogicalExtensionCodec {}),
        None,
        &provider_a,
    );
    let imported = import_as_foreign!(ffi_a, dyn LogicalExtensionCodec);
    let rebound = FFI_LogicalExtensionCodec::new(imported, None, &provider_b);

    // Session A is dropped; the caller believes the handle is bound to B.
    drop(provider_a);
    drop(ctx_a);

    match <Arc<TaskContext>>::try_from(&rebound.task_ctx_provider) {
        Ok(_) => {}
        Err(DataFusionError::Ffi(msg)) => {
            panic!("handle went stale with session A dropped: {msg}")
        }
        Err(e) => panic!("unexpected error: {e}"),
    }
}

handle went stale with session A dropped: TaskContextProvider went out of scope over FFI boundary.

3. FFI_PhysicalExtensionCodec::new behaves identically

#[test]
fn physical_codec_rebind_is_a_noop() {
    let (ctx_a, provider_a) = session();
    let (ctx_b, provider_b) = session();

    let ffi_a = FFI_PhysicalExtensionCodec::new(
        Arc::new(DefaultPhysicalExtensionCodec {}),
        None,
        &provider_a,
    );

    let imported = import_as_foreign!(ffi_a, dyn PhysicalExtensionCodec);
    assert!(
        (Arc::clone(&imported) as Arc<dyn Any>)
            .downcast_ref::<ForeignPhysicalExtensionCodec>()
            .is_some(),
        "expected the import to produce a ForeignPhysicalExtensionCodec"
    );

    let rebound = FFI_PhysicalExtensionCodec::new(imported, None, &provider_b);

    let bound_to: Arc<TaskContext> = (&rebound.task_ctx_provider).try_into().unwrap();
    assert_eq!(
        bound_to.session_id(),
        ctx_b.task_ctx().session_id(),
        "rebound codec still resolves against session A ({}), not B",
        ctx_a.task_ctx().session_id()
    );
}

assertion `left == right` failed: rebound codec still resolves against session A (2b9e4d4b-6b1f-44d5-85a3-8d078bdcd0dd), not B

4. FFI_TableProvider::new_with_ffi_codec discards logical_codec

#[test]
fn table_provider_rebind_discards_the_codec() {
    let (ctx_a, provider_a) = session();
    let (ctx_b, provider_b) = session();

    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
    let batch = RecordBatch::try_new(
        Arc::clone(&schema),
        vec![Arc::new(Float32Array::from(vec![2.0, 4.0]))],
    )
    .unwrap();
    let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap())
        as Arc<dyn TableProvider>;

    let ffi_a = FFI_TableProvider::new(table, true, None, &provider_a, None);
    let imported = import_as_foreign!(ffi_a, dyn TableProvider);
    assert!(
        imported.downcast_ref::<ForeignTableProvider>().is_some(),
        "expected the import to produce a ForeignTableProvider"
    );

    // Rebuild the codec against session B and re-wrap.
    let codec_b = FFI_LogicalExtensionCodec::new(
        Arc::new(DefaultLogicalExtensionCodec {}),
        None,
        &provider_b,
    );
    let rebound = FFI_TableProvider::new_with_ffi_codec(imported, true, None, codec_b);

    let bound_to: Arc<TaskContext> =
        (&rebound.logical_codec.task_ctx_provider).try_into().unwrap();
    assert_eq!(
        bound_to.session_id(),
        ctx_b.task_ctx().session_id(),
        "rebound provider's codec still resolves against session A ({}), not B",
        ctx_a.task_ctx().session_id()
    );
}

assertion `left == right` failed: rebound provider's codec still resolves against session A (6d4e6e51-f4de-45ab-816c-95a1b62c0e04), not B

Control: the sibling constructor does adopt

Added to query_planner.rs's own test module (FFI_QueryPlanner::logical_codec is private to that module). This one passes, confirming the two behaviours really do differ rather than the harness being wrong:

#[test]
fn test_rebind_foreign_query_planner_adopts_codecs() {
    let ctx_a = Arc::new(SessionContext::new());
    let ctx_b = Arc::new(SessionContext::new());
    let provider_b = Arc::clone(&ctx_b) as Arc<dyn TaskContextProvider>;

    let mut ffi_a = create_ffi_query_planner(Arc::clone(&ctx_a));
    ffi_a.library_marker_id = crate::mock_foreign_marker_id;
    let imported: Arc<dyn QueryPlanner + Send + Sync> = (&ffi_a).into();
    let any_ref: &dyn std::any::Any = imported.as_ref();
    assert!(any_ref.downcast_ref::<ForeignQueryPlanner>().is_some());

    let rebound = FFI_QueryPlanner::new_with_ffi_codecs(
        imported,
        FFI_LogicalExtensionCodec::new(
            Arc::new(DefaultLogicalExtensionCodec {}),
            None,
            &provider_b,
        ),
        FFI_PhysicalExtensionCodec::new(
            Arc::new(DefaultPhysicalExtensionCodec {}),
            None,
            &provider_b,
        ),
    );

    let bound_to: Arc<TaskContext> =
        (&rebound.logical_codec.task_ctx_provider).try_into().unwrap();
    assert_eq!(bound_to.session_id(), ctx_b.task_ctx().session_id());
}

FFI_SessionRef::new_with_ffi_codecs was not exercised by a test; its adopt-on-unwrap behaviour is visible directly at session/mod.rs:486-491.

Consequence in practice

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 —

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.

Test 2 above is the second failure mode: 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:

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;
}

runtime is not adopted, for the reason given above; it should instead be documented on each constructor so callers know it is only honored when a new wrapper is created.

If discarding is deliberate for any of the remaining arguments, then they 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: apache/datafusion-python#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.

Line numbers are against main at 1038d35.

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.