axum trait based factory

[washanhanzi]

Plan: Adopt Axum's Trait-Based Factory Pattern

Problem

Current approach uses closures for factories:

Box<dyn Fn(&S) -> BoxedCall<Req, Resp> + Send + Sync>

Issue: Clone is not object-safe, so we can't clone the factory without Arc or other workarounds.

Goal

Adopt Axum's pattern using a custom trait with clone_box() method to enable efficient cloning of type-erased factories.

---

Current vs Proposed

Current (Closure)

// Factory type
type Factory<S, Req, Resp> = Box<dyn Fn(&S) -> BoxedCall<Req, Resp> + Send + Sync>;

// In builder
pub struct HelloWorldServiceTonicCompatibleBuilder<S = ()> {
    pub say_hello: Option<Box<dyn Fn(&S) -> BoxedCall<...> + Send + Sync>>,
    // Cannot clone!
}

Proposed (Trait)

// Trait for type erasure with cloning
pub trait ErasedIntoBoxedCall<S, Req, Resp>: Send + Sync {
    fn clone_box(&self) -> Box<dyn ErasedIntoBoxedCall<S, Req, Resp>>;
    fn into_boxed_call(&self, state: &S) -> BoxedCall<Req, Resp>;
}

// Wrapper type
pub struct BoxedIntoCall<S, Req, Resp>(Box<dyn ErasedIntoBoxedCall<S, Req, Resp>>);

// In builder
pub struct HelloWorldServiceTonicCompatibleBuilder<S = ()> {
    pub say_hello: Option<BoxedIntoCall<S, HelloRequest, HelloResponse>>,
    // Can clone via clone_box()!
}

---

Implementation Plan

Step 1: Define Core Traits

// connectrpc-axum/src/handler/tonic.rs

/// Trait for type-erased factory with cloning support.
/// Mirrors Axum's `ErasedIntoRoute` pattern.
pub trait ErasedIntoBoxedCall<S, Req, Resp>: Send + Sync {
    /// Clone this factory into a new boxed trait object.
    fn clone_box(&self) -> Box<dyn ErasedIntoBoxedCall<S, Req, Resp>>;

    /// Create a BoxedCall by binding state.
    fn into_boxed_call(&self, state: &S) -> BoxedCall<Req, Resp>;
}

/// Trait for streaming variant.
pub trait ErasedIntoBoxedStreamCall<S, Req, Resp>: Send + Sync {
    fn clone_box(&self) -> Box<dyn ErasedIntoBoxedStreamCall<S, Req, Resp>>;
    fn into_boxed_stream_call(&self, state: &S) -> BoxedStreamCall<Req, Resp>;
}

Step 2: Define Concrete Factory Struct

/// Concrete factory that stores handler and conversion function.
/// Similar to Axum's `MakeErasedHandler`.
pub struct MakeBoxedCall<F, S, Req, Resp> {
    handler: F,
    make_call: fn(F, &S) -> BoxedCall<Req, Resp>,
    _marker: PhantomData<(S, Req, Resp)>,
}

impl<F, S, Req, Resp> MakeBoxedCall<F, S, Req, Resp>
where
    F: Clone,
{
    pub fn new(handler: F, make_call: fn(F, &S) -> BoxedCall<Req, Resp>) -> Self {
        Self {
            handler,
            make_call,
            _marker: PhantomData,
        }
    }
}

impl<F, S, Req, Resp> Clone for MakeBoxedCall<F, S, Req, Resp>
where
    F: Clone,
{
    fn clone(&self) -> Self {
        Self {
            handler: self.handler.clone(),
            make_call: self.make_call,
            _marker: PhantomData,
        }
    }
}

Step 3: Implement Trait for Concrete Struct

impl<F, S, Req, Resp> ErasedIntoBoxedCall<S, Req, Resp> for MakeBoxedCall<F, S, Req, Resp>
where
    F: Clone + Send + Sync + 'static,
    S: 'static,
    Req: 'static,
    Resp: 'static,
{
    fn clone_box(&self) -> Box<dyn ErasedIntoBoxedCall<S, Req, Resp>> {
        Box::new(self.clone())
    }

    fn into_boxed_call(&self, state: &S) -> BoxedCall<Req, Resp> {
        (self.make_call)(self.handler.clone(), state)
    }
}

Step 4: Define Wrapper Type

/// Type-erased factory wrapper with Clone support.
/// Similar to Axum's `BoxedIntoRoute`.
pub struct BoxedIntoCall<S, Req, Resp>(Box<dyn ErasedIntoBoxedCall<S, Req, Resp>>);

