You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iggy.apache.org by hg...@apache.org on 2025/07/07 19:41:12 UTC
(iggy) branch io_uring_tpc_direct_io updated (3f4bea29 -> 4c0267f5)
This is an automated email from the ASF dual-hosted git repository.
hgruszecki pushed a change to branch io_uring_tpc_direct_io
in repository https://gitbox.apache.org/repos/asf/iggy.git
discard 3f4bea29 feat(io_uring): implement directio and aligned PooledBuffer
new 4c0267f5 feat(io_uring): implement directio and aligned PooledBuffer
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (3f4bea29)
\
N -- N -- N refs/heads/io_uring_tpc_direct_io (4c0267f5)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
core/server/src/streaming/segments/segment.rs | 23 ++---------------------
core/server/src/streaming/utils/memory_pool.rs | 6 ------
core/server/src/streaming/utils/pooled_buffer.rs | 24 ------------------------
3 files changed, 2 insertions(+), 51 deletions(-)
(iggy) 01/01: feat(io_uring): implement directio and aligned PooledBuffer
Posted by hg...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
hgruszecki pushed a commit to branch io_uring_tpc_direct_io
in repository https://gitbox.apache.org/repos/asf/iggy.git
commit 4c0267f57a875228477b4133c3e1793b6c5e40d4
Author: Hubert Gruszecki <h....@gmail.com>
AuthorDate: Mon Jul 7 21:33:25 2025 +0200
feat(io_uring): implement directio and aligned PooledBuffer
---
Cargo.lock | 30 ++
core/server/Cargo.toml | 1 +
.../handlers/messages/send_messages_handler.rs | 42 +-
core/server/src/binary/handlers/utils.rs | 3 +-
core/server/src/binary/sender.rs | 7 +-
core/server/src/main.rs | 11 +
core/server/src/quic/quic_sender.rs | 2 +-
core/server/src/shard/mod.rs | 28 +-
core/server/src/shard/system/storage.rs | 2 +-
core/server/src/streaming/segments/direct_file.rs | 459 +++++++++++++++++++++
.../src/streaming/segments/indexes/index_reader.rs | 23 +-
.../src/streaming/segments/indexes/index_writer.rs | 97 +++--
.../streaming/segments/messages/messages_reader.rs | 22 +-
.../streaming/segments/messages/messages_writer.rs | 106 +++--
core/server/src/streaming/segments/messages/mod.rs | 28 +-
core/server/src/streaming/segments/mod.rs | 2 +
.../src/streaming/segments/reading_messages.rs | 24 +-
core/server/src/streaming/segments/segment.rs | 58 ++-
.../streaming/segments/types/messages_batch_mut.rs | 15 +-
.../src/streaming/segments/writing_messages.rs | 19 +-
core/server/src/streaming/utils/memory_pool.rs | 64 +--
core/server/src/streaming/utils/mod.rs | 2 +-
core/server/src/streaming/utils/pooled_buffer.rs | 152 ++++---
core/server/src/tcp/connection_handler.rs | 13 +-
core/server/src/tcp/sender.rs | 18 +-
core/server/src/tcp/tcp_sender.rs | 6 +-
core/server/src/tcp/tcp_tls_sender.rs | 2 +-
27 files changed, 867 insertions(+), 369 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 540217c4..20d0a61b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -307,6 +307,15 @@ dependencies = [
"memchr",
]
+[[package]]
+name = "aligned-vec"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b"
+dependencies = [
+ "equator",
+]
+
[[package]]
name = "alloc-no-stdlib"
version = "2.0.4"
@@ -2461,6 +2470,26 @@ dependencies = [
"log",
]
+[[package]]
+name = "equator"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc"
+dependencies = [
+ "equator-macro",
+]
+
+[[package]]
+name = "equator-macro"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.104",
+]
+
[[package]]
name = "equivalent"
version = "1.0.2"
@@ -7071,6 +7100,7 @@ name = "server"
version = "0.5.0"
dependencies = [
"ahash 0.8.12",
+ "aligned-vec",
"anyhow",
"async-channel",
"async_zip",
diff --git a/core/server/Cargo.toml b/core/server/Cargo.toml
index 47738185..1ce3aa05 100644
--- a/core/server/Cargo.toml
+++ b/core/server/Cargo.toml
@@ -36,6 +36,7 @@ mimalloc = ["dep:mimalloc"]
[dependencies]
ahash = { workspace = true }
+aligned-vec = "0.6"
anyhow = { workspace = true }
async_zip = { workspace = true }
axum = { workspace = true }
diff --git a/core/server/src/binary/handlers/messages/send_messages_handler.rs b/core/server/src/binary/handlers/messages/send_messages_handler.rs
index bb536c61..6d155bca 100644
--- a/core/server/src/binary/handlers/messages/send_messages_handler.rs
+++ b/core/server/src/binary/handlers/messages/send_messages_handler.rs
@@ -57,18 +57,19 @@ impl ServerCommandHandler for SendMessages {
let total_payload_size = length as usize - std::mem::size_of::<u32>();
let metadata_len_field_size = std::mem::size_of::<u32>();
- let metadata_length_buffer = PooledBuffer::with_capacity(4);
- let (result, metadata_len_buf) = sender.read(metadata_length_buffer.slice(0..4)).await;
- let metadata_len_buf = metadata_len_buf.into_inner();
- result?;
- let metadata_size = u32::from_le_bytes(metadata_len_buf[..].try_into().unwrap());
-
- let metadata_buffer = PooledBuffer::with_capacity(metadata_size as usize);
- let (result, metadata_buf) = sender
+ let mut metadata_length_buffer = PooledBuffer::with_capacity(4);
+
+ let metadata_length_buffer = sender
+ .read(metadata_length_buffer.slice(0..4))
+ .await?
+ .into_inner();
+ let metadata_size = u32::from_le_bytes(metadata_length_buffer[0..4].try_into().unwrap());
+
+ let mut metadata_buffer = PooledBuffer::with_capacity(metadata_size as usize);
+ let metadata_buf = sender
.read(metadata_buffer.slice(0..metadata_size as usize))
- .await;
- result?;
- let metadata_buf = metadata_buf.into_inner();
+ .await?
+ .into_inner();
let mut element_size = 0;
@@ -91,17 +92,20 @@ impl ServerCommandHandler for SendMessages {
);
let indexes_size = messages_count as usize * INDEX_SIZE;
- let indexes_buffer = PooledBuffer::with_capacity(indexes_size);
- let (result, indexes_buffer) = sender.read(indexes_buffer.slice(0..indexes_size)).await;
- result?;
- let indexes_buffer = indexes_buffer.into_inner();
+ let mut indexes_buffer = PooledBuffer::with_capacity(indexes_size + 512); // extra space for possible padding to not cause reallocations
+ let indexes_buffer = sender
+ .read(indexes_buffer.slice(0..indexes_size))
+ .await?
+ .into_inner();
let messages_size =
total_payload_size - metadata_size as usize - indexes_size - metadata_len_field_size;
- let messages_buffer = PooledBuffer::with_capacity(messages_size);
- let (result, messages_buffer) = sender.read(messages_buffer.slice(0..messages_size)).await;
- result?;
- let messages_buffer = messages_buffer.into_inner();
+
+ let mut messages_buffer = PooledBuffer::with_capacity(messages_size + 512); // extra space for possible padding to not cause reallocations
+ let messages_buffer = sender
+ .read(messages_buffer.slice(0..messages_size))
+ .await?
+ .into_inner();
let indexes = IggyIndexesMut::from_bytes(indexes_buffer, 0);
let batch = IggyMessagesBatchMut::from_indexes_and_messages(
diff --git a/core/server/src/binary/handlers/utils.rs b/core/server/src/binary/handlers/utils.rs
index 80bcede4..54411d91 100644
--- a/core/server/src/binary/handlers/utils.rs
+++ b/core/server/src/binary/handlers/utils.rs
@@ -32,8 +32,7 @@ pub async fn receive_and_validate(
let buffer = if length == 0 {
buffer
} else {
- let (result, buffer) = sender.read(buffer).await;
- result?;
+ let buffer = sender.read(buffer).await?;
buffer
};
diff --git a/core/server/src/binary/sender.rs b/core/server/src/binary/sender.rs
index 96bed3d5..c6f1a811 100644
--- a/core/server/src/binary/sender.rs
+++ b/core/server/src/binary/sender.rs
@@ -56,10 +56,7 @@ macro_rules! forward_async_methods {
}
pub trait Sender {
- fn read<B: IoBufMut>(
- &mut self,
- buffer: B,
- ) -> impl Future<Output = (Result<usize, IggyError>, B)>;
+ fn read<B: IoBufMut>(&mut self, buffer: B) -> impl Future<Output = Result<B, IggyError>>;
fn send_empty_ok_response(&mut self) -> impl Future<Output = Result<(), IggyError>>;
fn send_ok_response(&mut self, payload: &[u8]) -> impl Future<Output = Result<(), IggyError>>;
fn send_ok_response_vectored(
@@ -98,7 +95,7 @@ impl SenderKind {
}
forward_async_methods! {
- async fn read<B: IoBufMut>(&mut self, buffer: B) -> (Result<usize, IggyError>, B);
+ async fn read<B: IoBufMut>(&mut self, buffer: B) -> Result<B, IggyError>;
async fn send_empty_ok_response(&mut self) -> Result<(), IggyError>;
async fn send_ok_response(&mut self, payload: &[u8]) -> Result<(), IggyError>;
async fn send_ok_response_vectored(&mut self, length: &[u8], slices: Vec<PooledBuffer>) -> Result<(), IggyError>;
diff --git a/core/server/src/main.rs b/core/server/src/main.rs
index 2602e3c8..89fa68e7 100644
--- a/core/server/src/main.rs
+++ b/core/server/src/main.rs
@@ -57,6 +57,15 @@ use tracing::{error, info, instrument};
const COMPONENT: &str = "MAIN";
+fn thread_print_memory_pool_stats() {
+ std::thread::sleep(std::time::Duration::from_secs(5));
+ let pool = server::streaming::utils::memory_pool();
+ loop {
+ pool.log_stats();
+ std::thread::sleep(std::time::Duration::from_secs(5));
+ }
+}
+
#[instrument(skip_all, name = "trace_start_server")]
fn main() -> Result<(), ServerError> {
let startup_timestamp = Instant::now();
@@ -269,6 +278,8 @@ fn main() -> Result<(), ServerError> {
.expect("Error setting Ctrl-C handler");
*/
+ std::thread::spawn(thread_print_memory_pool_stats);
+
info!("Iggy server is running. Press Ctrl+C or send SIGTERM to shutdown.");
for (idx, handle) in handles.into_iter().enumerate() {
info!("Waiting for shard thread {} to complete...", idx);
diff --git a/core/server/src/quic/quic_sender.rs b/core/server/src/quic/quic_sender.rs
index 54bda604..30c4a6a7 100644
--- a/core/server/src/quic/quic_sender.rs
+++ b/core/server/src/quic/quic_sender.rs
@@ -37,7 +37,7 @@ pub struct QuicSender {
}
impl Sender for QuicSender {
- async fn read<B: IoBufMut>(&mut self, buffer: B) -> (Result<usize, IggyError>, B) {
+ async fn read<B: IoBufMut>(&mut self, buffer: B) -> Result<B, IggyError> {
//TODO: Fixme
// Not-so-nice code because quinn recv stream has different API for read_exact
/*
diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs
index 3d467779..078c9018 100644
--- a/core/server/src/shard/mod.rs
+++ b/core/server/src/shard/mod.rs
@@ -221,16 +221,6 @@ impl IggyShard {
}
async fn load_version(&self) -> Result<(), IggyError> {
- async fn update_system_info(
- storage: &Rc<SystemStorage>,
- system_info: &mut SystemInfo,
- version: &SemanticVersion,
- ) -> Result<(), IggyError> {
- system_info.update_version(version);
- storage.info.save(system_info).await?;
- Ok(())
- }
-
let current_version = &self.version;
let mut system_info;
let load_system_info = self.storage.info.load().await;
@@ -239,7 +229,7 @@ impl IggyShard {
if let IggyError::ResourceNotFound(_) = error {
info!("System info not found, creating...");
system_info = SystemInfo::default();
- update_system_info(&self.storage, &mut system_info, current_version).await?;
+ Self::update_system_info(&self.storage, &mut system_info, current_version).await?;
} else {
return Err(error);
}
@@ -248,24 +238,34 @@ impl IggyShard {
}
info!("Loaded {system_info}.");
- let loaded_version = SemanticVersion::from_str(&system_info.version.version)?;
+ let loaded_version = SemanticVersion::from_str(&system_info.version.version).unwrap();
if current_version.is_equal_to(&loaded_version) {
info!("System version {current_version} is up to date.");
} else if current_version.is_greater_than(&loaded_version) {
info!(
"System version {current_version} is greater than {loaded_version}, checking the available migrations..."
);
- update_system_info(&self.storage, &mut system_info, current_version).await?;
+ Self::update_system_info(&self.storage, &mut system_info, current_version).await?;
} else {
info!(
"System version {current_version} is lower than {loaded_version}, possible downgrade."
);
- update_system_info(&self.storage, &mut system_info, current_version).await?;
+ Self::update_system_info(&self.storage, &mut system_info, current_version).await?;
}
Ok(())
}
+ async fn update_system_info(
+ storage: &Rc<SystemStorage>,
+ system_info: &mut SystemInfo,
+ version: &SemanticVersion,
+ ) -> Result<(), IggyError> {
+ system_info.update_version(version);
+ storage.info.save(system_info).await?;
+ Ok(())
+ }
+
async fn load_state(&self) -> Result<SystemState, IggyError> {
let state_entries = self.state.init().await.with_error_context(|error| {
format!("{COMPONENT} (error: {error}) - failed to initialize state entries")
diff --git a/core/server/src/shard/system/storage.rs b/core/server/src/shard/system/storage.rs
index 736d8e48..ac4aca3e 100644
--- a/core/server/src/shard/system/storage.rs
+++ b/core/server/src/shard/system/storage.rs
@@ -65,7 +65,7 @@ impl SystemInfoStorage for FileSystemInfoStorage {
let file = file::open(&self.path)
.await
.map_err(|_| IggyError::CannotReadFile)?;
- let buffer = PooledBuffer::with_capacity(file_size);
+ let mut buffer = PooledBuffer::with_capacity(file_size);
let (result, buffer) = file
.read_exact_at(buffer.slice(0..file_size), 0)
.await
diff --git a/core/server/src/streaming/segments/direct_file.rs b/core/server/src/streaming/segments/direct_file.rs
new file mode 100644
index 00000000..a0f541dc
--- /dev/null
+++ b/core/server/src/streaming/segments/direct_file.rs
@@ -0,0 +1,459 @@
+/* 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 crate::streaming::utils::{ALIGNMENT, PooledBuffer};
+use compio::fs::{File, OpenOptions};
+use compio::io::AsyncWriteAtExt;
+use error_set::ErrContext;
+use iggy_common::IggyError;
+
+#[derive(Debug)]
+pub struct DirectFile {
+ file_path: String,
+ file: File,
+ file_position: u64,
+ tail: PooledBuffer,
+ tail_len: usize,
+}
+
+impl DirectFile {
+ pub async fn open(
+ file_path: &str,
+ initial_position: u64,
+ file_exists: bool,
+ ) -> Result<Self, IggyError> {
+ let mut file = OpenOptions::new()
+ .create(true)
+ .write(true)
+ .custom_flags(0x4000)
+ .open(file_path)
+ .await
+ .with_error_context(|err| {
+ format!("Failed to open file with O_DIRECT: {file_path}, error: {err}")
+ })
+ .map_err(|_| IggyError::CannotReadFile)?;
+
+ if !file_exists {
+ let init_buffer = PooledBuffer::with_capacity(ALIGNMENT);
+ let (write_result, _) = file.write_all_at(init_buffer, 0).await.into();
+ write_result
+ .with_error_context(|error| {
+ tracing::error!(
+ "Failed to initialize file with dummy block: {file_path}, error: {error}"
+ );
+ format!(
+ "Failed to initialize file with dummy block: {file_path}, error: {error}"
+ )
+ })
+ .map_err(|_| IggyError::CannotWriteToFile)?;
+
+ tracing::trace!("Successfully initialized new file with dummy block: {file_path}");
+ }
+
+ tracing::trace!(
+ "Successfully opened DirectFile: {}, position: {}, exists: {}",
+ file_path,
+ initial_position,
+ file_exists
+ );
+
+ Ok(Self {
+ file_path: file_path.to_string(),
+ file,
+ file_position: initial_position,
+ tail: PooledBuffer::with_capacity(ALIGNMENT),
+ tail_len: 0,
+ })
+ }
+
+ pub async fn get_file_size(&self) -> Result<u64, IggyError> {
+ self.file
+ .metadata()
+ .await
+ .with_error_context(|error| {
+ format!(
+ "Failed to get metadata of file: {}, error: {error}",
+ self.file_path
+ )
+ })
+ .map_err(|_| IggyError::CannotReadFileMetadata)
+ .map(|metadata| metadata.len())
+ }
+
+ fn new(file: File, file_path: String, initial_position: u64) -> Self {
+ Self {
+ file_path,
+ file,
+ file_position: initial_position,
+ tail: PooledBuffer::with_capacity(ALIGNMENT),
+ tail_len: 0,
+ }
+ }
+
+ pub async fn write_all(&mut self, mut data: &[u8]) -> Result<usize, IggyError> {
+ let initial_len = data.len();
+ tracing::trace!(
+ "DirectFile write_all called for file: {}, data_len: {}, position: {}, tail_len: {}",
+ self.file_path,
+ initial_len,
+ self.file_position,
+ self.tail_len
+ );
+
+ if self.tail_len > 0 {
+ let need = ALIGNMENT - self.tail_len;
+ let take = need.min(data.len());
+ self.tail.extend_from_slice(&data[..take]);
+ self.tail_len += take;
+ data = &data[take..];
+
+ if self.tail_len == ALIGNMENT {
+ self.flush_tail().await?;
+ }
+ }
+
+ if !data.is_empty() {
+ let whole_sectors_end = data.len() & !(ALIGNMENT - 1);
+ if whole_sectors_end > 0 {
+ let whole_sectors = &data[..whole_sectors_end];
+ let mut written = 0;
+
+ while written < whole_sectors.len() {
+ let chunk_size = (whole_sectors.len() - written).min(128 * 1024 * 1024);
+ let chunk = &whole_sectors[written..written + chunk_size];
+
+ let chunk_buffer = PooledBuffer::from(chunk);
+
+ let (result, _) = self
+ .file
+ .write_all_at(chunk_buffer, self.file_position)
+ .await
+ .into();
+
+ result.map_err(|e| {
+ tracing::error!("Failed to write to direct file: {} at position {}, chunk size: {}, error: {}",
+ self.file_path, self.file_position, chunk_size, e);
+ IggyError::CannotWriteToFile
+ })?;
+
+ self.file_position += chunk_size as u64;
+ written += chunk_size;
+ }
+
+ data = &data[whole_sectors_end..];
+ }
+ }
+
+ if !data.is_empty() {
+ self.tail.clear();
+ self.tail.extend_from_slice(data);
+ self.tail_len = data.len();
+ }
+
+ Ok(initial_len)
+ }
+
+ pub async fn flush(&mut self) -> Result<(), IggyError> {
+ if self.tail_len > 0 {
+ self.tail.resize(ALIGNMENT, 0);
+ self.flush_tail().await?;
+ }
+ Ok(())
+ }
+
+ pub async fn close(mut self) -> Result<(), IggyError> {
+ if self.tail_len > 0 {
+ self.tail.resize(ALIGNMENT, 0);
+ self.flush_tail().await?;
+ }
+ Ok(())
+ }
+
+ pub async fn sync_all(&self) -> Result<(), IggyError> {
+ Ok(())
+ }
+
+ pub fn position(&self) -> u64 {
+ self.file_position
+ }
+
+ pub fn tail_len(&self) -> usize {
+ self.tail_len
+ }
+
+ pub fn file_path(&self) -> &str {
+ &self.file_path
+ }
+
+ pub fn tail_buffer(&self) -> &PooledBuffer {
+ &self.tail
+ }
+
+ pub fn take_tail(&mut self) -> (PooledBuffer, usize) {
+ let tail = std::mem::replace(&mut self.tail, PooledBuffer::with_capacity(ALIGNMENT));
+ let tail_len = self.tail_len;
+ self.tail_len = 0;
+ (tail, tail_len)
+ }
+
+ pub fn set_tail(&mut self, tail: PooledBuffer, tail_len: usize) {
+ self.tail = tail;
+ self.tail_len = tail_len;
+ }
+
+ async fn flush_tail(&mut self) -> Result<(), IggyError> {
+ assert_eq!(self.tail.len(), ALIGNMENT);
+
+ let tail_buffer = std::mem::replace(&mut self.tail, PooledBuffer::with_capacity(ALIGNMENT));
+
+ let (result, returned_buf) = self
+ .file
+ .write_all_at(tail_buffer, self.file_position)
+ .await
+ .into();
+
+ result.map_err(|e| {
+ tracing::error!(
+ "Failed to flush tail for file: {} at position {}, tail size: {}, error: {}",
+ self.file_path,
+ self.file_position,
+ ALIGNMENT,
+ e
+ );
+ IggyError::CannotWriteToFile
+ })?;
+
+ self.file_position += ALIGNMENT as u64;
+ self.tail_len = 0;
+ self.tail = returned_buf;
+ self.tail.clear();
+ Ok(())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::configs::system::SystemConfig;
+ use crate::streaming::utils::MemoryPool;
+ use compio::fs::OpenOptions;
+ use compio::io::{AsyncReadAt, AsyncReadAtExt};
+ use std::sync::Arc;
+ use tempfile::tempdir;
+
+ #[test]
+ fn test_direct_file_small_writes() {
+ compio::runtime::Runtime::new().unwrap().block_on(async {
+ MemoryPool::init_pool(Arc::new(SystemConfig::default()));
+ let temp_dir = tempdir().unwrap();
+ let file_path = temp_dir.path().join("test_direct_io.bin");
+
+ let mut direct_file = DirectFile::open(file_path.to_str().unwrap(), 0, false)
+ .await
+ .unwrap();
+
+ for i in 0..10u64 {
+ let buf = i.to_le_bytes();
+ direct_file.write_all(&buf).await.unwrap();
+ }
+
+ direct_file.flush().await.unwrap();
+
+ let file = OpenOptions::new()
+ .read(true)
+ .custom_flags(0x4000)
+ .open(&file_path)
+ .await
+ .unwrap();
+
+ let mut read_buffer = vec![0u8; 512];
+ let (result, buf) = file.read_at(read_buffer, 0).await.into();
+ result.unwrap();
+ read_buffer = buf;
+
+ for i in 0..10u64 {
+ let start = i as usize * 8;
+ let num = u64::from_le_bytes(read_buffer[start..start + 8].try_into().unwrap());
+ assert_eq!(num, i, "Expected {} at position {}, got {}", i, start, num);
+ }
+ });
+ }
+
+ #[test]
+ fn test_direct_file_exact_sector_write() {
+ compio::runtime::Runtime::new().unwrap().block_on(async {
+ MemoryPool::init_pool(Arc::new(SystemConfig::default()));
+ let temp_dir = tempdir().unwrap();
+ let file_path = temp_dir.path().join("test_exact_sector.bin");
+
+ let mut direct_file = DirectFile::open(file_path.to_str().unwrap(), 0, false)
+ .await
+ .unwrap();
+
+ let data = vec![42u8; ALIGNMENT];
+ direct_file.write_all(&data).await.unwrap();
+
+ assert_eq!(direct_file.tail_len(), 0);
+ assert_eq!(direct_file.position(), ALIGNMENT as u64);
+
+ direct_file.close().await.unwrap();
+
+ let file = OpenOptions::new()
+ .read(true)
+ .custom_flags(0x4000)
+ .open(&file_path)
+ .await
+ .unwrap();
+
+ let mut read_buffer = vec![0u8; ALIGNMENT];
+ let (result, buf) = file.read_at(read_buffer, 0).await.into();
+ result.unwrap();
+ read_buffer = buf;
+
+ assert_eq!(read_buffer, vec![42u8; ALIGNMENT]);
+ });
+ }
+
+ #[test]
+ fn test_direct_file_multiple_sector_writes() {
+ compio::runtime::Runtime::new().unwrap().block_on(async {
+ MemoryPool::init_pool(Arc::new(SystemConfig::default()));
+ let temp_dir = tempdir().unwrap();
+ let file_path = temp_dir.path().join("test_multiple_sectors.bin");
+
+ let mut direct_file = DirectFile::open(file_path.to_str().unwrap(), 0, false)
+ .await
+ .unwrap();
+
+ let data1 = vec![1u8; ALIGNMENT * 2];
+ let data2 = vec![2u8; ALIGNMENT * 3];
+ let data3 = vec![3u8; ALIGNMENT];
+
+ direct_file.write_all(&data1).await.unwrap();
+ direct_file.write_all(&data2).await.unwrap();
+ direct_file.write_all(&data3).await.unwrap();
+
+ direct_file.close().await.unwrap();
+
+ let file = OpenOptions::new()
+ .read(true)
+ .custom_flags(0x4000)
+ .open(&file_path)
+ .await
+ .unwrap();
+
+ let mut read_buffer = vec![0u8; ALIGNMENT * 6];
+ let (result, buf) = file.read_at(read_buffer, 0).await.into();
+ result.unwrap();
+ read_buffer = buf;
+
+ assert_eq!(&read_buffer[0..ALIGNMENT * 2], &vec![1u8; ALIGNMENT * 2]);
+ assert_eq!(
+ &read_buffer[ALIGNMENT * 2..ALIGNMENT * 5],
+ &vec![2u8; ALIGNMENT * 3]
+ );
+ assert_eq!(
+ &read_buffer[ALIGNMENT * 5..ALIGNMENT * 6],
+ &vec![3u8; ALIGNMENT]
+ );
+ });
+ }
+
+ #[test]
+ fn test_direct_file_unaligned_write() {
+ compio::runtime::Runtime::new().unwrap().block_on(async {
+ MemoryPool::init_pool(Arc::new(SystemConfig::default()));
+ let temp_dir = tempdir().unwrap();
+ let file_path = temp_dir.path().join("test_unaligned.bin");
+
+ let mut direct_file = DirectFile::open(file_path.to_str().unwrap(), 0, false)
+ .await
+ .unwrap();
+
+ let data = vec![77u8; 1000];
+ direct_file.write_all(&data).await.unwrap();
+
+ assert_eq!(direct_file.tail_len(), 1000 % ALIGNMENT);
+ assert_eq!(direct_file.position(), ALIGNMENT as u64);
+
+ direct_file.flush().await.unwrap();
+
+ assert_eq!(direct_file.tail_len(), 0);
+ assert_eq!(direct_file.position(), (ALIGNMENT * 2) as u64);
+
+ let file = OpenOptions::new()
+ .read(true)
+ .custom_flags(0x4000)
+ .open(&file_path)
+ .await
+ .unwrap();
+
+ let mut read_buffer = vec![0u8; ALIGNMENT * 2];
+ let (result, buf) = file.read_at(read_buffer, 0).await.into();
+ result.unwrap();
+ read_buffer = buf;
+
+ assert_eq!(&read_buffer[0..1000], &vec![77u8; 1000]);
+ assert_eq!(
+ &read_buffer[1000..ALIGNMENT * 2],
+ &vec![0u8; ALIGNMENT * 2 - 1000]
+ );
+ });
+ }
+
+ #[test]
+ fn test_direct_file_cross_sector_boundary() {
+ compio::runtime::Runtime::new().unwrap().block_on(async {
+ MemoryPool::init_pool(Arc::new(SystemConfig::default()));
+ let temp_dir = tempdir().unwrap();
+ let file_path = temp_dir.path().join("test_cross_boundary.bin");
+
+ let mut direct_file = DirectFile::open(file_path.to_str().unwrap(), 0, false)
+ .await
+ .unwrap();
+
+ let data1 = vec![1u8; 400];
+ let data2 = vec![2u8; 200];
+ let data3 = vec![3u8; 100];
+
+ direct_file.write_all(&data1).await.unwrap();
+ direct_file.write_all(&data2).await.unwrap();
+ direct_file.write_all(&data3).await.unwrap();
+
+ assert_eq!(direct_file.tail_len(), 700 % ALIGNMENT);
+
+ direct_file.flush().await.unwrap();
+
+ let file = OpenOptions::new()
+ .read(true)
+ .custom_flags(0x4000)
+ .open(&file_path)
+ .await
+ .unwrap();
+
+ let mut read_buffer = vec![0u8; ALIGNMENT * 2];
+ let (result, buf) = file.read_at(read_buffer, 0).await.into();
+ result.unwrap();
+ read_buffer = buf;
+
+ assert_eq!(&read_buffer[0..400], &vec![1u8; 400]);
+ assert_eq!(&read_buffer[400..600], &vec![2u8; 200]);
+ assert_eq!(&read_buffer[600..700], &vec![3u8; 100]);
+ });
+ }
+}
diff --git a/core/server/src/streaming/segments/indexes/index_reader.rs b/core/server/src/streaming/segments/indexes/index_reader.rs
index ce235ef4..2a9fef00 100644
--- a/core/server/src/streaming/segments/indexes/index_reader.rs
+++ b/core/server/src/streaming/segments/indexes/index_reader.rs
@@ -21,6 +21,7 @@ use crate::streaming::utils::PooledBuffer;
use bytes::BytesMut;
use compio::{
BufResult,
+ buf::{IntoInner, IoBuf},
fs::{File, OpenOptions},
io::AsyncReadAtExt,
};
@@ -336,19 +337,15 @@ impl IndexReader {
len: u32,
use_pool: bool,
) -> Result<PooledBuffer, std::io::Error> {
- if use_pool {
- let mut buf = PooledBuffer::with_capacity(len as usize);
- unsafe { buf.set_len(len as usize) };
- let (result, buf) = self.file.read_exact_at(buf, offset as u64).await.into();
- result?;
- Ok(buf)
- } else {
- let mut buf = BytesMut::with_capacity(len as usize);
- unsafe { buf.set_len(len as usize) };
- let (result, buf) = self.file.read_exact_at(buf, offset as u64).await.into();
- result?;
- Ok(PooledBuffer::from_existing(buf))
- }
+ let mut buf = PooledBuffer::with_capacity(len as usize);
+ let (result, buf) = self
+ .file
+ .read_exact_at(buf.slice(0..len as usize), offset as u64)
+ .await
+ .into();
+ result?;
+
+ Ok(buf.into_inner())
}
/// Gets the nth index from the index file.
diff --git a/core/server/src/streaming/segments/indexes/index_writer.rs b/core/server/src/streaming/segments/indexes/index_writer.rs
index 416cde21..26f3a132 100644
--- a/core/server/src/streaming/segments/indexes/index_writer.rs
+++ b/core/server/src/streaming/segments/indexes/index_writer.rs
@@ -16,25 +16,20 @@
* under the License.
*/
-use compio::fs::File;
-use compio::fs::OpenOptions;
-use compio::io::AsyncWriteAtExt;
+use crate::streaming::segments::DirectFile;
+use crate::streaming::utils::PooledBuffer;
use error_set::ErrContext;
-use iggy_common::INDEX_SIZE;
-use iggy_common::IggyError;
+use iggy_common::{INDEX_SIZE, IggyError};
use std::sync::{
Arc,
atomic::{AtomicU64, Ordering},
};
use tracing::trace;
-use crate::streaming::utils::PooledBuffer;
-
/// A dedicated struct for writing to the index file.
#[derive(Debug)]
pub struct IndexWriter {
- file_path: String,
- file: File,
+ direct_file: DirectFile,
index_size_bytes: Arc<AtomicU64>,
fsync: bool,
}
@@ -47,39 +42,34 @@ impl IndexWriter {
fsync: bool,
file_exists: bool,
) -> Result<Self, IggyError> {
- let file = OpenOptions::new()
- .create(true)
- .write(true)
- .open(file_path)
- .await
- .with_error_context(|error| format!("Failed to open index file: {file_path}. {error}"))
- .map_err(|_| IggyError::CannotReadFile)?;
+ let file_position = if file_exists {
+ let current_size = index_size_bytes.load(Ordering::Acquire);
+ (current_size + 511) & !511
+ } else {
+ index_size_bytes.store(0, Ordering::Release);
+ 0
+ };
- if file_exists {
- let _ = file.sync_all().await.with_error_context(|error| {
- format!("Failed to fsync index file after creation: {file_path}. {error}",)
- });
-
- let actual_index_size = file
- .metadata()
- .await
- .with_error_context(|error| {
- format!("Failed to get metadata of index file: {file_path}. {error}")
- })
- .map_err(|_| IggyError::CannotReadFileMetadata)?
- .len();
+ trace!(
+ "Opening index file for writing: {file_path}, file_position: {}",
+ file_position
+ );
+
+ let mut direct_file = DirectFile::open(file_path, file_position, file_exists).await?;
+ if file_exists {
+ let actual_index_size = direct_file.get_file_size().await?;
index_size_bytes.store(actual_index_size, Ordering::Release);
- }
- trace!(
- "Opened index file for writing: {file_path}, size: {}",
- index_size_bytes.load(Ordering::Acquire)
- );
+ trace!(
+ "Opened existing index file: {file_path}, size: {}, file_position: {}",
+ actual_index_size,
+ direct_file.position()
+ );
+ }
Ok(Self {
- file_path: file_path.to_string(),
- file,
+ direct_file,
index_size_bytes,
fsync,
})
@@ -92,42 +82,49 @@ impl IndexWriter {
}
let count = indexes.len() / INDEX_SIZE;
- let len = indexes.len();
+ let actual_len = indexes.len();
+
+ trace!(
+ "Saving {count} indexes to file: {} (size: {} bytes)",
+ self.direct_file.file_path(),
+ actual_len
+ );
- let position = self.index_size_bytes.load(Ordering::Relaxed);
- self.file
- .write_all_at(indexes, position)
+ let bytes_written = self
+ .direct_file
+ .write_all(&indexes)
.await
- .0
.with_error_context(|error| {
format!(
"Failed to write {} indexes to file: {}. {error}",
- count, self.file_path
+ count,
+ self.direct_file.file_path()
)
})
.map_err(|_| IggyError::CannotSaveIndexToSegment)?;
+ let new_logical_size = self.index_size_bytes.load(Ordering::Relaxed) + bytes_written as u64;
self.index_size_bytes
- .fetch_add(len as u64, Ordering::Release);
+ .store(new_logical_size, Ordering::Release);
if self.fsync {
let _ = self.fsync().await;
}
- trace!(
- "Saved {count} indexes of size {} to file: {}",
- INDEX_SIZE * count,
- self.file_path
- );
+
+ trace!("Saved {count} indexes. Logical size: {}", new_logical_size);
Ok(())
}
pub async fn fsync(&self) -> Result<(), IggyError> {
- self.file
+ self.direct_file
.sync_all()
.await
.with_error_context(|error| {
- format!("Failed to fsync index file: {}. {error}", self.file_path)
+ format!(
+ "Failed to fsync index file: {}. {error}",
+ self.direct_file.file_path()
+ )
})
.map_err(|_| IggyError::CannotWriteToFile)?;
Ok(())
diff --git a/core/server/src/streaming/segments/messages/messages_reader.rs b/core/server/src/streaming/segments/messages/messages_reader.rs
index e81eb793..5ad23502 100644
--- a/core/server/src/streaming/segments/messages/messages_reader.rs
+++ b/core/server/src/streaming/segments/messages/messages_reader.rs
@@ -19,6 +19,7 @@
use crate::streaming::segments::{IggyIndexesMut, IggyMessagesBatchMut};
use crate::streaming::utils::PooledBuffer;
use bytes::BytesMut;
+use compio::buf::{IntoInner, IoBuf};
use compio::fs::{File, OpenOptions};
use compio::io::AsyncReadAtExt;
use error_set::ErrContext;
@@ -176,18 +177,13 @@ impl MessagesReader {
len: u32,
use_pool: bool,
) -> Result<PooledBuffer, std::io::Error> {
- if use_pool {
- let mut buf = PooledBuffer::with_capacity(len as usize);
- unsafe { buf.set_len(len as usize) };
- let (result, buf) = self.file.read_exact_at(buf, offset as u64).await.into();
- result?;
- Ok(buf)
- } else {
- let mut buf = BytesMut::with_capacity(len as usize);
- unsafe { buf.set_len(len as usize) };
- let (result, buf) = self.file.read_exact_at(buf, offset as u64).await.into();
- result?;
- Ok(PooledBuffer::from_existing(buf))
- }
+ let mut buf = PooledBuffer::with_capacity(len as usize);
+ let (result, buf) = self
+ .file
+ .read_exact_at(buf.slice(0..len as usize), offset as u64)
+ .await
+ .into();
+ result?;
+ Ok(buf.into_inner())
}
}
diff --git a/core/server/src/streaming/segments/messages/messages_writer.rs b/core/server/src/streaming/segments/messages/messages_writer.rs
index f3d43d1e..ba9501f8 100644
--- a/core/server/src/streaming/segments/messages/messages_writer.rs
+++ b/core/server/src/streaming/segments/messages/messages_writer.rs
@@ -16,8 +16,9 @@
* under the License.
*/
-use crate::streaming::segments::{IggyMessagesBatchSet, messages::write_batch};
-use compio::fs::{File, OpenOptions};
+use crate::streaming::segments::{
+ DirectFile, IggyMessagesBatchSet, messages::write_batch_with_direct_file,
+};
use error_set::ErrContext;
use iggy_common::{IggyByteSize, IggyError};
use std::sync::{
@@ -29,68 +30,51 @@ use tracing::{error, trace};
/// A dedicated struct for writing to the messages file.
#[derive(Debug)]
pub struct MessagesWriter {
- file_path: String,
- /// Holds the file for synchronous writes; when asynchronous persistence is enabled, this will be None.
- file: Option<File>,
- /// When set, asynchronous writes are handled by this persister task.
+ direct_file: Option<DirectFile>,
messages_size_bytes: Arc<AtomicU64>,
fsync: bool,
}
impl MessagesWriter {
- /// Opens the messages file in write mode.
- ///
- /// If the server confirmation is set to `NoWait`, the file handle is transferred to the
- /// persister task (and stored in `persister_task`) so that writes are done asynchronously.
- /// Otherwise, the file is retained in `self.file` for synchronous writes.
pub async fn new(
file_path: &str,
messages_size_bytes: Arc<AtomicU64>,
fsync: bool,
file_exists: bool,
) -> Result<Self, IggyError> {
- let file = OpenOptions::new()
- .create(true)
- .write(true)
- .open(file_path)
- .await
- .with_error_context(|err| {
- format!("Failed to open messages file: {file_path}, error: {err}")
- })
- .map_err(|_| IggyError::CannotReadFile)?;
+ let file_position = if file_exists {
+ let current_size = messages_size_bytes.load(Ordering::Acquire);
+ (current_size + 511) & !511
+ } else {
+ messages_size_bytes.store(0, Ordering::Release);
+ 0
+ };
- if file_exists {
- let _ = file.sync_all().await.with_error_context(|error| {
- format!("Failed to fsync messages file after creation: {file_path}, error: {error}")
- });
+ trace!(
+ "Opening messages file for writing: {file_path}, file_position: {}",
+ file_position
+ );
- let actual_messages_size = file
- .metadata()
- .await
- .with_error_context(|error| {
- format!("Failed to get metadata of messages file: {file_path}, error: {error}")
- })
- .map_err(|_| IggyError::CannotReadFileMetadata)?
- .len();
+ let mut direct_file = DirectFile::open(file_path, file_position, file_exists).await?;
+ if file_exists {
+ let actual_messages_size = direct_file.get_file_size().await?;
messages_size_bytes.store(actual_messages_size, Ordering::Release);
- }
- trace!(
- "Opened messages file for writing: {file_path}, size: {}",
- messages_size_bytes.load(Ordering::Acquire)
- );
+ trace!(
+ "Opened existing messages file: {file_path}, size: {}, file_position: {}",
+ actual_messages_size,
+ direct_file.position()
+ );
+ }
- let file = Some(file);
Ok(Self {
- file_path: file_path.to_string(),
- file,
+ direct_file: Some(direct_file),
messages_size_bytes,
fsync,
})
}
- /// Append a batch of messages to the messages file.
pub async fn save_batch_set(
&mut self,
batch_set: IggyMessagesBatchSet,
@@ -98,45 +82,51 @@ impl MessagesWriter {
let messages_size = batch_set.size();
let messages_count = batch_set.count();
let containers_count = batch_set.containers_count();
- trace!(
- "Saving batch set of size {messages_size} bytes ({containers_count} containers, {messages_count} messages) to messages file: {}",
- self.file_path
- );
- let position = self.messages_size_bytes.load(Ordering::Relaxed);
- if let Some(ref mut file) = self.file {
- write_batch(file, position, batch_set)
+ let actual_written = if let Some(ref mut direct_file) = self.direct_file {
+ trace!(
+ "Saving batch set of size {messages_size} bytes ({containers_count} containers, {messages_count} messages) to messages file: {}",
+ direct_file.file_path()
+ );
+
+ write_batch_with_direct_file(direct_file, batch_set)
.await
.with_error_context(|error| {
format!(
"Failed to write batch to messages file: {}. {error}",
- self.file_path
+ direct_file.file_path()
)
- })?;
+ })?
} else {
- error!("File handle is not available for synchronous write.");
+ tracing::error!("File handle is not available for synchronous write.");
return Err(IggyError::CannotWriteToFile);
- }
+ };
if self.fsync {
let _ = self.fsync().await;
}
+ let logical_size = self.messages_size_bytes.load(Ordering::Relaxed) + actual_written as u64;
self.messages_size_bytes
- .fetch_add(messages_size as u64, Ordering::Release);
+ .store(logical_size, Ordering::Release);
+
trace!(
- "Written batch set of size {messages_size} bytes ({containers_count} containers, {messages_count} messages) to disk messages file: {}",
- self.file_path
+ "Written batch set of size {messages_size} bytes to disk. Logical size: {}",
+ logical_size
);
Ok(IggyByteSize::from(messages_size as u64))
}
pub async fn fsync(&self) -> Result<(), IggyError> {
- if let Some(file) = self.file.as_ref() {
- file.sync_all()
+ if let Some(direct_file) = self.direct_file.as_ref() {
+ direct_file
+ .sync_all()
.await
.with_error_context(|error| {
- format!("Failed to fsync messages file: {}. {error}", self.file_path)
+ format!(
+ "Failed to fsync messages file: {}. {error}",
+ direct_file.file_path()
+ )
})
.map_err(|_| IggyError::CannotWriteToFile)?;
}
diff --git a/core/server/src/streaming/segments/messages/mod.rs b/core/server/src/streaming/segments/messages/mod.rs
index a3989a2e..537ecb78 100644
--- a/core/server/src/streaming/segments/messages/mod.rs
+++ b/core/server/src/streaming/segments/messages/mod.rs
@@ -19,25 +19,27 @@
mod messages_reader;
mod messages_writer;
-use super::IggyMessagesBatchSet;
-use compio::{fs::File, io::AsyncWriteAtExt};
+use super::{DirectFile, IggyMessagesBatchSet};
+use crate::streaming::utils::PooledBuffer;
use iggy_common::IggyError;
pub use messages_reader::MessagesReader;
pub use messages_writer::MessagesWriter;
-/// Vectored write a batches of messages to file
-async fn write_batch(
- file: &mut File,
- position: u64,
+async fn write_batch_with_direct_file(
+ direct_file: &mut DirectFile,
mut batches: IggyMessagesBatchSet,
) -> Result<usize, IggyError> {
- let total_written = batches.iter().map(|b| b.size() as usize).sum();
- let batches = batches
- .iter_mut()
- .map(|b| b.take_messages())
- .collect::<Vec<_>>();
- let (result, _) = file.write_vectored_all_at(batches, position).await.into();
- result.map_err(|_| IggyError::CannotWriteToFile)?;
+ let total_written: usize = batches.iter().map(|b| b.size() as usize).sum();
+ let mut messages_count = 0;
+
+ for batch in batches.iter_mut() {
+ messages_count += batch.count();
+ let messages = batch.take_messages();
+ direct_file.write_all(&messages).await?;
+ }
+
+ tracing::trace!("Saved {} messages", messages_count);
+
Ok(total_written)
}
diff --git a/core/server/src/streaming/segments/mod.rs b/core/server/src/streaming/segments/mod.rs
index f3105e54..9fbd4abe 100644
--- a/core/server/src/streaming/segments/mod.rs
+++ b/core/server/src/streaming/segments/mod.rs
@@ -16,6 +16,7 @@
* under the License.
*/
+mod direct_file;
mod indexes;
mod messages;
mod messages_accumulator;
@@ -24,6 +25,7 @@ mod segment;
mod types;
mod writing_messages;
+pub use direct_file::DirectFile;
pub use indexes::IggyIndexesMut;
pub use messages_accumulator::MessagesAccumulator;
pub use segment::Segment;
diff --git a/core/server/src/streaming/segments/reading_messages.rs b/core/server/src/streaming/segments/reading_messages.rs
index 5f0f6bb4..fa46c183 100644
--- a/core/server/src/streaming/segments/reading_messages.rs
+++ b/core/server/src/streaming/segments/reading_messages.rs
@@ -273,8 +273,16 @@ impl Segment {
relative_start_offset: u32,
count: u32,
) -> Result<Option<IggyIndexesMut>, IggyError> {
- let indexes = if !self.indexes.is_empty() {
- self.indexes.slice_by_offset(relative_start_offset, count)
+ let indexes = if let Some(ref indexes) = self.indexes {
+ if !indexes.is_empty() {
+ indexes.slice_by_offset(relative_start_offset, count)
+ } else {
+ self.index_reader
+ .as_ref()
+ .expect("Index reader not initialized")
+ .load_from_disk_by_offset(relative_start_offset, count)
+ .await?
+ }
} else {
self.index_reader
.as_ref()
@@ -290,8 +298,16 @@ impl Segment {
timestamp: u64,
count: u32,
) -> Result<Option<IggyIndexesMut>, IggyError> {
- let indexes = if !self.indexes.is_empty() {
- self.indexes.slice_by_timestamp(timestamp, count)
+ let indexes = if let Some(ref indexes) = self.indexes {
+ if !indexes.is_empty() {
+ indexes.slice_by_timestamp(timestamp, count)
+ } else {
+ self.index_reader
+ .as_ref()
+ .unwrap()
+ .load_from_disk_by_timestamp(timestamp, count)
+ .await?
+ }
} else {
self.index_reader
.as_ref()
diff --git a/core/server/src/streaming/segments/segment.rs b/core/server/src/streaming/segments/segment.rs
index 08b30ef6..937f4cae 100644
--- a/core/server/src/streaming/segments/segment.rs
+++ b/core/server/src/streaming/segments/segment.rs
@@ -62,7 +62,7 @@ pub struct Segment {
pub(super) message_expiry: IggyExpiry,
pub(super) accumulator: MessagesAccumulator,
pub(super) config: Arc<SystemConfig>,
- pub(super) indexes: IggyIndexesMut,
+ pub(super) indexes: Option<IggyIndexesMut>,
pub(super) messages_size: Arc<AtomicU64>,
pub(super) indexes_size: Arc<AtomicU64>,
}
@@ -92,10 +92,6 @@ impl Segment {
_ => message_expiry,
};
- // In order to preserve BytesMut buffer between restarts, initialize it with a capacity 0.
- // We don't care whether server startup would be couple of seconds longer.
- let indexes_capacity = if fresh { SIZE_16MB / INDEX_SIZE } else { 0 };
-
Segment {
stream_id,
topic_id,
@@ -109,7 +105,7 @@ impl Segment {
last_index_position: 0,
max_size_bytes: config.segment.size,
message_expiry,
- indexes: IggyIndexesMut::with_capacity(indexes_capacity, 0),
+ indexes: None,
accumulator: MessagesAccumulator::default(),
is_closed: false,
messages_writer: None,
@@ -145,7 +141,7 @@ impl Segment {
self.last_index_position = log_size_bytes as _;
- self.indexes = self
+ let loaded_indexes = self
.index_reader
.as_ref()
.unwrap()
@@ -154,17 +150,21 @@ impl Segment {
.with_error_context(|error| format!("Failed to load indexes for {self}. {error}"))
.map_err(|_| IggyError::CannotReadFile)?;
- let last_index_offset = if self.indexes.is_empty() {
+ let last_index_offset = if loaded_indexes.is_empty() {
0_u64
} else {
- self.indexes.last().unwrap().offset() as u64
+ loaded_indexes.last().unwrap().offset() as u64
};
+ if !loaded_indexes.is_empty() {
+ self.indexes = Some(loaded_indexes);
+ }
+
self.end_offset = self.start_offset + last_index_offset;
info!(
"Loaded {} indexes for segment with start offset: {}, end offset: {}, and partition with ID: {}, topic with ID: {}, and stream with ID: {}.",
- self.indexes.count(),
+ self.indexes.as_ref().map_or(0, |idx| idx.count()),
self.start_offset,
self.end_offset,
self.partition_id,
@@ -300,27 +300,8 @@ impl Segment {
}
pub async fn shutdown_writing(&mut self) {
- if let Some(log_writer) = self.messages_writer.take() {
- //TODO: Fixme not sure whether we should spawn a task here.
- compio::runtime::spawn(async move {
- let _ = log_writer.fsync().await;
- });
- } else {
- warn!(
- "Log writer already closed when calling close() for {}",
- self
- );
- }
-
- if let Some(index_writer) = self.index_writer.take() {
- //TODO: Fixme not sure whether we should spawn a task here.
- compio::runtime::spawn(async move {
- let _ = index_writer.fsync().await;
- drop(index_writer)
- });
- } else {
- warn!("Index writer already closed when calling close()");
- }
+ let _ = self.messages_writer.take().map(|mut writer| {});
+ let _ = self.index_writer.take().map(|mut writer| {});
}
pub async fn delete(&mut self) -> Result<(), IggyError> {
@@ -382,6 +363,14 @@ impl Segment {
self.message_expiry = message_expiry;
}
+ /// Ensure indexes are initialized with proper capacity
+ pub fn ensure_indexes(&mut self) {
+ if self.indexes.is_none() {
+ let capacity = SIZE_16MB / INDEX_SIZE;
+ self.indexes = Some(IggyIndexesMut::with_capacity(capacity, 0));
+ }
+ }
+
pub fn is_closed(&self) -> bool {
self.is_closed
}
@@ -412,8 +401,7 @@ impl Segment {
/// Explicitly drop the old indexes to ensure memory is freed
pub fn drop_indexes(&mut self) {
- let old_indexes = std::mem::replace(&mut self.indexes, IggyIndexesMut::empty());
- drop(old_indexes);
+ self.indexes = None;
}
}
@@ -487,7 +475,7 @@ mod tests {
assert_eq!(segment.messages_file_path(), messages_file_path);
assert_eq!(segment.index_file_path(), index_path);
assert_eq!(segment.message_expiry, message_expiry);
- assert!(segment.indexes.is_empty());
+ assert!(segment.indexes.is_none());
assert!(!segment.is_closed());
assert!(!segment.is_full().await);
}
@@ -530,6 +518,6 @@ mod tests {
true,
);
- assert!(segment.indexes.is_empty());
+ assert!(segment.indexes.is_none());
}
}
diff --git a/core/server/src/streaming/segments/types/messages_batch_mut.rs b/core/server/src/streaming/segments/types/messages_batch_mut.rs
index 764faa11..7152f53d 100644
--- a/core/server/src/streaming/segments/types/messages_batch_mut.rs
+++ b/core/server/src/streaming/segments/types/messages_batch_mut.rs
@@ -488,7 +488,7 @@ impl IggyMessagesBatchMut {
/// subsequent messages in the new buffer.
#[allow(clippy::too_many_arguments)]
fn rebuild_indexes_for_chunk(
- new_buffer: &BytesMut,
+ new_buffer: &PooledBuffer,
new_indexes: &mut IggyIndexesMut,
offset_in_new_buffer: &mut u32,
chunk_start: usize,
@@ -571,7 +571,7 @@ impl IggyMessagesBatchMut {
for &(start, end) in &boundaries_to_remove {
if start > last_pos {
- let keep_len = start - last_pos;
+ let keep_len: usize = start - last_pos;
let chunk = source.split_to(keep_len);
let chunk_start_in_new_buffer = new_buffer.len();
new_buffer.put(chunk);
@@ -684,7 +684,8 @@ impl IggyMessagesBatchMut {
prev_offset = message.header().offset();
prev_position = index.position();
- messages_size += message.size();
+ let msg_size = message.size();
+ messages_size += msg_size;
messages_count += 1;
}
@@ -822,11 +823,3 @@ impl Index<usize> for IggyMessagesBatchMut {
&self.messages[start..end]
}
}
-
-impl Deref for IggyMessagesBatchMut {
- type Target = BytesMut;
-
- fn deref(&self) -> &Self::Target {
- &self.messages
- }
-}
diff --git a/core/server/src/streaming/segments/writing_messages.rs b/core/server/src/streaming/segments/writing_messages.rs
index e95d1e2c..8c6a785c 100644
--- a/core/server/src/streaming/segments/writing_messages.rs
+++ b/core/server/src/streaming/segments/writing_messages.rs
@@ -82,7 +82,8 @@ impl Segment {
let batch_size = batches.size();
let batch_count = batches.count();
- batches.append_indexes_to(&mut self.indexes);
+ self.ensure_indexes();
+ batches.append_indexes_to(self.indexes.as_mut().unwrap());
let saved_bytes = self
.messages_writer
@@ -98,7 +99,7 @@ impl Segment {
self.last_index_position += saved_bytes.as_bytes_u64() as u32;
- let unsaved_indexes_slice = self.indexes.unsaved_slice();
+ let unsaved_indexes_slice = self.indexes.as_ref().unwrap().unsaved_slice();
let len = unsaved_indexes_slice.len();
self.index_writer
.as_mut()
@@ -109,17 +110,23 @@ impl Segment {
format!("Failed to save index of {} indexes to {self}. {error}", len)
})?;
- self.indexes.mark_saved();
+ self.indexes.as_mut().unwrap().mark_saved();
if self.config.segment.cache_indexes == CacheIndexesConfig::None {
- self.indexes.clear();
+ if let Some(indexes) = self.indexes.as_mut() {
+ indexes.clear();
+ }
}
self.check_and_handle_segment_full().await?;
trace!(
- "Saved {} messages on disk in segment with start offset: {} for partition with ID: {}, total bytes written: {}.",
- unsaved_messages_count, self.start_offset, self.partition_id, saved_bytes
+ "Saved {} messages on disk in segment with start offset: {}, end offset: {}, for partition with ID: {}, total bytes written: {}.",
+ unsaved_messages_count,
+ self.start_offset,
+ self.end_offset,
+ self.partition_id,
+ saved_bytes
);
Ok(unsaved_messages_count)
diff --git a/core/server/src/streaming/utils/memory_pool.rs b/core/server/src/streaming/utils/memory_pool.rs
index 2cc031d3..df0783b1 100644
--- a/core/server/src/streaming/utils/memory_pool.rs
+++ b/core/server/src/streaming/utils/memory_pool.rs
@@ -17,7 +17,7 @@
*/
use crate::configs::system::SystemConfig;
-use bytes::BytesMut;
+use aligned_vec::{AVec, ConstAlign};
use crossbeam::queue::ArrayQueue;
use human_repr::HumanCount;
use once_cell::sync::OnceCell;
@@ -26,15 +26,18 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tracing::{error, info, trace, warn};
+pub const ALIGNMENT: usize = 512;
+pub type Align512 = ConstAlign<ALIGNMENT>;
+pub type AlignedBuffer = AVec<u8, Align512>;
+
/// Global memory pool instance. Use `memory_pool()` to access it.
pub static MEMORY_POOL: OnceCell<MemoryPool> = OnceCell::new();
/// Total number of distinct bucket sizes.
-const NUM_BUCKETS: usize = 32;
+const NUM_BUCKETS: usize = 31;
/// Array of bucket sizes in ascending order. Each entry is a distinct buffer size (in bytes).
const BUCKET_SIZES: [usize; NUM_BUCKETS] = [
- 256,
512,
1024,
2 * 1024,
@@ -49,8 +52,8 @@ const BUCKET_SIZES: [usize; NUM_BUCKETS] = [
768 * 1024,
1024 * 1024,
1536 * 1024,
- 2 * 1024 * 1024, // Above 2MiB everything should be rounded up to the next power of 2 to take advantage of hugepages
- 4 * 1024 * 1024, // (environment variables MIMALLOC_ALLOW_LARGE_OS_PAGES=1 and MIMALLOC_LARGE_OS_PAGES=1).
+ 2 * 1024 * 1024,
+ 4 * 1024 * 1024,
6 * 1024 * 1024,
8 * 1024 * 1024,
10 * 1024 * 1024,
@@ -75,7 +78,7 @@ pub fn memory_pool() -> &'static MemoryPool {
.expect("Memory pool not initialized - MemoryPool::init_pool should be called first")
}
-/// A memory pool that maintains fixed-size buckets for reusing `BytesMut` buffers.
+/// A memory pool that maintains fixed-size buckets for reusing `AlignedBuffer` buffers.
///
/// Each bucket corresponds to a particular size in `BUCKET_SIZES`. The pool tracks:
/// - Buffers currently in use (`in_use`)
@@ -97,7 +100,7 @@ pub struct MemoryPool {
/// Array of queues for reusable buffers. Each queue can store up to `bucket_capacity` buffers.
/// The length of each queue (`buckets[i].len()`) is how many **free** buffers are currently available.
/// Free doesn't mean the buffer is allocated, it just means it's not in use.
- buckets: [Arc<ArrayQueue<BytesMut>>; NUM_BUCKETS],
+ buckets: [Arc<ArrayQueue<AlignedBuffer>>; NUM_BUCKETS],
/// Number of buffers **in use** for each bucket size (grow/shrink as they are acquired/released).
in_use: [Arc<AtomicUsize>; NUM_BUCKETS],
@@ -163,16 +166,16 @@ impl MemoryPool {
MEMORY_POOL.get_or_init(|| MemoryPool::new(is_enabled, memory_limit, bucket_capacity));
}
- /// Acquire a `BytesMut` buffer with at least `capacity` bytes.
+ /// Acquire a `AlignedBuffer` buffer with at least `capacity` bytes.
///
/// - If a bucket can fit `capacity`, try to pop from its free buffer queue; otherwise create a new buffer.
/// - If `memory_limit` would be exceeded, allocate outside the pool.
///
/// Returns a tuple of (buffer, was_pool_allocated) where was_pool_allocated indicates if the buffer
/// was allocated from the pool (true) or externally (false).
- pub fn acquire_buffer(&self, capacity: usize) -> (BytesMut, bool) {
+ pub fn acquire_buffer(&self, capacity: usize) -> (AlignedBuffer, bool) {
if !self.is_enabled {
- return (BytesMut::with_capacity(capacity), false);
+ return (AlignedBuffer::with_capacity(ALIGNMENT, capacity), false);
}
let current = self.pool_current_size();
@@ -193,12 +196,12 @@ impl MemoryPool {
new_size, current, self.memory_limit
);
self.inc_external_allocations();
- return (BytesMut::with_capacity(new_size), false);
+ return (AlignedBuffer::with_capacity(ALIGNMENT, new_size), false);
}
self.inc_bucket_alloc(idx);
self.inc_bucket_in_use(idx);
- (BytesMut::with_capacity(new_size), true)
+ (AlignedBuffer::with_capacity(ALIGNMENT, new_size), true)
}
None => {
if current + capacity > self.memory_limit {
@@ -207,16 +210,16 @@ impl MemoryPool {
capacity, current, self.memory_limit
);
self.inc_external_allocations();
- return (BytesMut::with_capacity(capacity), false);
+ return (AlignedBuffer::with_capacity(ALIGNMENT, capacity), false);
}
self.inc_external_allocations();
- (BytesMut::with_capacity(capacity), false)
+ (AlignedBuffer::with_capacity(ALIGNMENT, capacity), false)
}
}
}
- /// Return a `BytesMut` buffer previously acquired from the pool.
+ /// Return a `AlignedBuffer` buffer previously acquired from the pool.
///
/// - If `current_capacity` differs from `original_capacity`, increments `resize_events`.
/// - If a matching bucket exists, place it back in that bucket's queue (if space is available).
@@ -224,7 +227,7 @@ impl MemoryPool {
/// - The `was_pool_allocated` flag indicates if this buffer was originally allocated from the pool.
pub fn release_buffer(
&self,
- buffer: BytesMut,
+ buffer: AlignedBuffer,
original_capacity: usize,
was_pool_allocated: bool,
) {
@@ -235,10 +238,6 @@ impl MemoryPool {
let current_capacity = buffer.capacity();
if current_capacity != original_capacity {
self.inc_resize_events();
- trace!(
- "Buffer capacity {} != original {} when returning",
- current_capacity, original_capacity
- );
}
if was_pool_allocated {
@@ -438,11 +437,11 @@ impl MemoryPool {
/// Return a buffer to the pool by calling `release_buffer` with the original capacity.
/// This extension trait makes it easy to do `some_bytes.return_to_pool(orig_cap, was_pool_allocated)`.
-pub trait BytesMutExt {
+pub trait AlignedBufferExt {
fn return_to_pool(self, original_capacity: usize, was_pool_allocated: bool);
}
-impl BytesMutExt for BytesMut {
+impl AlignedBufferExt for AlignedBuffer {
fn return_to_pool(self, original_capacity: usize, was_pool_allocated: bool) {
memory_pool().release_buffer(self, original_capacity, was_pool_allocated);
}
@@ -667,27 +666,6 @@ mod tests {
);
}
- // Test put_bytes
- {
- let initial_events = pool.resize_events();
- let mut buffer = PooledBuffer::with_capacity(4 * 1024);
- let orig_bucket_idx = pool.best_fit(buffer.capacity()).unwrap();
- let orig_in_use = pool.bucket_current_elements(orig_bucket_idx);
-
- buffer.put_bytes(0, 64 * 1024); // 64KiB of zeros
-
- assert_eq!(
- pool.resize_events(),
- initial_events + 1,
- "put_bytes should trigger resize event"
- );
- assert_eq!(
- pool.bucket_current_elements(orig_bucket_idx),
- orig_in_use - 1,
- "put_bytes should update bucket accounting"
- );
- }
-
// Test extend_from_slice
{
let initial_events = pool.resize_events();
diff --git a/core/server/src/streaming/utils/mod.rs b/core/server/src/streaming/utils/mod.rs
index e441a051..d858d684 100644
--- a/core/server/src/streaming/utils/mod.rs
+++ b/core/server/src/streaming/utils/mod.rs
@@ -25,5 +25,5 @@ pub mod random_id;
mod memory_pool;
mod pooled_buffer;
-pub use memory_pool::{MemoryPool, memory_pool};
+pub use memory_pool::{ALIGNMENT, MemoryPool, memory_pool};
pub use pooled_buffer::PooledBuffer;
diff --git a/core/server/src/streaming/utils/pooled_buffer.rs b/core/server/src/streaming/utils/pooled_buffer.rs
index ee6f2014..b3769bf2 100644
--- a/core/server/src/streaming/utils/pooled_buffer.rs
+++ b/core/server/src/streaming/utils/pooled_buffer.rs
@@ -16,8 +16,8 @@
* under the License.
*/
-use super::memory_pool::{BytesMutExt, memory_pool};
-use bytes::{Buf, BufMut, BytesMut};
+use super::memory_pool::{AlignedBuffer, AlignedBufferExt, memory_pool};
+use crate::streaming::utils::memory_pool::ALIGNMENT;
use compio::buf::{IoBuf, IoBufMut, SetBufInit};
use std::ops::{Deref, DerefMut};
@@ -26,7 +26,7 @@ pub struct PooledBuffer {
from_pool: bool,
original_capacity: usize,
original_bucket_idx: Option<usize>,
- inner: BytesMut,
+ inner: AlignedBuffer,
}
impl Default for PooledBuffer {
@@ -42,13 +42,14 @@ impl PooledBuffer {
///
/// * `capacity` - The capacity of the buffer
pub fn with_capacity(capacity: usize) -> Self {
- let (buffer, was_pool_allocated) = memory_pool().acquire_buffer(capacity);
+ let (mut buffer, was_pool_allocated) = memory_pool().acquire_buffer(capacity);
let original_capacity = buffer.capacity();
let original_bucket_idx = if was_pool_allocated {
memory_pool().best_fit(original_capacity)
} else {
None
};
+
Self {
from_pool: was_pool_allocated,
original_capacity,
@@ -57,27 +58,13 @@ impl PooledBuffer {
}
}
- /// Creates a new pooled buffer from an existing `BytesMut`.
- ///
- /// # Arguments
- ///
- /// * `existing` - The existing `BytesMut` buffer
- pub fn from_existing(existing: BytesMut) -> Self {
- Self {
- from_pool: false,
- original_capacity: existing.capacity(),
- original_bucket_idx: None,
- inner: existing,
- }
- }
-
/// Creates an empty pooled buffer.
pub fn empty() -> Self {
Self {
from_pool: false,
original_capacity: 0,
original_bucket_idx: None,
- inner: BytesMut::new(),
+ inner: AlignedBuffer::new(ALIGNMENT),
}
}
@@ -90,6 +77,11 @@ impl PooledBuffer {
let current_capacity = self.inner.capacity();
if current_capacity != self.original_capacity {
+ tracing::error!(
+ "Pooled buffer resized from {} to {}",
+ self.original_capacity,
+ current_capacity
+ );
memory_pool().inc_resize_events();
if let Some(orig_idx) = self.original_bucket_idx {
@@ -131,49 +123,111 @@ impl PooledBuffer {
}
}
- /// Wrapper for put_bytes which might cause resize
- pub fn put_bytes(&mut self, byte: u8, len: usize) {
+ /// Wrapper for put_slice which might cause resize
+ pub fn put_slice(&mut self, src: &[u8]) {
let before_cap = self.inner.capacity();
- self.inner.put_bytes(byte, len);
+ self.extend_from_slice(src);
if self.inner.capacity() != before_cap {
self.check_for_resize();
}
}
- /// Wrapper for put_slice which might cause resize
- pub fn put_slice(&mut self, src: &[u8]) {
+ /// Wrapper for put_u32_le which might cause resize
+ pub fn put_u32_le(&mut self, value: u32) {
let before_cap = self.inner.capacity();
- self.inner.put_slice(src);
+ self.reserve(4);
+ self.inner.extend_from_slice(&value.to_le_bytes());
if self.inner.capacity() != before_cap {
self.check_for_resize();
}
}
- /// Wrapper for put_u32_le which might cause resize
- pub fn put_u32_le(&mut self, value: u32) {
+ /// Wrapper for put_u64_le which might cause resize
+ pub fn put_u64_le(&mut self, value: u64) {
let before_cap = self.inner.capacity();
- self.inner.put_u32_le(value);
+ self.reserve(8);
+ self.inner.extend_from_slice(&value.to_le_bytes());
if self.inner.capacity() != before_cap {
self.check_for_resize();
}
}
- /// Wrapper for put_u64_le which might cause resize
- pub fn put_u64_le(&mut self, value: u64) {
+ /// Get a slice of the buffer's contents
+ pub fn as_slice(&self) -> &[u8] {
+ self.inner.as_slice()
+ }
+
+ /// Get the length of the buffer
+ pub fn len(&self) -> usize {
+ self.inner.len()
+ }
+
+ /// Check if the buffer is empty
+ pub fn is_empty(&self) -> bool {
+ self.inner.is_empty()
+ }
+
+ /// Get the capacity of the buffer
+ pub fn capacity(&self) -> usize {
+ self.inner.capacity()
+ }
+
+ /// Clear the buffer
+ pub fn clear(&mut self) {
+ self.inner.clear()
+ }
+
+ /// Resize the buffer
+ pub fn resize(&mut self, new_len: usize, value: u8) {
let before_cap = self.inner.capacity();
- self.inner.put_u64_le(value);
+ self.inner.resize(new_len, value);
if self.inner.capacity() != before_cap {
self.check_for_resize();
}
}
+
+ /// Split the buffer at the given index, returning the data before the split
+ /// and keeping the data after the split in self
+ pub fn split_to(&mut self, at: usize) -> Vec<u8> {
+ if at > self.inner.len() {
+ panic!("split_to out of bounds");
+ }
+
+ let mut result = Vec::with_capacity(at);
+ result.extend_from_slice(&self.inner[..at]);
+
+ let remaining = self.inner.len() - at;
+ let mut new_data = Vec::with_capacity(remaining);
+ new_data.extend_from_slice(&self.inner[at..]);
+
+ self.inner.clear();
+ self.inner.extend_from_slice(&new_data);
+
+ result
+ }
+
+ /// Put bytes from a slice
+ pub fn put<T: AsRef<[u8]>>(&mut self, data: T) {
+ self.extend_from_slice(data.as_ref());
+ }
+
+ /// Align the buffer length to the next 512-byte boundary by padding with zeros
+ pub fn align(&mut self) {
+ let current_len = self.inner.len();
+ let aligned_len = (current_len + 511) & !511;
+ if aligned_len > current_len {
+ let padding = aligned_len - current_len;
+ self.resize(aligned_len, 0);
+ }
+ }
}
impl Deref for PooledBuffer {
- type Target = BytesMut;
+ type Target = AlignedBuffer;
fn deref(&self) -> &Self::Target {
&self.inner
@@ -189,12 +243,18 @@ impl DerefMut for PooledBuffer {
impl Drop for PooledBuffer {
fn drop(&mut self) {
if self.from_pool {
- let buf = std::mem::take(&mut self.inner);
+ let buf = std::mem::replace(&mut self.inner, AlignedBuffer::new(ALIGNMENT));
buf.return_to_pool(self.original_capacity, true);
}
}
}
+impl AsRef<[u8]> for PooledBuffer {
+ fn as_ref(&self) -> &[u8] {
+ self.inner.as_slice()
+ }
+}
+
impl From<&[u8]> for PooledBuffer {
fn from(slice: &[u8]) -> Self {
let mut buf = PooledBuffer::with_capacity(slice.len());
@@ -203,30 +263,10 @@ impl From<&[u8]> for PooledBuffer {
}
}
-impl Buf for PooledBuffer {
- fn remaining(&self) -> usize {
- self.inner.remaining()
- }
-
- fn chunk(&self) -> &[u8] {
- self.inner.chunk()
- }
-
- fn advance(&mut self, cnt: usize) {
- self.inner.advance(cnt)
- }
-
- fn chunks_vectored<'t>(&'t self, dst: &mut [std::io::IoSlice<'t>]) -> usize {
- self.inner.chunks_vectored(dst)
- }
-}
-
impl SetBufInit for PooledBuffer {
unsafe fn set_buf_init(&mut self, len: usize) {
- if self.inner.len() <= len {
- unsafe {
- self.inner.set_len(len);
- }
+ unsafe {
+ self.inner.set_len(len);
}
}
}
@@ -239,7 +279,7 @@ unsafe impl IoBufMut for PooledBuffer {
unsafe impl IoBuf for PooledBuffer {
fn as_buf_ptr(&self) -> *const u8 {
- self.inner.as_buf_ptr()
+ self.inner.as_ptr()
}
fn buf_len(&self) -> usize {
diff --git a/core/server/src/tcp/connection_handler.rs b/core/server/src/tcp/connection_handler.rs
index 9c6803b8..1f1bb482 100644
--- a/core/server/src/tcp/connection_handler.rs
+++ b/core/server/src/tcp/connection_handler.rs
@@ -49,7 +49,7 @@ pub(crate) async fn handle_connection(
loop {
let read_future = sender.read(length_buffer.clone());
- let (read_length, initial_buffer) = futures::select! {
+ let initial_buffer = futures::select! {
_ = stop_receiver.recv().fuse() => {
info!("Connection stop signal received for session: {}", session);
let _ = sender.send_error_response(IggyError::Disconnected).await;
@@ -57,8 +57,8 @@ pub(crate) async fn handle_connection(
}
result = read_future.fuse() => {
match result {
- (Ok(read_length), initial_buffer) => (read_length, initial_buffer),
- (Err(error), _) => {
+ Ok(initial_buffer) => initial_buffer,
+ Err(error) => {
if error.as_code() == IggyError::ConnectionClosed.as_code() {
return Err(ConnectionError::from(error));
} else {
@@ -71,9 +71,9 @@ pub(crate) async fn handle_connection(
}
};
- if read_length != INITIAL_BYTES_LENGTH {
+ if initial_buffer.len() != INITIAL_BYTES_LENGTH {
sender.send_error_response(IggyError::CommandLengthError(format!(
- "Unable to read the TCP request length, expected: {INITIAL_BYTES_LENGTH} bytes, received: {read_length} bytes."
+ "Unable to read the TCP request length, expected: {INITIAL_BYTES_LENGTH} bytes, received: {} bytes.", initial_buffer.len()
))).await?;
continue;
}
@@ -81,8 +81,7 @@ pub(crate) async fn handle_connection(
let initial_buffer = initial_buffer.freeze();
let length =
u32::from_le_bytes(initial_buffer[0..INITIAL_BYTES_LENGTH].try_into().unwrap());
- let (res, code_buffer) = sender.read(code_buffer.clone()).await;
- let _ = res?;
+ let code_buffer = sender.read(code_buffer.clone()).await?;
let code_buffer = code_buffer.freeze();
let code: u32 =
u32::from_le_bytes(code_buffer[0..INITIAL_BYTES_LENGTH].try_into().unwrap());
diff --git a/core/server/src/tcp/sender.rs b/core/server/src/tcp/sender.rs
index 2c0ecfa9..f757a69b 100644
--- a/core/server/src/tcp/sender.rs
+++ b/core/server/src/tcp/sender.rs
@@ -23,31 +23,25 @@ use compio::{
io::{AsyncRead, AsyncReadAtExt, AsyncReadExt, AsyncWriteExt},
};
use iggy_common::IggyError;
-use nix::libc;
-use std::io::IoSlice;
use tracing::{debug, error};
use crate::streaming::utils::PooledBuffer;
const STATUS_OK: &[u8] = &[0; 4];
-pub(crate) async fn read<T, B>(stream: &mut T, buffer: B) -> (Result<usize, IggyError>, B)
+pub(crate) async fn read<T, B>(stream: &mut T, buffer: B) -> Result<B, IggyError>
where
T: AsyncReadExt + AsyncWriteExt + Unpin,
B: IoBufMut,
{
let BufResult(result, buffer) = stream.read_exact(buffer).await;
- match (result, buffer) {
- (Ok(_), buffer) => (Ok(buffer.buf_len()), buffer),
- // TODO: How to handle this ?(Ok(0), buffer) => (Err(IggyError::ConnectionClosed), buffer),
- // `read_exact` from compio doesn't return how many bytes it read.
- (Err(error), buffer) => {
+ match result {
+ Ok(_) => Ok(buffer),
+ Err(error) => {
if error.kind() == std::io::ErrorKind::UnexpectedEof {
- //error!("Got some error tho.. {}", error);
- (Err(IggyError::ConnectionClosed), buffer)
+ Err(IggyError::ConnectionClosed)
} else {
- //error!("Got some other error tho.. {}", error);
- (Err(IggyError::TcpError), buffer)
+ Err(IggyError::TcpError)
}
}
}
diff --git a/core/server/src/tcp/tcp_sender.rs b/core/server/src/tcp/tcp_sender.rs
index d0a1c27c..ccdc95d2 100644
--- a/core/server/src/tcp/tcp_sender.rs
+++ b/core/server/src/tcp/tcp_sender.rs
@@ -20,13 +20,11 @@ use crate::binary::sender::Sender;
use crate::streaming::utils::PooledBuffer;
use crate::tcp::COMPONENT;
use crate::{server_error::ServerError, tcp::sender};
-use bytes::BytesMut;
-use compio::buf::{IoBuf, IoBufMut};
+use compio::buf::IoBufMut;
use compio::io::AsyncWrite;
use compio::net::TcpStream;
use error_set::ErrContext;
use iggy_common::IggyError;
-use nix::libc;
#[derive(Debug)]
pub struct TcpSender {
@@ -34,7 +32,7 @@ pub struct TcpSender {
}
impl Sender for TcpSender {
- async fn read<B: IoBufMut>(&mut self, buffer: B) -> (Result<usize, IggyError>, B) {
+ async fn read<B: IoBufMut>(&mut self, buffer: B) -> Result<B, IggyError> {
sender::read(&mut self.stream, buffer).await
}
diff --git a/core/server/src/tcp/tcp_tls_sender.rs b/core/server/src/tcp/tcp_tls_sender.rs
index 8dbb3972..f0dfcdf4 100644
--- a/core/server/src/tcp/tcp_tls_sender.rs
+++ b/core/server/src/tcp/tcp_tls_sender.rs
@@ -35,7 +35,7 @@ pub struct TcpTlsSender {
}
impl Sender for TcpTlsSender {
- async fn read<B: IoBufMut>(&mut self, buffer: B) -> (Result<usize, IggyError>, B) {
+ async fn read<B: IoBufMut>(&mut self, buffer: B) -> Result<B, IggyError> {
todo!();
sender::read(&mut self.stream, buffer).await
}