rocketmq-store 0.6.0

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use bytes::Buf;
use bytes::BufMut;
use cheetah_string::CheetahString;
use rocketmq_common::common::config::TopicConfig;
use rocketmq_common::common::message::message_batch::MessageExtBatch;
use rocketmq_common::common::message::message_ext_broker_inner::MessageExtBrokerInner;
use rocketmq_common::common::sys_flag::message_sys_flag::MessageSysFlag;
use rocketmq_common::utils::message_utils;
use rocketmq_common::CRC32Utils::crc32;
use rocketmq_common::MessageDecoder::create_crc32;
use rocketmq_common::MessageUtils::build_batch_message_id;
use rocketmq_rust::SyncUnsafeCellWrapper;

use crate::base::message_result::AppendMessageResult;
use crate::base::message_status_enum::AppendMessageStatus;
use crate::base::put_message_context::PutMessageContext;
use crate::config::message_store_config::MessageStoreConfig;
use crate::log_file::commit_log::get_message_num;
use crate::log_file::commit_log::CommitLog;
use crate::log_file::commit_log::BLANK_MAGIC_CODE;
use crate::log_file::commit_log::CRC32_RESERVED_LEN;
use crate::log_file::mapped_file::MappedFile;

/// Write messages callback interface
pub trait AppendMessageCallback {
    /// After message serialization, write MappedByteBuffer
    ///
    /// # Arguments
    ///
    /// * `file_from_offset` - The offset of the file
    /// * `byte_buffer` - The buffer to write
    /// * `max_blank` - The maximum blank space
    /// * `msg` - The message to write
    /// * `put_message_context` - The context of putting message
    ///
    /// # Returns
    ///
    /// The number of bytes written
    fn do_append<MF: MappedFile>(
        &self,
        file_from_offset: i64,
        mapped_file: &MF,
        max_blank: i32,
        msg: &mut MessageExtBrokerInner,
        put_message_context: &PutMessageContext,
    ) -> AppendMessageResult;

    fn do_append_batch<MF: MappedFile>(
        &self,
        file_from_offset: i64,
        mapped_file: &MF,
        max_blank: i32,
        msg: &mut MessageExtBatch,
        put_message_context: &mut PutMessageContext,
        enabled_append_prop_crc: bool,
    ) -> AppendMessageResult;
}

// File at the end of the minimum fixed length empty
const END_FILE_MIN_BLANK_LENGTH: i32 = 4 + 4;

pub(crate) struct DefaultAppendMessageCallback {
    msg_store_item_memory: SyncUnsafeCellWrapper<bytes::BytesMut>,
    crc32_reserved_length: i32,
    message_store_config: Arc<MessageStoreConfig>,
    topic_config_table: Arc<parking_lot::Mutex<HashMap<CheetahString, TopicConfig>>>,
}

impl DefaultAppendMessageCallback {
    pub fn new(
        message_store_config: Arc<MessageStoreConfig>,
        topic_config_table: Arc<parking_lot::Mutex<HashMap<CheetahString, TopicConfig>>>,
    ) -> Self {
        let crc32_reserved_length = if message_store_config.enabled_append_prop_crc {
            CRC32_RESERVED_LEN
        } else {
            0
        };
        Self {
            msg_store_item_memory: SyncUnsafeCellWrapper::new(bytes::BytesMut::with_capacity(
                END_FILE_MIN_BLANK_LENGTH as usize,
            )),
            crc32_reserved_length,
            message_store_config,
            topic_config_table,
        }
    }
}

