Converting FunctionTrace from C to Rust

I maintain FunctionTrace, a graphical Python profiler that provides a clear view of your application’s execution with minimal performance impact.

I recently¹ converted FunctionTrace’s Python implementation from a C extension into a Rust extension backed by PyO3. While there are various resources for creating new Python extensions written in Rust, I found very little information on how to incrementally migrate an existing extension. This writeup details the somewhat sketchy but effective approach I took to do a gradual migration from C to Rust.

Why migrate?

FunctionTrace was initially written in 2020 and was composed of two parts: a small Python script that wraps a meaty C extension to trace Python execution, and a server that collects trace data from running Python processes and turns it into something useable².

At the time, writing a C extension was the only realistic approach to writing a high-performance Python extension. In the years since then, this has changed. PyO3 has become an increasingly popular option for writing Python extensions and is used by extensions including cryptography and pydantic. A couple years ago, I started keeping an eye on PyO3 and thinking about whether it would be worth it to migrate.

While I enjoy writing C³, some of its restrictions make it challenging to maintain a Python extension. Unfortunately, C doesn’t have a particularly enjoyable story around using third-party libraries, and its type system leaves something to be desired (particularly around dangerous operations). FunctionTrace supports multiple operating systems, meaning I either need to pull in libraries, write per-OS code, or stick to the lowest common denominator when implementing features. FunctionTrace also supports a range of Python versions, some of which require documentation to safely access from C and have different internal representations of data.

None of these issues are insurmountable, but as a single developer maintaining the project mostly as a hobby, I found them annoying. If I switched to PyO3, I expected to have significantly more guidance from the compiler around risky operations, as well as be able to pull in libraries to handle functionality that I’d simply left out of FunctionTrace previously.

Finally, I really enjoy writing Rust, and figured I’d simply be able to move faster if all of FunctionTrace was written in Rust, rather than switching back and forth between Rust for the server and C for the extension.

Starting the migration

Over the years, I had poked at PyO3 a couple times to see how challenging it would be to switch FunctionTrace to use it. However, I kept running into the problem that most of FunctionTrace’s code either runs at module initialization or is tightly tied to data structures setup then. Without any standalone functions I could migrate to Rust, it seemed like I’d be committed to rewriting upwards of a thousand lines of rather gnarly C code at once if I wanted to do the conversion⁴.

This time, I instead embraced the unsafe and came up with a way for C and Rust extensions to coexist!

The main insight is Python modules can be looked at as an implementation of dlopen. Modules are all loaded into the same process and therefore address space, so data can be directly shared between extensions without needing to go via the Python API. This allows us to rewrite chunks of C functionality into Rust and expose the Rust functionality via a function pointer to C without ever getting Python involved!

I added the initial implementation in 26170f07, which set the scaffolding for using this approach:

I created a new internal Python module via PyO3 named _functiontrace_rs. I moved a chunk of functionality around socket messaging to this module in the Mpack_Flush function, then exposed a function pointer to it via a public Python function message_flush.

#[unsafe(no_mangle)]
extern "C" fn Mpack_Flush(writer: *const MpackWriter, buffer: *const u8, bytes: usize) {
  // ... data loading functionality ...

  if let Err(e) = socket.write_all(data) {
    panic!("Socket send failed: {e}");
  }
}

#[pyfunction]
fn message_flush() -> PyResult<usize> {
  Ok(Mpack_Flush as usize)
}

In the C extension, I then import _functiontrace_rs and call message_flush() to get a function pointer to Mpack_Flush, allowing it to be used as if it was a normal C function.

PyObject* rust = PyImport_ImportModule("_functiontrace_rs");

PyObject* msg_flush = PyObject_GetAttrString(rust, "message_flush");
void* mpack_flush = PyLong_AsVoidPtr(PyObject_CallFunctionObjArgs(msg_flush, NULL));

mpack_writer_set_context(&initWriter, mpack_flush);

After this change, C code was still responsible for all the tracing and interesting work, but server communication was now handled in Rust! At this point, I published a new release on PyPi to see if the new dependency on Rust caused anyone issues, and resolved to wait a bit before doing further migration.

Migrating the rest of the owl

Most of the full migration was able to follow this same pattern — I’d find a piece of functionality I could move into Rust, move it, then expose a function pointer back to the C. Along the way, I ran into a few patterns that were useful:

Embrace unsafe

Working with external data like this is one of the things that unsafe was made for. Rust obviously can’t check any of the interactions with the C extension, so it’s fully up to you to ensure you’re doing things safely. Read the Rustonomicon and understand what you’re doing, and leave SAFETY comments where applicable. Otherwise, expect to blow your foot off⁵.

Sharing a more complex structure

Exposing a one function pointer at a time via a Python function works, but it gets tedious if you’re working with multiple functions or other pieces of data at a time. It’s easier to make a shared struct containing the information that both sides need.

// The exposed apis of the Rust extension.
struct {
  void* (*message_initialization)(char* socket_path);
  void (*message_shutdown)(void* context);
  void (*message_flush)(Writer* writer, const char* buffer, size_t bytes);
}* rust;

Then, the Rust code can allocate this structure via Box::leak(), ensuring it’s always safe for the C extension to access.

#[repr(C)]
struct RustFunctions {
    message_initialization: usize,
    message_shutdown: usize,
    message_flush: usize,
}

#[pyfunction]
fn c_api() -> usize {
  let api = Box::new(RustFunctions {
    message_initialization: message_initialization as usize,
    // ...
  });

  Box::into_raw(api) as usize
}

I added a version of this in 1d0ee59f, then heavily used it throughout the remainder of the migration.

Sharing parts of a structure

It’s tempting to think of structs as monolithic objects that contain a set of fields. At the end of the day though, it’s all just bytes. To “safely”⁶ interoperate with C, your Rust code often only needs to understand the prefixes of C structs. This means if you put your interesting fields at the start of a struct, you may be able to share it with Rust without caring about things like architecture-specific layouts or conditional compilation.

For example, in f7397709 I create an MpackWriter that exposes only a small prefix of the fields in an mpack_writer_t. However, these fields are enough for me to move most meaningful operations on mpack_writer_ts to Rust, paving the way for fully removing mpack_writer_t and supporting code from the C codebase.

Sharing can go both ways

We’ve discussed exposing Rust to the C extension, but the C extension can also expose itself to Rust. f7397709 implements this with a set_config function pointer that’s exposed to C, allowing the C extension to pass arbitrary pointers to the Rust code. This was particularly useful for sharing globals between the two modules.

Afterwards

The actual migration followed the techniques I mentioned, with the C extension steadily growing smaller over a series of 27 commits until it was finally eliminated in a beautiful -13808 line change. Because I slowly migrated each piece, it was easy to run tests after each change, detect if I’d broken something, and know exactly what to fix it if so. I only ran into a couple bugs, each of which involving me misunderstanding the safety invariants I needed to uphold for an unsafe function.

PyO3 feels more restrictive than writing C. This is partly because there aren’t safe wrappers around much of the functionality that I use, and partly because it doesn’t allow me to access structures that won’t work in some versions of Python. This seems like a fair trade.

Beyond that, using PyO3 is a substantial improvement over working directly against the Python C API. Error handling is cleaner, and the code in general is less verbose. Where there are safe abstractions, such as for exposing functions to Python or running with the GIL held, it feels downright easy to use. I’d recommend using it rather than C for Python extensions, and would love to see more Python libraries using it to speed up inefficient operations.

Now that FunctionTrace is entirely written in Rust, I’m excited for some of the future opportunities that I could have done in C but realistically wouldn’t have. I’ve already added logging to make it easier for people to report and debug issues they run into. I’d also like to switch to a shared-memory implementation for client-server communication, further reducing overhead and finally allowing FunctionTrace to run on Windows.

Please let me know if you run into issues using the new Rust implementation of FunctionTrace!