The Arc Subfield Pattern

While working on my Horizon/NX reimplementation, I came across a somewhat fun pattern. Horizon/NX has a structure that looks like this:

struct KPort {
    KAutoObject refcount;
    KServerPort serverside;
    KClientPort clientside;
};

This structure is allocated via a SlabHeap. Both the Client and Server are handed separately, yet they all refer to the same (global) refcount. Once that refcount drops to 0, the whole KPort (including both substructures) will be deallocated.

Horizon/NX is written in C, so all operations are manual anyways. But porting this pattern to Rust, while making it as Rust-y as possible (read: using RAII), is non-trivial. We basically want the ability to map an Arc<T> to an Arc<T.field>. Hmm...

The boilerplate

Let's start by writing some boilerplate:

// Their content don't really matter...
pub struct ServerPort(());
pub struct ClientPort(());

pub struct Port {
    server: ServerPort,
    client: ClientPort
}

impl Port {
    pub fn new() -> Arc<Port> {
        Arc::new(Port {
            server: ServerPort,
            client: ClientPort
        })
    }
}

Nothing crazy so far. Here's an important observation: it should be impossible to construct ServerPort/ClientPort for the construct I'm about to show to be safe!

Now, we want to create a wrapper around ClientPort that represents the shared ownership with a Port. (you'd want to do this for ServerPort also).

struct ClientPortArc(*const ClientPort);

impl Port {
    fn client(this: Arc<Self>) -> ClientPortArc {
        unsafe {
            // Safety: no actual dereferencing is happening. We are merely
            // getting the address of one of the field.
            ClientPortArc(&(*Arc::into_raw(s)).client)
        }
    }
}

impl Deref for ClientPortArc {
    type Target = ClientPort;
    fn deref(&self) -> &ClientPort {
        unsafe {
            // Safety: We guarantee that the raw pointer is valid at
            // construction. Furthermore, since we leak an Arc with into_raw,
            // we know the pointer shouldn't get freed from under our feets.
            &*self.0
        }
    }
}
            

We now have our wrapper type. Now, we need to drop it. This is where the magic comes in. We have a pointer to a subfield. We know that subfield comes from a wrapper struct (Port). We now want to get a pointer back to the wrapper struct, hopefully in a generic way (let's not hardcode offsets, please).

Introducing: container_of

container_of is a magical macro. I first met it in the Linux Kernel, in their implementation of intrusive linked lists. It takes three arguments: a wrapper type name, the name of a member in that type, and a pointer to that member. Through pointer arithmetic, it will deduce a pointer to the wrapper type.

In other words, we can do this:

let client: ClientPortArc = panic!("whatever");
let port = container_of!(client.0, Port, client);

And get back a pointer to the Port that contained the ClientPortArc. I won't bore you with the details of how the container_of macro actually works (it's not complicated, but kinda messy, and has some interesting implications with UB. Which is especially complicated in Rust given what's UB and what isn't is still not defined). But here's a link if you want all the details. As for the rust version, it can be found in the intrusive-collections crate.

Goint back to our ClientPort. We can now use this to recover the Arc we leaked, and let it drop the refcount!

impl Drop for ClientPortArc {
    fn drop(&mut self) {
        unsafe {
            // Safety: We are guaranteed that ClientPortArc was created from an
            // Arc, and is built from within the wrapper.
            Arc::from_raw(container_of!(self.0, Port, client))
        }
    }
}

This should all work, but it's a lot of boilerplate to write each time we want to create such a wrapper. Can we do better?

Macros. They're good for you.

Each time I find such a pattern, I try to wrap it in a macro, and thoroughly document it. Now, there is one very important thing about macros: they can be extremely opaque to readers. As such, it's important for the macro designer to make it look as close as possible to normal rust code, so we can match people's expectations.

Let's start with how we want our macro to look:

arc_wrapper! {
    /// A wrapper around a ClientPort, providing shared ownership with the
    /// parent Port.
    pub struct ClientPortArc(ClientPort) wrapping Port, client {
        /// Get an owned reference to the Client part of this Port.
        fn client();
    }
}

It's kinda messy, but I couldn't think up of a better way. The function defined inside the {} will be added to Port, and will have the same prototype as our client function from earlier.

Now, let's try to implement the macro.

macro_rules! arc_wrapper {
    ($(#[$meta:meta])* $vis:vis struct $ty:ident($innerty:path) wrapping $container:path, $field:ident {
        $(#[$methodmeta:meta])*
        fn $methodname();
    }) => {

Nothing crazy so far. Something interesting to keep in mind: doc-comments in the form of /// will get turned into #[doc = "comment"] before being passed to macros. Hence why we take a list of meta here. Also, the $vis macro type will only be stable from 1.30 onward.

        $(#[$meta])*
        $vis struct $ty(*const $innerty);
        
        impl $container {
            $(#[$methodmeta])*
            $vis fn $field(s: Arc<Self>) -> $ty {
                unsafe { $ty(&(*Arc::into_raw(s)).$field) }
            }
        }

We're defining our new type, and creating a getter on the container struct to return our wrapper.

        impl Deref for $ty {
            type Target = $innerty;
            
            fn deref(&self) -> &$innerty {
                unsafe {
                    &*self.0
                }
            }
        }
        
        impl Drop for $ty {
            fn drop(&mut self) {
                unsafe {
                    Arc::from_raw(container_of!(self.0, $container, $field));
                }
            }
        }
    }
}