impl AppendMessageCallback for DefaultAppendMessageCallback {
    fn do_append<MF: MappedFile>(
        &self,
        file_from_offset: i64,
        mapped_file: &MF,
        max_blank: i32,
        msg_inner: &mut MessageExtBrokerInner,
        put_message_context: &PutMessageContext,
    ) -> AppendMessageResult {
        let mut pre_encode_buffer = msg_inner.encoded_buff.take().unwrap(); // Assuming get_encoded_buff returns Option<ByteBuffer>
        let is_multi_dispatch_msg = self.message_store_config.enable_multi_dispatch
            && CommitLog::is_multi_dispatch_msg(msg_inner);
        if is_multi_dispatch_msg {
            unimplemented!("Multi dispatch message is not supported yet");
        }

        let msg_len = i32::from_be_bytes(pre_encode_buffer[0..4].try_into().unwrap());
        //physic offset
        let wrote_offset = file_from_offset + mapped_file.get_wrote_position() as i64;
        let addr = msg_inner.message_ext_inner.store_host;
        let msg_id_supplier =
            move || -> String { message_utils::build_message_id(addr, wrote_offset) };

        let mut queue_offset = msg_inner.queue_offset();
        let message_num = get_message_num(&self.topic_config_table, msg_inner);
        // Transaction messages that require special handling
        if let MessageSysFlag::TRANSACTION_PREPARED_TYPE
        | MessageSysFlag::TRANSACTION_ROLLBACK_TYPE =
            MessageSysFlag::get_transaction_value(msg_inner.sys_flag())
        {
            queue_offset = 0;
        }

        // Determines whether there is sufficient free space
        if (msg_len + END_FILE_MIN_BLANK_LENGTH) > max_blank {
            let bytes = self.msg_store_item_memory.mut_from_ref();
            bytes.clear();
            bytes.put_i32(max_blank);
            bytes.put_i32(BLANK_MAGIC_CODE);
            let instant = Instant::now();
            mapped_file.write_bytes_segment(bytes.as_ref(), wrote_offset as usize, 0, bytes.len());
            return AppendMessageResult {
                status: AppendMessageStatus::EndOfFile,
                wrote_offset,
                /* only wrote 8 bytes, but declare wrote maxBlank for compute write position */
                wrote_bytes: max_blank,
                store_timestamp: msg_inner.store_timestamp(),
                logics_offset: queue_offset,
                msg_num: message_num as i32,
                msg_id_supplier: Some(Arc::new(msg_id_supplier)),
                page_cache_rt: instant.elapsed().as_millis() as i64,
                ..Default::default()
            };
        }

        /*        let mut pos = 4 // 1 TOTALSIZE
        + 4// 2 MAGICCODE
        + 4// 3 BODYCRC
        + 4 // 4 QUEUEID
        + 4; // 5 FLAG*/
        let mut pos = 20; // TOTALSIZE +  MAGICCODE + BODYCRC + QUEUEID + FLAG
        pre_encode_buffer[pos..(pos + 8)].copy_from_slice(&queue_offset.to_be_bytes()); // 6 QUEUEOFFSET
        pos += 8;
        pre_encode_buffer[pos..(pos + 8)].copy_from_slice(&wrote_offset.to_be_bytes()); // 7 PHYSICALOFFSET
        let ip_len = if msg_inner.sys_flag() & MessageSysFlag::BORNHOST_V6_FLAG == 0 {
            4 + 4
        } else {
            16 + 4
        };
        // 8 SYSFLAG, 9 BORNTIMESTAMP, 10 BORNHOST
        pos += 8 + 4 + 8 + ip_len;

        // 11 STORETIMESTAMP refresh store time stamp in lock
        pre_encode_buffer[pos..(pos + 8)]
            .copy_from_slice(&msg_inner.store_timestamp().to_be_bytes());

        if self.message_store_config.enabled_append_prop_crc {
            // 18 CRC32
            let check_size = msg_len - self.crc32_reserved_length;
            let crc32 = crc32(&pre_encode_buffer[..check_size as usize]);
            create_crc32(
                &mut pre_encode_buffer[check_size as usize..msg_len as usize],
                crc32,
            );
        }

        //let bytes = pre_encode_buffer.freeze();
        let instant = Instant::now();
        mapped_file.append_message_bytes_no_position_update_ref(pre_encode_buffer.chunk());
        AppendMessageResult {
            status: AppendMessageStatus::PutOk,
            wrote_offset,
            wrote_bytes: msg_len,
            store_timestamp: msg_inner.store_timestamp(),
            logics_offset: queue_offset,
            msg_num: message_num as i32,
            msg_id_supplier: Some(Arc::new(Box::new(msg_id_supplier))),
            page_cache_rt: instant.elapsed().as_millis() as i64,
            ..Default::default()
        }
    }

