use mongodb::bson::doc;
use mongodb::{Client, Collection};
use tokio::sync::mpsc;
use crate::{Log, LogError, is_heartbeat};
const ASYNC_LOG_CHANNEL_CAPACITY: usize = 1024;
fn io_err<E: std::fmt::Display>(e: E) -> LogError {
LogError::Io(e.to_string())
}
fn now_unix() -> i64 {
time::OffsetDateTime::now_utc().unix_timestamp()
}
enum Entry {
Message {
direction: &'static str,
text: String,
},
Event {
text: String,
},
}
#[derive(Debug, Clone)]
pub struct MongoLogConfig {
pub uri: String,
pub database: String,
pub incoming_collection: String,
pub outgoing_collection: String,
pub event_collection: String,
pub include_heartbeats: bool,
pub session_id: String,
}
impl MongoLogConfig {
pub fn new(uri: impl Into<String>) -> Self {
Self {
uri: uri.into(),
database: "truefix".to_owned(),
incoming_collection: "log_incoming".to_owned(),
outgoing_collection: "log_outgoing".to_owned(),
event_collection: "log_event".to_owned(),
include_heartbeats: true,
session_id: "default".to_owned(),
}
}
}
pub struct MongoLog {
tx: std::sync::Mutex<Option<mpsc::Sender<Entry>>>,
task: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
include_heartbeats: bool,
}
impl MongoLog {
pub async fn connect(uri: &str) -> Result<Self, LogError> {
Self::connect_with_config(MongoLogConfig::new(uri)).await
}
pub async fn connect_with_config(config: MongoLogConfig) -> Result<Self, LogError> {
let client = Client::with_uri_str(&config.uri).await.map_err(io_err)?;
let db = client.database(&config.database);
let incoming: Collection<mongodb::bson::Document> =
db.collection(&config.incoming_collection);
let outgoing: Collection<mongodb::bson::Document> =
db.collection(&config.outgoing_collection);
let event: Collection<mongodb::bson::Document> = db.collection(&config.event_collection);
let (tx, mut rx) = mpsc::channel::<Entry>(ASYNC_LOG_CHANNEL_CAPACITY);
let session_id = config.session_id;
let task = tokio::spawn(async move {
while let Some(entry) = rx.recv().await {
match entry {
Entry::Message { direction, text } => {
let collection = if direction == "I" {
&incoming
} else {
&outgoing
};
let _ = collection
.insert_one(doc! {
"text": text,
"logged_at": now_unix(),
"session_id": &session_id,
})
.await;
}
Entry::Event { text } => {
let _ = event
.insert_one(doc! {
"text": text,
"logged_at": now_unix(),
"session_id": &session_id,
})
.await;
}
}
}
});
Ok(Self {
tx: std::sync::Mutex::new(Some(tx)),
task: std::sync::Mutex::new(Some(task)),
include_heartbeats: config.include_heartbeats,
})
}
fn send(&self, entry: Entry) {
if let Ok(guard) = self.tx.lock()
&& let Some(tx) = guard.as_ref()
{
let _ = tx.try_send(entry);
}
}
}
#[async_trait::async_trait]
impl Log for MongoLog {
fn on_incoming(&self, message: &str) {
if is_heartbeat(message) && !self.include_heartbeats {
return;
}
self.send(Entry::Message {
direction: "I",
text: message.to_owned(),
});
}
fn on_outgoing(&self, message: &str) {
if is_heartbeat(message) && !self.include_heartbeats {
return;
}
self.send(Entry::Message {
direction: "O",
text: message.to_owned(),
});
}
fn on_event(&self, text: &str) {
self.send(Entry::Event {
text: text.to_owned(),
});
}
async fn shutdown(&self) {
let tx = self.tx.lock().ok().and_then(|mut guard| guard.take());
drop(tx);
let task = self.task.lock().ok().and_then(|mut guard| guard.take());
if let Some(task) = task {
let _ = task.await;
}
}
}