use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::fmt::Display;
use std::fmt::Formatter;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::Condvar;
use std::sync::Mutex as StdMutex;
use std::thread;
use std::time::Duration;
use cheetah_string::CheetahString;
use parking_lot::RwLock;
use rocketmq_error::RocketMQError;
use tokio::sync::Notify;
use tracing::error;
use tracing::info;
use tracing::warn;
use crate::base::transient_store_pool::TransientStorePool;
use crate::config::flush_disk_type::FlushDiskType;
use crate::config::message_store_config::MessageStoreConfig;
use crate::log_file::mapped_file::default_mapped_file_impl::DefaultMappedFile;
use crate::log_file::mapped_file::MappedFile;
const WAIT_TIMEOUT: Duration = Duration::from_secs(5);
#[derive(Clone, Copy)]
struct WarmMappedFileConfig {
enabled: bool,
flush_disk_type: FlushDiskType,
mapped_file_size_commit_log: usize,
flush_least_pages_when_warm_mapped_file: usize,
}
impl WarmMappedFileConfig {
fn disabled() -> Self {
Self {
enabled: false,
flush_disk_type: FlushDiskType::AsyncFlush,
mapped_file_size_commit_log: usize::MAX,
flush_least_pages_when_warm_mapped_file: 0,
}
}
fn from_message_store_config(message_store_config: &MessageStoreConfig) -> Self {
Self {
enabled: message_store_config.warm_mapped_file_enable,
flush_disk_type: message_store_config.flush_disk_type,
mapped_file_size_commit_log: message_store_config.mapped_file_size_commit_log,
flush_least_pages_when_warm_mapped_file: message_store_config.flush_least_pages_when_warm_mapped_file,
}
}
fn should_warm(self, file_size: u64) -> bool {
self.enabled && file_size as usize >= self.mapped_file_size_commit_log
}
}
pub struct AllocateMappedFileService {
request_table: Arc<RwLock<HashMap<String, Arc<AllocateRequest>>>>,
request_queue: Arc<RwLock<BinaryHeap<Arc<AllocateRequest>>>>,
has_exception: Arc<AtomicBool>,
stopped: Arc<AtomicBool>,
notify: Arc<Notify>,
worker_wakeup: Arc<(StdMutex<()>, Condvar)>,
worker_handle: Arc<parking_lot::Mutex<Option<thread::JoinHandle<()>>>>,
transient_store_pool: Option<Arc<TransientStorePool>>,
transient_store_pool_enable: bool,
fast_fail_if_no_buffer: bool,
warm_mapped_file_config: WarmMappedFileConfig,
}
impl Clone for AllocateMappedFileService {
fn clone(&self) -> Self {
Self {
request_table: self.request_table.clone(),
request_queue: self.request_queue.clone(),
has_exception: self.has_exception.clone(),
stopped: self.stopped.clone(),
notify: self.notify.clone(),
worker_wakeup: self.worker_wakeup.clone(),
worker_handle: self.worker_handle.clone(),
transient_store_pool: self.transient_store_pool.clone(),
transient_store_pool_enable: self.transient_store_pool_enable,
fast_fail_if_no_buffer: self.fast_fail_if_no_buffer,
warm_mapped_file_config: self.warm_mapped_file_config,
}
}
}
impl Default for AllocateMappedFileService {
fn default() -> Self {
Self::new()
}
}
impl AllocateMappedFileService {
pub fn new_with_config(
transient_store_pool: Option<Arc<TransientStorePool>>,
transient_store_pool_enable: bool,
fast_fail_if_no_buffer: bool,
) -> Self {
let request_table = Arc::new(RwLock::new(HashMap::new()));
let request_queue = Arc::new(RwLock::new(BinaryHeap::new()));
let has_exception = Arc::new(AtomicBool::new(false));
let stopped = Arc::new(AtomicBool::new(false));
let notify = Arc::new(Notify::new());
let worker_wakeup = Arc::new((StdMutex::new(()), Condvar::new()));
Self {
request_table,
request_queue,
has_exception,
stopped,
notify,
worker_wakeup,
worker_handle: Arc::new(parking_lot::Mutex::new(None)),
transient_store_pool,
transient_store_pool_enable,
fast_fail_if_no_buffer,
warm_mapped_file_config: WarmMappedFileConfig::disabled(),
}
}
pub fn new_with_message_store_config(
transient_store_pool: Option<Arc<TransientStorePool>>,
transient_store_pool_enable: bool,
fast_fail_if_no_buffer: bool,
message_store_config: &MessageStoreConfig,
) -> Self {
let mut service = Self::new_with_config(
transient_store_pool,
transient_store_pool_enable,
fast_fail_if_no_buffer,
);
service.warm_mapped_file_config = WarmMappedFileConfig::from_message_store_config(message_store_config);
service
}
pub fn new() -> Self {
Self::new_with_config(None, false, false)
}
pub fn is_started(&self) -> bool {
self.worker_handle.lock().is_some() && !self.stopped.load(Ordering::Acquire)
}
#[cfg(test)]
pub(crate) fn should_warm_mapped_file(&self, file_size: u64) -> bool {
self.warm_mapped_file_config.should_warm(file_size)
}
#[cfg(test)]
pub(crate) fn has_request(&self, file_path: &str) -> bool {
self.request_table.read().contains_key(file_path)
}
pub fn start(&self) {
{
let worker_handle = self.worker_handle.lock();
if worker_handle.is_some() {
return;
}
}
self.stopped.store(false, Ordering::Relaxed);
let request_table = self.request_table.clone();
let request_queue = self.request_queue.clone();
let has_exception = self.has_exception.clone();
let stopped = self.stopped.clone();
let transient_store_pool = self.transient_store_pool.clone();
let worker_wakeup = self.worker_wakeup.clone();
let warm_mapped_file_config = self.warm_mapped_file_config;
match thread::Builder::new()
.name("allocate-mapped-file-service".to_string())
.spawn(move || {
Self::run_worker(
request_table,
request_queue,
has_exception,
stopped,
transient_store_pool,
worker_wakeup,
warm_mapped_file_config,
);
}) {
Ok(handle) => {
*self.worker_handle.lock() = Some(handle);
info!("AllocateMappedFileService started");
}
Err(error) => {
self.has_exception.store(true, Ordering::Relaxed);
error!("AllocateMappedFileService failed to start worker thread: {}", error);
}
}
}
fn run_worker(
request_table: Arc<RwLock<HashMap<String, Arc<AllocateRequest>>>>,
request_queue: Arc<RwLock<BinaryHeap<Arc<AllocateRequest>>>>,
has_exception: Arc<AtomicBool>,
stopped: Arc<AtomicBool>,
transient_store_pool: Option<Arc<TransientStorePool>>,
worker_wakeup: Arc<(StdMutex<()>, Condvar)>,
warm_mapped_file_config: WarmMappedFileConfig,
) {
info!("AllocateMappedFileService: service started");
while !stopped.load(Ordering::Relaxed) {
while !stopped.load(Ordering::Relaxed)
&& Self::mmap_operation(
&request_table,
&request_queue,
&has_exception,
&transient_store_pool,
warm_mapped_file_config,
)
{}
if stopped.load(Ordering::Relaxed) {
break;
}
if request_queue.read().is_empty() {
let (lock, condvar) = &*worker_wakeup;
let guard = lock.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
match condvar.wait_timeout(guard, Duration::from_millis(100)) {
Ok((_guard, _timeout)) => {}
Err(poisoned) => {
let (_guard, _timeout) = poisoned.into_inner();
}
}
}
}
info!("AllocateMappedFileService: service end");
}
fn mmap_operation(
request_table: &Arc<RwLock<HashMap<String, Arc<AllocateRequest>>>>,
request_queue: &Arc<RwLock<BinaryHeap<Arc<AllocateRequest>>>>,
has_exception: &Arc<AtomicBool>,
transient_store_pool: &Option<Arc<TransientStorePool>>,
warm_mapped_file_config: WarmMappedFileConfig,
) -> bool {
let req = {
let mut queue = request_queue.write();
queue.pop()
};
let req = match req {
Some(r) => r,
None => return false, };
let expected_request = {
let table = request_table.read();
table.get(&req.file_path).cloned()
};
let expected_request = match expected_request {
Some(r) => r,
None => {
warn!(
"this mmap request expired, maybe cause timeout {} {}",
req.file_path, req.file_size
);
return true;
}
};
if !Arc::ptr_eq(&expected_request, &req) {
warn!(
"never expected here, maybe cause timeout {} {}",
req.file_path, req.file_size
);
return true;
}
if req.mapped_file.read().is_some() {
return true;
}
let result = Self::create_mapped_file(&req, transient_store_pool, warm_mapped_file_config);
match result {
Ok(mapped_file) => {
*req.mapped_file.write() = Some(mapped_file);
has_exception.store(false, Ordering::Relaxed);
req.complete();
true
}
Err(e) => {
error!(
"AllocateMappedFileService: failed to create mapped file {}: {}",
req.file_path, e
);
has_exception.store(true, Ordering::Relaxed);
request_queue.write().push(req);
thread::sleep(Duration::from_millis(1));
false
}
}
}
fn create_mapped_file(
req: &AllocateRequest,
transient_store_pool: &Option<Arc<TransientStorePool>>,
warm_mapped_file_config: WarmMappedFileConfig,
) -> Result<Arc<DefaultMappedFile>, RocketMQError> {
let start = std::time::Instant::now();
let file_path = req.file_path.clone();
let file_size = req.file_size as u64;
let transient_pool = transient_store_pool.clone();
let mapped_file = if let Some(pool) = transient_pool {
DefaultMappedFile::try_new_with_transient_store_pool(
CheetahString::from_string(file_path.clone()),
file_size,
(*pool).clone(),
)
} else {
DefaultMappedFile::try_new(CheetahString::from_string(file_path.clone()), file_size)
}
.map_err(|error| RocketMQError::StorageWriteFailed {
path: req.file_path.clone(),
reason: error.to_string(),
})?;
if warm_mapped_file_config.should_warm(file_size) {
mapped_file.warm_mapped_file(
warm_mapped_file_config.flush_disk_type,
warm_mapped_file_config.flush_least_pages_when_warm_mapped_file,
);
}
let elapsed = start.elapsed();
if elapsed.as_millis() > 10 {
let queue_size = 0; warn!(
"create mappedFile spent time(ms) {} queue size {} {} {}",
elapsed.as_millis(),
queue_size,
req.file_path,
req.file_size
);
}
Ok(Arc::new(mapped_file))
}
fn notify_worker(&self) {
self.notify.notify_one();
let (_, condvar) = &*self.worker_wakeup;
condvar.notify_one();
}
pub async fn put_request_and_return_mapped_file(
&self,
next_file_path: String,
next_next_file_path: String,
file_size: i32,
) -> Result<Option<Arc<DefaultMappedFile>>, RocketMQError> {
let mut can_submit_requests = 2;
if self.transient_store_pool_enable {
if let Some(ref pool) = self.transient_store_pool {
if self.fast_fail_if_no_buffer {
let queue_size = self.request_queue.read().len();
can_submit_requests = pool.available_buffer_nums().saturating_sub(queue_size);
}
}
}
let next_req = Arc::new(AllocateRequest::new(next_file_path.clone(), file_size));
let next_put_ok = {
let mut table = self.request_table.write();
if table.contains_key(&next_file_path) {
false
} else {
table.insert(next_file_path.clone(), next_req.clone());
true
}
};
if next_put_ok {
if can_submit_requests == 0 {
warn!(
"[NOTIFYME]TransientStorePool is not enough, so create mapped file error, RequestQueueSize: {}, \
StorePoolSize: {}",
self.request_queue.read().len(),
self.transient_store_pool
.as_ref()
.map_or(0, |p| p.available_buffer_nums())
);
self.request_table.write().remove(&next_file_path);
return Ok(None);
}
self.request_queue.write().push(next_req.clone());
self.notify_worker();
can_submit_requests -= 1;
}
if !next_next_file_path.is_empty() {
let next_next_req = Arc::new(AllocateRequest::new(next_next_file_path.clone(), file_size));
let next_next_put_ok = {
let mut table = self.request_table.write();
if table.contains_key(&next_next_file_path) {
false
} else {
table.insert(next_next_file_path.clone(), next_next_req.clone());
true
}
};
if next_next_put_ok {
if can_submit_requests == 0 {
warn!(
"[NOTIFYME]TransientStorePool is not enough, so skip preallocate mapped file, \
RequestQueueSize: {}, StorePoolSize: {}",
self.request_queue.read().len(),
self.transient_store_pool
.as_ref()
.map_or(0, |p| p.available_buffer_nums())
);
self.request_table.write().remove(&next_next_file_path);
} else {
self.request_queue.write().push(next_next_req);
self.notify_worker();
}
}
}
if self.has_exception.load(Ordering::Relaxed) {
warn!("AllocateMappedFileService has exception, so return null");
return Ok(None);
}
let result = {
let table = self.request_table.read();
table.get(&next_file_path).cloned()
};
if let Some(req) = result {
let wait_result = tokio::time::timeout(WAIT_TIMEOUT, req.wait()).await;
match wait_result {
Ok(()) => {
self.request_table.write().remove(&next_file_path);
let mapped_file = req.mapped_file.read().clone();
Ok(mapped_file)
}
Err(_) => {
warn!("create mmap timeout {} {}", req.file_path, req.file_size);
Ok(None)
}
}
} else {
error!("find preallocate mmap failed, this never happen");
Ok(None)
}
}
pub async fn submit_request(
&self,
file_path: String,
file_size: u64,
) -> Result<Arc<DefaultMappedFile>, RocketMQError> {
let result = self
.put_request_and_return_mapped_file(
file_path.clone(),
String::new(), file_size as i32,
)
.await?;
result.ok_or_else(|| RocketMQError::StorageWriteFailed {
path: file_path.clone(),
reason: "Allocation failed or timed out".to_string(),
})
}
pub async fn allocate_mapped_file(
&self,
file_path: String,
file_size: u64,
) -> Result<Arc<DefaultMappedFile>, RocketMQError> {
self.submit_request(file_path, file_size).await
}
pub fn allocate_mapped_file_blocking(
&self,
file_path: String,
file_size: u64,
) -> Result<Arc<DefaultMappedFile>, RocketMQError> {
let result =
self.put_request_and_return_mapped_file_blocking(file_path.clone(), String::new(), file_size as i32)?;
result.ok_or_else(|| RocketMQError::StorageWriteFailed {
path: file_path,
reason: "Allocation failed or timed out".to_string(),
})
}
fn put_request_and_return_mapped_file_blocking(
&self,
next_file_path: String,
next_next_file_path: String,
file_size: i32,
) -> Result<Option<Arc<DefaultMappedFile>>, RocketMQError> {
let next_req = Arc::new(AllocateRequest::new(next_file_path.clone(), file_size));
let next_put_ok = {
let mut table = self.request_table.write();
if table.contains_key(&next_file_path) {
false
} else {
table.insert(next_file_path.clone(), next_req.clone());
true
}
};
if next_put_ok {
self.request_queue.write().push(next_req.clone());
self.notify_worker();
}
if !next_next_file_path.is_empty() {
let next_next_req = Arc::new(AllocateRequest::new(next_next_file_path.clone(), file_size));
let next_next_put_ok = {
let mut table = self.request_table.write();
if table.contains_key(&next_next_file_path) {
false
} else {
table.insert(next_next_file_path.clone(), next_next_req.clone());
true
}
};
if next_next_put_ok {
self.request_queue.write().push(next_next_req);
self.notify_worker();
}
}
if self.has_exception.load(Ordering::Relaxed) {
warn!("AllocateMappedFileService has exception, so return null");
return Ok(None);
}
let result = {
let table = self.request_table.read();
table.get(&next_file_path).cloned()
};
if let Some(req) = result {
if req.wait_blocking(WAIT_TIMEOUT) {
self.request_table.write().remove(&next_file_path);
Ok(req.mapped_file.read().clone())
} else {
warn!("create mmap timeout {} {}", req.file_path, req.file_size);
Ok(None)
}
} else {
error!("find preallocate mmap failed, this never happen");
Ok(None)
}
}
pub fn submit_request_in_background(&self, file_path: String, file_size: u64) {
let mut can_submit_request = true;
if self.transient_store_pool_enable && self.fast_fail_if_no_buffer {
if let Some(ref pool) = self.transient_store_pool {
let queue_size = self.request_queue.read().len();
can_submit_request = pool.available_buffer_nums().saturating_sub(queue_size) > 0;
}
}
if !can_submit_request {
warn!(
"[NOTIFYME]TransientStorePool is not enough, so skip background preallocate mapped file, \
RequestQueueSize: {}, StorePoolSize: {}",
self.request_queue.read().len(),
self.transient_store_pool
.as_ref()
.map_or(0, |pool| pool.available_buffer_nums())
);
return;
}
let req = Arc::new(AllocateRequest::new(file_path.clone(), file_size as i32));
let put_ok = {
let mut table = self.request_table.write();
if let std::collections::hash_map::Entry::Vacant(entry) = table.entry(file_path) {
entry.insert(req.clone());
true
} else {
false
}
};
if put_ok {
self.request_queue.write().push(req);
self.notify_worker();
}
}
pub async fn shutdown(&self) {
info!("AllocateMappedFileService: shutting down");
self.stopped.store(true, Ordering::Relaxed);
self.notify_worker();
let (_, condvar) = &*self.worker_wakeup;
condvar.notify_all();
let handle = self.worker_handle.lock().take();
if let Some(handle) = handle {
let _ = tokio::task::spawn_blocking(move || handle.join()).await;
}
let table = self.request_table.read();
for req in table.values() {
if let Some(ref mapped_file) = *req.mapped_file.read() {
info!("delete pre allocated mapped file, {}", req.file_path);
mapped_file.destroy(1000);
}
}
info!("AllocateMappedFileService: shutdown complete");
}
pub fn get_service_name(&self) -> &'static str {
"AllocateMappedFileService"
}
pub fn has_exception(&self) -> bool {
self.has_exception.load(Ordering::Relaxed)
}
}
struct AllocateRequest {
file_path: String,
file_size: i32,
completion: Arc<Notify>,
blocking_completion: Arc<(StdMutex<()>, Condvar)>,
completed: Arc<AtomicBool>,
mapped_file: Arc<RwLock<Option<Arc<DefaultMappedFile>>>>,
}
impl AllocateRequest {
fn new(file_path: String, file_size: i32) -> Self {
Self {
file_path,
file_size,
completion: Arc::new(Notify::new()),
blocking_completion: Arc::new((StdMutex::new(()), Condvar::new())),
completed: Arc::new(AtomicBool::new(false)),
mapped_file: Arc::new(RwLock::new(None)),
}
}
async fn wait(&self) {
if !self.completed.load(Ordering::Acquire) {
self.completion.notified().await;
}
}
fn wait_blocking(&self, timeout: Duration) -> bool {
if self.completed.load(Ordering::Acquire) {
return true;
}
let (lock, condvar) = &*self.blocking_completion;
let guard = lock.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
match condvar.wait_timeout_while(guard, timeout, |_| !self.completed.load(Ordering::Acquire)) {
Ok((_guard, _timeout)) => {}
Err(poisoned) => {
let (_guard, _timeout) = poisoned.into_inner();
}
}
self.completed.load(Ordering::Acquire)
}
fn complete(&self) {
self.completed.store(true, Ordering::Release);
self.completion.notify_waiters();
let (_, condvar) = &*self.blocking_completion;
condvar.notify_all();
}
fn file_offset(&self) -> i64 {
if let Some(separator_idx) = self.file_path.rfind(std::path::MAIN_SEPARATOR) {
if let Ok(offset) = self.file_path[(separator_idx + 1)..].parse::<i64>() {
return offset;
}
}
0
}
}
impl Display for AllocateRequest {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"AllocateRequest[file_path={},file_size={}]",
self.file_path, self.file_size
)
}
}
impl PartialEq for AllocateRequest {
fn eq(&self, other: &Self) -> bool {
self.file_path == other.file_path && self.file_size == other.file_size
}
}
impl Eq for AllocateRequest {}
impl PartialOrd for AllocateRequest {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for AllocateRequest {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
other.file_offset().cmp(&self.file_offset())
}
}
#[cfg(test)]
mod tests {
use tempfile::tempdir;
use super::*;
use crate::config::message_store_config::MessageStoreConfig;
#[tokio::test]
async fn allocate_mapped_file_blocking_works_inside_runtime() {
let temp_dir = tempdir().expect("temp dir");
let file_path = temp_dir.path().join("00000000000000000000");
let service = AllocateMappedFileService::new();
assert!(!service.is_started());
service.start();
assert!(service.is_started());
let mapped_file = service
.allocate_mapped_file_blocking(file_path.to_string_lossy().to_string(), 1024)
.expect("allocate mapped file");
assert_eq!(mapped_file.get_file_size(), 1024);
assert!(file_path.exists(), "mapped file should be created on disk");
service.shutdown().await;
assert!(!service.is_started());
}
#[test]
fn warm_mapped_file_config_follows_commitlog_file_size_threshold() {
let config = MessageStoreConfig {
warm_mapped_file_enable: true,
mapped_file_size_commit_log: 1024,
flush_least_pages_when_warm_mapped_file: 1,
..MessageStoreConfig::default()
};
let service = AllocateMappedFileService::new_with_message_store_config(None, false, false, &config);
assert!(!service.should_warm_mapped_file(1023));
assert!(service.should_warm_mapped_file(1024));
}
}