io_engine/

merge.rs

use crate::embedded_list::*;
use crate::tasks::*;
use crossfire::BlockingTxTrait;
use io_buffer::Buffer;
use std::io;
use std::mem::transmute;
use std::os::fd::RawFd;

pub struct EventMergeBuffer<C: IoCallback> {
    pub merge_size_limit: usize,
    last_event: Option<Box<IOEvent<C>>>,
    merged_events: Option<EmbeddedList>,
    merged_offset: i64,
    merged_data_size: usize,
    merged_count: usize,
    list_node_offset: usize,
}

impl<C: IoCallback> EventMergeBuffer<C> {
    pub fn new(merge_size_limit: usize) -> Self {
        Self {
            merge_size_limit,
            last_event: None,
            merged_count: 0,
            merged_events: None,
            merged_offset: -1,
            merged_data_size: 0,
            list_node_offset: std::mem::offset_of!(IOEvent<C>, node),
        }
    }

    #[inline(always)]
    pub fn may_add_event(&mut self, event: &Box<IOEvent<C>>) -> bool {
        if self.merged_count > 0 {
            if self.merged_data_size + event.get_size() > self.merge_size_limit {
                return false;
            }
            return self.merged_offset + self.merged_data_size as i64 == event.offset;
        } else {
            return true;
        }
    }

    // return full
    #[inline(always)]
    pub fn push_event(&mut self, event: Box<IOEvent<C>>) -> bool {
        if self.merged_count == 0 {
            self.merged_offset = event.offset;
            self.merged_data_size = event.get_size();
            self.last_event = Some(event);
            self.merged_count = 1;
            return false;
        } else {
            self.merged_data_size += event.get_size();
            if self.merged_count == 1 {
                let last_event = self.last_event.take().unwrap();
                let mut list = EmbeddedList::new(self.list_node_offset);
                last_event.push_to_list(&mut list);
                event.push_to_list(&mut list);
                self.merged_events = Some(list);
            } else {
                let merged_events = self.merged_events.as_mut().unwrap();
                event.push_to_list(merged_events);
            }
            self.merged_count += 1;
            return self.merged_data_size >= self.merge_size_limit;
        }
    }

    #[inline(always)]
    pub fn len(&self) -> usize {
        self.merged_count
    }

    #[inline(always)]
    pub fn take(&mut self) -> (EmbeddedList, i64, usize) {
        log_debug_assert!(self.merged_count > 1, "merged_count {}", self.merged_count);
        let tasks = self.merged_events.take().unwrap();
        let merged_data_size = self.merged_data_size;
        let merged_offset = self.merged_offset;
        self.merged_offset = -1;
        self.merged_data_size = 0;
        self.merged_count = 0;
        (tasks, merged_offset, merged_data_size)
    }

    #[inline(always)]
    pub fn take_one(&mut self) -> Box<IOEvent<C>> {
        log_debug_assert_eq!(self.merged_count, 1);
        self.merged_offset = -1;
        self.merged_data_size = 0;
        self.merged_count = 0;
        self.last_event.take().unwrap()
    }
}

/// Try to merge sequential IOEvent.
///
/// NOTE: Assuming all the IOEvents are of the same file, only debug mode will do validation.
pub struct IOMergeSubmitter<C: IoCallback, S: BlockingTxTrait<Box<IOEvent<C>>>> {
    fd: RawFd,
    buffer: EventMergeBuffer<C>,
    sender: S,
    action: IOAction,
}

impl<C: IoCallback, S: BlockingTxTrait<Box<IOEvent<C>>>> IOMergeSubmitter<C, S> {
    pub fn new(fd: RawFd, sender: S, merge_size_limit: usize, action: IOAction) -> Self {
        log_assert!(merge_size_limit > 0);
        Self { fd, buffer: EventMergeBuffer::new(merge_size_limit), sender, action }
    }

    /// On debug mode, will validate event.fd and event.action.
    pub fn add_event(&mut self, event: Box<IOEvent<C>>) -> Result<(), io::Error> {
        log_debug_assert_eq!(self.fd, event.fd);
        log_debug_assert_eq!(event.action, self.action);
        let event_size = event.get_size();

        if event_size >= self.buffer.merge_size_limit || !self.buffer.may_add_event(&event) {
            if let Err(e) = self._flush() {
                event.callback();
                return Err(e);
            }
        }
        if self.buffer.push_event(event) {
            self._flush()?;
        }
        return Ok(());
    }

    pub fn flush(&mut self) -> Result<(), io::Error> {
        self._flush()
    }

    #[inline(always)]
    fn _flush(&mut self) -> Result<(), io::Error> {
        let batch_len = self.buffer.len();
        if batch_len == 0 {
            return Ok(());
        }
        if batch_len == 1 {
            let submit_event = self.buffer.take_one();
            trace!("mio: submit {:?} not merged", submit_event);
            self.sender
                .send(submit_event)
                .map_err(|_| io::Error::new(io::ErrorKind::Other, "Queue closed"))?;
        } else {
            let (mut events, offset, size) = self.buffer.take();
            log_assert!(size > 0);
            match Buffer::aligned(size as i32) {
                Ok(mut buffer) => {
                    trace!("mio: merged {} ev into {} @{}", events.get_length(), size, offset);
                    if self.action == IOAction::Write {
                        let mut _offset = 0;
                        for _event in events.iter::<IOEvent<C>>() {
                            let event: &IOEvent<C> = unsafe { transmute(_event) };
                            let b = event.get_buf_ref();
                            let _size = b.len();
                            buffer.copy_from(_offset, b);
                            _offset += _size;
                        }
                    }
                    let mut event = IOEvent::<C>::new(self.fd, buffer, self.action, offset);
                    event.set_subtasks(events);
                    self.sender
                        .send(event)
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, "Queue closed"))?;
                }
                Err(_) => {
                    // commit separately
                    warn!("mio: alloc buffer size {} failed", size);
                    let mut e: Option<io::Error> = None;
                    while let Some(event) = IOEvent::<C>::pop_from_list(&mut events) {
                        if let Err(_e) = self
                            .sender
                            .send(event)
                            .map_err(|_| io::Error::new(io::ErrorKind::Other, "Queue closed"))
                        {
                            e.replace(_e);
                        }
                    }
                    if let Some(_e) = e {
                        return Err(_e);
                    }
                }
            }
        }
        Ok(())
    }
}