    fn do_append_batch<MF: MappedFile>(
        &self,
        file_from_offset: i64,
        mapped_file: &MF,
        max_blank: i32,
        msg_batch: &mut MessageExtBatch,
        put_message_context: &mut PutMessageContext,
        enabled_append_prop_crc: bool,
    ) -> AppendMessageResult {
        //physic offset--The starting point for writing this message file.If, while writing a
        // message, it is found that the length is insufficient,the remaining length of the file
        // and the end-of-file marker should be rewritten at this point.
        let wrote_offset = file_from_offset + mapped_file.get_wrote_position() as i64;
        // Record ConsumeQueue information
        let queue_offset = msg_batch.message_ext_broker_inner.queue_offset();
        let begin_queue_offset = queue_offset;

        let begin_time_mills = Instant::now();

        // Assuming get_encoded_buff returns Option<ByteBuffer>
        let mut messages_byte_buffer = msg_batch.encoded_buff.take().unwrap();
        let sys_flag = msg_batch.message_ext_broker_inner.sys_flag();
        //born host length
        let born_host_length = if sys_flag & MessageSysFlag::BORNHOST_V6_FLAG == 0 {
            4 + 4
        } else {
            16 + 4
        };
        //store host length
        let store_host_length = if sys_flag & MessageSysFlag::STOREHOSTADDRESS_V6_FLAG == 0 {
            4 + 4
        } else {
            16 + 4
        };
        let addr = msg_batch.message_ext_broker_inner.store_host();
        let batch_size = put_message_context.get_batch_size();
        let phy_ops = put_message_context.get_phy_pos().to_vec();
        let msg_id_supplier = move || -> String {
            build_batch_message_id(addr, store_host_length, batch_size as usize, &phy_ops)
        };
        let mut total_msg_len = 0;
        let mut msg_num = 0;
        let mut msg_pos = 0;
        let mut index = 0;
        while total_msg_len < messages_byte_buffer.len() as i32 {
            let msg_len = i32::from_be_bytes(
                messages_byte_buffer[total_msg_len as usize..(total_msg_len + 4) as usize]
                    .try_into()
                    .unwrap(),
            );
            total_msg_len += msg_len;
            if total_msg_len + END_FILE_MIN_BLANK_LENGTH > max_blank {
                let bytes = self.msg_store_item_memory.mut_from_ref();
                bytes.clear();
                bytes.put_i32(max_blank);
                bytes.put_i32(BLANK_MAGIC_CODE);
                mapped_file.write_bytes_segment(
                    bytes.as_ref(),
                    wrote_offset as usize,
                    0,
                    bytes.len(),
                );
                return AppendMessageResult {
                    status: AppendMessageStatus::EndOfFile,
                    wrote_offset,
                    wrote_bytes: max_blank,
                    msg_id_supplier: Some(Arc::new(Box::new(msg_id_supplier))),
                    store_timestamp: msg_batch.message_ext_broker_inner.store_timestamp(),
                    logics_offset: begin_queue_offset,
                    page_cache_rt: begin_time_mills.elapsed().as_millis() as i64,
                    ..Default::default()
                };
            }
            let mut pos = msg_pos + 20;
            messages_byte_buffer[pos..(pos + 8)].copy_from_slice(&queue_offset.to_be_bytes());
            pos += 8;
            let phy_pos = wrote_offset + total_msg_len as i64 - msg_len as i64;
            messages_byte_buffer[pos..(pos + 8)].copy_from_slice(&phy_pos.to_be_bytes());
            pos += 8 + 4 + 8 + born_host_length;
            messages_byte_buffer[pos..(pos + 8)].copy_from_slice(
                &msg_batch
                    .message_ext_broker_inner
                    .store_timestamp()
                    .to_be_bytes(),
            );
            if enabled_append_prop_crc {
                let _check_size = msg_len - self.crc32_reserved_length;
            }
            put_message_context.get_phy_pos_mut()[index] = phy_pos;
            msg_num += 1;
            msg_pos += msg_len as usize;
            index += 1;
        }

        let bytes = messages_byte_buffer.freeze();
        mapped_file.append_message_bytes_no_position_update(&bytes);
        AppendMessageResult {
            status: AppendMessageStatus::PutOk,
            wrote_offset,
            wrote_bytes: total_msg_len,
            msg_id_supplier: Some(Arc::new(Box::new(msg_id_supplier))),
            store_timestamp: msg_batch.message_ext_broker_inner.store_timestamp(),
            logics_offset: begin_queue_offset,
            page_cache_rt: begin_time_mills.elapsed().as_millis() as i64,
            msg_num,
            ..Default::default()
        }
    }
}