impl<S, Req, Resp> BoxedIntoCall<S, Req, Resp> {
    /// Create from a handler using the IntoFactory trait.
    pub fn from_handler<F, T>(handler: F) -> Self
    where
        F: Clone + Send + Sync + 'static,
        TonicCompatibleHandlerWrapper<F>: IntoFactory<T, Req, Resp, S>,
        S: Clone + Send + Sync + 'static,
        Req: 'static,
        Resp: 'static,
    {
        let make_call = |f: F, state: &S| {
            let wrapper = TonicCompatibleHandlerWrapper(f);
            // Use existing IntoFactory logic
            let factory = wrapper.into_factory();
            factory(state)
        };

        Self(Box::new(MakeBoxedCall::new(handler, make_call)))
    }

    /// Bind state to produce a BoxedCall.
    pub fn into_boxed_call(self, state: &S) -> BoxedCall<Req, Resp> {
        self.0.into_boxed_call(state)
    }

    /// Bind state without consuming (for repeated use).
    pub fn make_boxed_call(&self, state: &S) -> BoxedCall<Req, Resp> {
        self.0.into_boxed_call(state)
    }
}

impl<S, Req, Resp> Clone for BoxedIntoCall<S, Req, Resp> {
    fn clone(&self) -> Self {
        Self(self.0.clone_box())
    }
}

Step 5: Update Generated Code

Before:

pub struct HelloWorldServiceTonicCompatibleBuilder<S = ()> {
    pub say_hello: Option<
        Box<dyn Fn(&S) -> BoxedCall<HelloRequest, HelloResponse> + Send + Sync>,
    >,
}

After:

pub struct HelloWorldServiceTonicCompatibleBuilder<S = ()> {
    pub say_hello: Option<BoxedIntoCall<S, HelloRequest, HelloResponse>>,
}

impl<S> Clone for HelloWorldServiceTonicCompatibleBuilder<S>
where
    S: Clone,
{
    fn clone(&self) -> Self {
        Self {
            say_hello: self.say_hello.clone(),  // Now works!
            router: self.router.clone(),
        }
    }
}

Step 6: Update Builder Methods

Before:

pub fn say_hello<F, T>(mut self, handler: F) -> Self
where
    TonicCompatibleHandlerWrapper<F>: IntoFactory<T, ...>,
{
    let wrapper = TonicCompatibleHandlerWrapper(handler);
    let factory = wrapper.into_factory();
    self.say_hello = Some(factory);  // Box<dyn Fn...>
    self
}

After:

pub fn say_hello<F, T>(mut self, handler: F) -> Self
where
    F: Clone + Send + Sync + 'static,
    TonicCompatibleHandlerWrapper<F>: IntoFactory<T, ...>,
{
    self.say_hello = Some(BoxedIntoCall::from_handler(handler));
    self
}

---

Streaming Variant

Same pattern for streaming:

pub struct BoxedIntoStreamCall<S, Req, Resp>(
    Box<dyn ErasedIntoBoxedStreamCall<S, Req, Resp>>
);

pub struct MakeBoxedStreamCall<F, S, Req, Resp> {
    handler: F,
    make_call: fn(F, &S) -> BoxedStreamCall<Req, Resp>,
    _marker: PhantomData<(S, Req, Resp)>,
}

// ... similar implementations

---

File Changes Summary

File	Changes
connectrpc-axum/src/handler/tonic.rs	Add trait + struct + wrapper
connectrpc-build/src/prost/connect.rs	Update codegen to use BoxedIntoCall
connectrpc-build/src/prost/tonic_compat.rs	Update builder generation

---

Benefits

1. Cloneable Builders: HelloWorldServiceTonicCompatibleBuilder can derive Clone
2. Matches Axum Pattern: Familiar to Axum users
3. No Arc Overhead: Clone via clone_box() is efficient
4. Type Safety: Same compile-time guarantees

Drawbacks

1. More Code: Additional trait + struct + wrapper
2. Complexity: Harder to understand than simple closure
3. Generated Code Size: Slightly larger output

---

Alternative: Keep Closure, Add Arc

If cloning isn't a hard requirement, simpler alternative:

pub say_hello: Option<Arc<dyn Fn(&S) -> BoxedCall<...> + Send + Sync>>,

Pros: Minimal changes
Cons: Arc overhead, not idiomatic with Axum

---

Recommendation

Adopt the trait-based approach if:

• Users need to clone builders
• Consistency with Axum patterns is important
• Performance of cloning matters

Keep closure approach if:

• Cloning builders is rare/unnecessary
• Simplicity is preferred
• Current approach works for all use cases

---

Questions to Resolve

1. Is builder cloning actually needed? Check if users clone builders in practice.
2. Performance impact? Benchmark clone_box() vs Arc::clone().
3. Code generation complexity? Ensure prost-build changes are manageable.