#![cfg(feature = "mongodb")]
#![allow(dead_code)]
use crate::tina::data::date_time::LocalDateTime;
use crate::tina::data::AppResult;
use crate::tina::log::CloseableAppender;
use crate::tina::util::string::AsStr;
use crate::tina::{constant::Constants, data::date_time::DateTimeExt};
use crate::{app_error_from, app_error_from_none_static, app_system_error};
use anyhow::anyhow;
use futures_util::FutureExt;
use log::{Level, LevelFilter, Record};
use log4rs::append::Append;
use mongodb::bson::{doc, to_bson, Bson, DateTime};
use mongodb::options::IndexOptions;
use mongodb::{Collection, Database, IndexModel};
use once_cell::sync::Lazy;
use std::borrow::Cow;
use std::fmt::Arguments;
use std::panic::AssertUnwindSafe;
use std::sync::{Arc, RwLock};
use std::thread::Thread;
use std::time::Duration;
use tokio::runtime::Runtime;
use tokio::sync::mpsc::{channel, Sender};
use tokio::sync::Mutex;
static LOGGER_TARGET: Lazy<RwLock<Vec<(String, String)>>> = Lazy::new(|| RwLock::new(Vec::new()));
pub fn add_logger(name: &str, label: &str) {
let mut lock = LOGGER_TARGET.write().expect("write LOGGER_TARGET failed");
lock.push((name.to_string(), label.to_string()));
}
pub fn get_loggers() -> AppResult<Vec<(String, String)>> {
let lock = LOGGER_TARGET.read().map_err(app_error_from_none_static!())?;
Ok((*lock).clone())
}
#[derive(Debug)]
pub(crate) struct MongodbAppender {
database: Database,
collection: Collection<Bson>,
level: LevelFilter,
dispatch_runtime: Arc<std::sync::Mutex<Option<Runtime>>>,
queue: Arc<Mutex<Vec<Bson>>>,
sender: Arc<Sender<bool>>,
}
impl MongodbAppender {
const UTC_TIME_INDEX: &str = "idx_utc_time";
const QUEUE_CAPACITY: usize = 2000;
pub async fn new(appender_name: &str, database: Database, level: LevelFilter, expires: Duration) -> AppResult<Self> {
let collection = Self::init_collection(appender_name, database.clone(), expires).await?;
let dispatch_runtime = tokio::runtime::Builder::new_current_thread()
.worker_threads(1)
.thread_name(format!("{}分发线程", appender_name))
.enable_all()
.build()
.map_err(app_error_from!())?;
let dispatch_runtime = Arc::new(std::sync::Mutex::new(Some(dispatch_runtime)));
let execute_runtime = tokio::runtime::Builder::new_current_thread()
.worker_threads(1)
.thread_name(format!("{}处理执行线程", appender_name))
.enable_all()
.build()
.map_err(app_error_from!())?;
let queue = Arc::new(Mutex::new(Vec::with_capacity(Self::QUEUE_CAPACITY)));
let queue1 = queue.clone();
let collection1 = collection.clone();
let (sender, receiver) = channel::<bool>(Self::QUEUE_CAPACITY);
let receiver = Arc::new(Mutex::new(receiver));
let fut_fn = move || {
let recerver1 = receiver.clone();
let queue2 = queue1.clone();
let collection2 = collection1.clone();
AssertUnwindSafe(async move {
let mut recv = recerver1.lock().await;
loop {
if let Some(flag) = recv.recv().await {
if !flag {
break;
}
let data = {
let mut list = queue2.lock().await;
let mut l = Vec::with_capacity(list.len());
while let Some(d) = list.pop() {
l.push(d);
}
l
};
if !data.is_empty() {
if let Err(err) = collection2.insert_many(data, None).await {
println!("{:?}", err);
}
}
}
}
})
.catch_unwind()
};
let appender_name = appender_name.to_string();
std::thread::Builder::new()
.name(format!("{}处理调度线程", appender_name))
.spawn(move || {
if execute_runtime.block_on(fut_fn()).is_err() {
println!("panic occurred in {} execute thread", appender_name);
}
execute_runtime.shutdown_background();
})
.map_err(app_error_from!())?;
Ok(Self {
database,
collection,
level,
dispatch_runtime,
queue,
sender: Arc::new(sender),
})
}
async fn init_collection(appender_name: &str, database: Database, expires: Duration) -> AppResult<Collection<Bson>> {
let collection_names = database.list_collection_names(None).await.map_err(app_error_from!())?;
let collection = match collection_names.contains(&appender_name.to_string()) {
true => database.collection::<Bson>(appender_name),
false => {
database
.create_collection(appender_name, None)
.await
.map_err(|err| app_system_error!("创建日志{}失败, 请使用mongodb 5.0 以上版本再试! reason: {:?}", appender_name, err))?;
database.collection::<Bson>(appender_name)
}
};
{
let index_name = Self::UTC_TIME_INDEX;
let index = IndexModel::builder()
.keys(doc! { "utc_time": 1 })
.options(Some(IndexOptions::builder().name(Some(index_name.to_string())).expire_after(Some(expires)).build()))
.build();
create_index(index_name, index, &collection).await?;
}
Ok(collection)
}
fn wait_for_shutdown(&self) {
println!("停止MongodbAppender开始: {}", self.collection.name());
match self.dispatch_runtime.lock() {
Ok(mut dispatch_runtime) => {
if let Some(dispatch_runtime) = dispatch_runtime.take() {
dispatch_runtime.block_on(async {
if let Err(err) = self.sender.send(false).await {
println!("{:?}", err);
}
});
dispatch_runtime.shutdown_background();
}
}
Err(_) => {
println!("停止MongodbAppender异常: 锁定dispatch_runtime失败");
}
}
println!("停止MongodbAppender完成: {}", self.collection.name());
}
}
impl CloseableAppender for MongodbAppender {
fn shutdown(&self) {
self.wait_for_shutdown();
}
}
impl Drop for MongodbAppender {
fn drop(&mut self) {
self.wait_for_shutdown();
}
}
async fn create_index(index_name: &str, index: IndexModel, collection: &Collection<Bson>) -> AppResult<()> {
match collection.create_index(index.clone(), None).await {
Ok(create_result) => {
let create_index_name = create_result.index_name.as_str();
if create_index_name != index_name {
collection.drop_index(create_index_name, None).await.map_err(app_error_from!())?;
return Err(app_system_error!("创建日志索引失败: {}", index_name));
}
}
Err(_) => {
println!("删除旧索引开始: {}", index_name);
if collection.drop_index(index_name, None).await.is_err() {}
println!("删除旧索引完成: {}", index_name);
println!("创建新索引开始: {}", index_name);
let create_result = collection.create_index(index, None).await.map_err(app_error_from!())?;
let create_index_name = create_result.index_name.as_str();
if create_index_name != index_name {
collection.drop_index(create_index_name, None).await.map_err(app_error_from!())?;
return Err(app_system_error!("创建日志索引失败: {}", index_name));
}
println!("创建新索引完成: {}", index_name);
}
}
Ok(())
}
fn ser_display<T, S>(v: &T, s: S) -> Result<S::Ok, S::Error>
where
T: std::fmt::Display,
S: serde::ser::Serializer,
{
s.collect_str(v)
}
struct JsonWriter(Vec<u8>);
impl std::io::Write for JsonWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl log4rs::encode::Write for JsonWriter {}
#[derive(serde::Serialize)]
struct JsonRecord<'a> {
time: Bson,
#[serde(serialize_with = "ser_display")]
message: &'a Arguments<'a>,
#[serde(skip_serializing_if = "Option::is_none")]
module_path: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
file: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
line: Option<u32>,
level: Level,
target: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
thread: Option<Cow<'a, str>>,
thread_id: Cow<'a, str>,
#[serde(skip_serializing_if = "Option::is_none")]
request_id: Option<Cow<'a, str>>,
utc_time: Bson,
}
impl<'a> JsonRecord<'a> {
fn from_record(record: &'a Record, current_thread: &'a Thread) -> AppResult<Self> {
let log_time = match log_mdc::remove("business_thread_time") {
None => LocalDateTime::now(),
Some(v) => LocalDateTime::parse_from_str(v.as_str(), "%Y-%m-%d %H:%M:%S%.3f")?,
};
let time = Bson::DateTime(DateTime::from_millis(log_time.naive_local().timestamp_millis()));
let utc_time = Bson::DateTime(DateTime::from_millis(log_time.naive_utc().timestamp_millis()));
let thread_id = match log_mdc::remove("business_thread_id") {
None => Cow::Owned(thread_id::get().to_string()),
Some(v) => Cow::Owned(v),
};
let thread = match log_mdc::remove("business_thread_name") {
None => current_thread.name().map(Cow::Borrowed),
Some(v) => Some(Cow::Owned(v)),
};
let request_id = log_mdc::remove(Constants::REQUEST_ID_HEADER_VALUE).map(Cow::Owned);
Ok(Self {
time,
message: record.args(),
module_path: record.module_path(),
file: record.file(),
line: record.line(),
level: record.level(),
target: record.target(),
thread,
thread_id,
request_id,
utc_time,
})
}
}
impl Append for MongodbAppender {
fn append(&self, record: &Record) -> anyhow::Result<()> {
if self.level < record.level() {
return Ok(());
}
if record.module_path().as_str().starts_with("mongodb::") {
return Ok(());
}
let thread = std::thread::current();
let encode = JsonRecord::from_record(record, &thread).map_err(anyhow::Error::new)?;
let bson = to_bson(&encode).map_err(anyhow::Error::new)?;
let lock = self.dispatch_runtime.lock().map_err(|err| anyhow!("{:?}", err))?;
if let Some(dispatch_runtime) = lock.as_ref() {
dispatch_runtime.block_on(async {
{
let mut datas = self.queue.lock().await;
datas.push(bson);
}
if let Err(err) = self.sender.send(true).await {
println!("{:?}", err);
}
});
}
Ok(())
}
fn flush(&self) {}
}