Didn't want to pass around the client because it would need to go around multiple tasks and threads. So I kept a static client using lazy static. However I can't use await in the initialization block.
What can I do to work around this?
Suggestions for doing it completely differently without lazy_static is also welcome.
use std::env;
use futures::stream::StreamExt;
use mongodb::{
bson::{doc, Bson},
options::ClientOptions,
Client,
};
lazy_static! {
static ref MONGO: Option<Client> = {
if let Ok(token) = env::var("MONGO_AUTH") {
if let Ok(client_options) = ClientOptions::parse(&token).await {
^^^^^
if let Ok(client) = Client::with_options(client_options) {
return Some(client);
}
}
}
return None;
};
}
I think the key problem is that lazy_static always uses a synchronous primitive for initialization, so its always going to block on the initialization.
What about building our own async initialization locking instead? Here's a snippet using once_cell to provide the mutable statics:
use once_cell::sync::OnceCell;
use tokio::sync::Mutex;
static MONGO: OnceCell<Client> = OnceCell::new();
static MONGO_INITIALIZED: OnceCell<Mutex<bool>> = OnceCell::new();
pub async fn get_mongo() -> &'static Client {
// this is racy, but that's OK: it's just a fast case
if let Some(v) = MONGO.get() {
return v;
}
// it hasn't been initialized yet, so let's grab the lock & try to
// initialize it
let initializing_mutex = MONGO_INITIALIZED.get_or_init(|| Mutex::new(false));
// this will wait if another task is currently initializing the client
let mut initialized = initializing_mutex.lock().await;
// if initialized is true, then someone else initialized it while we waited,
// and we can just skip this part.
if !*initialized {
// no one else has initialized it yet, so
let client = Client /* async code to initialize client here! */ ;
MONGO.set(client).expect(
"no one else should be initializing this \
as we hold MONGO_INITIALIZED lock",
);
*initialized = true;
drop(initialized);
}
MONGO.get().unwrap()
}
I'll let you look through that, but the gist of it is that we can use a tokio-based Mutex to ensure that only one task initializes the mongo database at a time, and besides that just use OnceCell for actually storing the initialized value.
OnceCell is very similar to lazy_static!, the main difference is that it requires a slightly newer rustc version, isn't macro-based, and forces whoever is accessing it to provide the initialization code, rather than building it into the definition.
We use two OnceCells since we have one for the actual data, and a separate one storing only a lock for initialization. I've written it this way so that in the fast case when it is already initialized, we don't have to lock the mutex or use it at all to access the data.
This code could also probably be improved - I don't think we need the boolean inside MONGO_INITIALIZED for instance, since it's redundant with MONGO's state. But regardless it should work.
I understand you approach fairly well actually. I tried something very similar. Instead of using a bool I tried to approach it by wrapping the client in an Option.
OnceCell<Mutex<Option>>
So it would be a None at first, and After initializng it would be a Some(client). It didn't work out, probably because I'm a total amature
would this work?
I'll try out your approach. If it works I'll try using an Option again.
I think it should! It might be a bit inconvenient, though, since you'll end up having to return a MutexGuard<'static, Client> every time you want to use the client.
The other thing is that since the client is inside a Mutex, now, you'll only be able to access it from one task at a time. If that's what you want, then maybe a Mutex would be OK or even preferable?
But if not, then keeping the client in its own OnceCell outside of the mutex allows us to access it freely from multiple tasks at the same time, and without any expensive locking. This is the main reason I separated them out rather than putting the client under an option in a single OnceCell.
It worked!
I understand your concern about having to return an entire mutexguard every time. I think I'll stick to this current approach.
This is what I'm finally using. I'm returning the result of .get(), an Option<&Client> directly without unwrapping it since I can deal with the client being uninitialized, a None, where I'm using it. So other than ditching the expects and unraps it's your implementation.
Thanks a lot!
static MONGO: OnceCell<Client> = OnceCell::new();
static MONGO_INITIALIZED: OnceCell<tokio::sync::Mutex<bool>> = OnceCell::new();
pub async fn get_mongo() -> Option<&'static Client> {
// this is racy, but that's OK: it's just a fast case
let client_option = MONGO.get();
if let Some(_) = client_option {
return client_option;
}
// it hasn't been initialized yet, so let's grab the lock & try to
// initialize it
let initializing_mutex = MONGO_INITIALIZED.get_or_init(|| tokio::sync::Mutex::new(false));
// this will wait if another task is currently initializing the client
let mut initialized = initializing_mutex.lock().await;
// if initialized is true, then someone else initialized it while we waited,
// and we can just skip this part.
if !*initialized {
// no one else has initialized it yet, so
if let Ok(token) = env::var("MONGO_AUTH") {
if let Ok(client_options) = ClientOptions::parse(&token).await {
if let Ok(client) = Client::with_options(client_options) {
if let Ok(_) = MONGO.set(client) {
*initialized = true;
}
}
}
}
}
drop(initialized);
MONGO.get()
}
The way I’ve done something similar in the past is using a OnceCell. Initializing it at the beginning of main so I know a value has been set and then in a helper function always clone the client using CELL.get().unwrap().clone() so that I have an owned client which is easier to use in async code.
Yea. That and instead of worrying about initialization in the helper function, just initializing at program start. I mostly just wanted to point out the clone every time since I find that a lot easier to use in async. Obviously you could just clone at the call site so maybe it doesn’t matter.
