Bridging between Tokio and Python

In the last few months, I have been working on new project. It is a rust library, that provides both a Rust, and Python library. I chose to build the Rust library with tokio I’m using Postgres to store state, and my preferred database library for Rust is sqlx. I also like Tokio because I wanted a smooth integration with the surrounding Tokio ecosystem for building networked applications, like Axum and Tonic.

I also wanted to use this library in Python applications. I frequently work with both Python and Rust, and I wanted an opportunity to learn how to use pyo3. The docs for pyo3 are solid and it explains how to do most of what I needed to do which was integrate tokio and pyo3. The other tool that I needed to use was maturin. Like pyo3, the documentation is excellent, and I was impressed at how thorough the project layout docs were as they took the time to cover both a pure rust package, and a mixed python/rust package.

The application and worker

My library contains an Application struct. The application takes configuration data, and acts as a factory for activities. The most important of these activities is the worker that consumes scheduled tasks and executes them. The Python library uses a Storage object from the core library that contains the core operations for creating, and running tasks. What I initially wanted to do was:

Show Plain Text
  1. // In the Python binding
  2. use pyo3::{exceptions::PyValueError, prelude::*};
  3. use crate::config::Config;
  4. use coreapp::storage::Storage;
  5.  
  6. #[pyclass]
  7. struct Application {
  8.     config: Config;
  9.     storage: Arc<Storage>;
  10. }
  11.  
  12. #[pymethods]
  13. impl Application {
  14.     #[new]
  15.     fn py_new(config: Config) -> Self {
  16.         let storage = Storage::new(config.clone().into());
  17.  
  18.         Application {
  19.             config: config,
  20.             storage: Arc::new(storage),
  21.         }
  22.     }
  23.  
  24.     fn create_worker(&self, worker_id: String) -> Worker {
  25.         Worker::new(
  26.             config: self.config.clone(),
  27.             storage: self.storage.clone(),
  28.             worker_id,
  29.         )
  30.     }
  31. }
  32.  
  33. #[pyclass]
  34. struct Worker {
  35.     config: Config;
  36.     storage: Arc<Storage>;
  37.     worker_id: String;
  38. }
  39.  
  40. #[pymethods]
  41. impl Worker {
  42.     #[new]
  43.     fn py_new(config: Config, storage: Arc<Storage>, worker_id: String) -> Self {
  44.         Worker {
  45.             config,
  46.             storage,
  47.             worker_id,
  48.         }
  49.     }
  50.  
  51.     /// Mark a run as complete.
  52.     fn complete_run(&self, run_id: String, run_result: Vec<u8>) -> PyResult<()> {
  53.         self.storage
  54.             .complete_run(run_id, &run_result)
  55.             .map_err(|e| PyValueError::new_err(format!("Could not complete_run: {e:?}")))
  56.     }
  57.  
  58.     /// Other methods
  59. }

The first problem I encountered was with the Worker::new method. pyo3 enforces a few rules on values that pass between rust and Python. The most relevant to my scenarios was that rust structs need to implement IntoPyObject, and I can’t do that for Tokio because of the orphan rule. Instead, I needed to remove the Python bound constructor and directly initialize the struct:

Show Plain Text
  1. #[pymethods]
  2. impl Application {
  3.     fn create_worker(&self, worker_id: String) -> Worker {
  4.         Worker {
  5.             config: self.config.clone(),
  6.             storage: self.storage.clone(),
  7.             worker_id,
  8.         }
  9.     }
  10. }

This approach let me use Rust’s smart pointers (like Arc) to share state with other application components. The drawback of this is that the Worker can’t ever be constructed in Python, it requires Rust structs, which can only be made from within Rust code. At first I was concerned that this would be a problem, but in practice it hasn’t mattered. In fact it has helped constrain the design and encourage simplicity in the boundaries between Python and Rust.

Using tokio in sync Python

Having an application class that acts like a factory for other components has been working out quite well as it has made passing a Tokio runtime around straightforward. For a period of time I had an intermediary struct that encapsulated the runtime + storage interaction, but it was so much boilerplate code. I found passing the runtime and storage around as a pair of parameters more ergonomic:

Show Plain Text
  1. #[pyclass]
  2. struct Application {
  3.     config: Config;
  4.     storage: Arc<Storage>;
  5.     runtime: Arc<tokio::runtime::Runtime>;
  6. }
  7.  
  8. #[pymethods]
  9. impl Application {
  10.     #[new]
  11.     fn py_new(config: Config) -> Self {
  12.         let runtime = tokio::runtime::Builder::new_current_thread()
  13.             .enable_all()
  14.             .build()
  15.             .unwrap();
  16.         let storage = runtime.block_on(async {
  17.             Storage::new(config.clone().into())
  18.         });
  19.  
  20.         Application {
  21.             config: config,
  22.             storage: Arc::new(storage),
  23.             runtime: Arc::new(runtime),
  24.         }
  25.     }

Then later when a worker needs to be created, the runtime and storage can be shared:

Show Plain Text
  1. #[pymethods]
  2. impl Application {
  3.     // other methods...
  4.  
  5.     fn create_worker(&self, worker_id: String) -> WorkerInner {
  6.         WorkerInner {
  7.             config: self.config.clone(),
  8.             storage: self.storage.clone(),
  9.             runtime: self.runtime.clone(),
  10.             worker_id,
  11.         }
  12.     }
  13.  
  14.     fn spawn_task(
  15.         &self,
  16.         task_name: &str,
  17.         params: &[u8],
  18.     ) -> PyResult<SpawnResult> {
  19.         let result = self.runtime.block_on(
  20.             self.storage.spawn_task(task_name, params)
  21.         );
  22.         result
  23.             .map(|v| v.into())
  24.             .map_err(|v| PyValueError::new_err(format!("Could not spawn task: {v:?}")))
  25.     }
  26.  
  27. }

The spawn_task method shows how the runtime and storage get used from a synchronous Python interface. The core library doesn’t return an error type that is compatible with Python, but that is solvable with map_err, and the Ok value is translated from the core API, to a Python value object using the Into trait. I am happy that I ended up with an implementation that isn’t too complicated given that I’m building a Python extension in Rust. I wanted to share this solution as I didn’t easily find another example when I was researching how to solve my initial problem.

I wanted to share this, as when I was initially stuck, I didn’t find any results, and it took some time to solve. I likely could have solved this faster with an agent, but I wouldn’t have learned something new, and been able to write about it. I hope to have this library ready as a development preview soon.

Comments

There are no comments, be the first!

Have your say: