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/09/24 08:11:35 UTC
(iggy) 01/01: task2
This is an automated email from the ASF dual-hosted git repository.
hgruszecki pushed a commit to branch io_uring_tpc_task_2
in repository https://gitbox.apache.org/repos/asf/iggy.git
commit 72d0c54f9498e9e12191ba07727b3a88465e1e8d
Author: Hubert Gruszecki <h....@gmail.com>
AuthorDate: Tue Sep 23 18:41:14 2025 +0200
task2
---
core/common/src/error/iggy_error.rs | 3 +
.../commands/clean_personal_access_tokens.rs | 75 -----
core/server/src/channels/commands/mod.rs | 5 -
core/server/src/channels/commands/print_sysinfo.rs | 70 -----
core/server/src/channels/commands/save_messages.rs | 86 ------
.../src/channels/commands/verify_heartbeats.rs | 80 ------
core/server/src/main.rs | 33 ++-
core/server/src/quic/quic_server.rs | 15 +-
core/server/src/shard/background_tasks.rs | 157 +++++++++++
core/server/src/shard/builder.rs | 3 +-
core/server/src/shard/listener.rs | 85 ++++++
core/server/src/shard/listeners/mod.rs | 3 +
core/server/src/shard/listeners/tcp.rs | 308 +++++++++++++++++++++
core/server/src/shard/mod.rs | 131 ++++-----
core/server/src/shard/system_tasks.rs | 46 +++
core/server/src/shard/task_registry.rs | 108 --------
core/server/src/shard/tasks.rs | 285 +++++++++++++++++++
core/server/src/shard/tasks/messages.rs | 59 ----
core/server/src/shard/tasks/mod.rs | 1 -
core/server/src/tcp/tcp_listener.rs | 28 +-
core/server/src/tcp/tcp_tls_listener.rs | 40 ++-
21 files changed, 1011 insertions(+), 610 deletions(-)
diff --git a/core/common/src/error/iggy_error.rs b/core/common/src/error/iggy_error.rs
index 614757f1..35ab50d1 100644
--- a/core/common/src/error/iggy_error.rs
+++ b/core/common/src/error/iggy_error.rs
@@ -476,6 +476,9 @@ pub enum IggyError {
#[error("Cannot bind to socket with addr: {0}")]
CannotBindToSocket(String) = 12000,
+
+ #[error("Task shutdown failed")]
+ TaskShutdownError = 13000,
}
impl IggyError {
diff --git a/core/server/src/channels/commands/clean_personal_access_tokens.rs b/core/server/src/channels/commands/clean_personal_access_tokens.rs
deleted file mode 100644
index d846aa27..00000000
--- a/core/server/src/channels/commands/clean_personal_access_tokens.rs
+++ /dev/null
@@ -1,75 +0,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::rc::Rc;
-
-use crate::shard::IggyShard;
-use iggy_common::{IggyError, IggyTimestamp};
-use tracing::{debug, info, trace};
-
-pub async fn clear_personal_access_tokens(shard: Rc<IggyShard>) -> Result<(), IggyError> {
- let config = &shard.config.personal_access_token.cleaner;
- if !config.enabled {
- info!("Personal access token cleaner is disabled.");
- return Ok(());
- }
-
- info!(
- "Personal access token cleaner is enabled, expired tokens will be deleted every: {}.",
- config.interval
- );
-
- let interval = config.interval.get_duration();
- let mut interval_timer = compio::time::interval(interval);
-
- loop {
- interval_timer.tick().await;
- trace!("Cleaning expired personal access tokens...");
-
- let users = shard.users.borrow();
- let now = IggyTimestamp::now();
- let mut deleted_tokens_count = 0;
-
- for (_, user) in users.iter() {
- let expired_tokens = user
- .personal_access_tokens
- .iter()
- .filter(|token| token.is_expired(now))
- .map(|token| token.token.clone())
- .collect::<Vec<_>>();
-
- for token in expired_tokens {
- debug!(
- "Personal access token: {} for user with ID: {} is expired.",
- token, user.id
- );
- deleted_tokens_count += 1;
- user.personal_access_tokens.remove(&token);
- debug!(
- "Deleted personal access token: {} for user with ID: {}.",
- token, user.id
- );
- }
- }
-
- info!(
- "Deleted {} expired personal access tokens.",
- deleted_tokens_count
- );
- }
-}
diff --git a/core/server/src/channels/commands/mod.rs b/core/server/src/channels/commands/mod.rs
index 0130170b..31bd66e6 100644
--- a/core/server/src/channels/commands/mod.rs
+++ b/core/server/src/channels/commands/mod.rs
@@ -15,8 +15,3 @@
* specific language governing permissions and limitations
* under the License.
*/
-
-pub mod clean_personal_access_tokens;
-pub mod print_sysinfo;
-pub mod save_messages;
-pub mod verify_heartbeats;
diff --git a/core/server/src/channels/commands/print_sysinfo.rs b/core/server/src/channels/commands/print_sysinfo.rs
deleted file mode 100644
index 0c3992d2..00000000
--- a/core/server/src/channels/commands/print_sysinfo.rs
+++ /dev/null
@@ -1,70 +0,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::rc::Rc;
-
-use crate::shard::IggyShard;
-use crate::streaming::utils::memory_pool;
-use human_repr::HumanCount;
-use iggy_common::IggyError;
-use tracing::{error, info, trace};
-
-pub async fn print_sys_info(shard: Rc<IggyShard>) -> Result<(), IggyError> {
- let config = &shard.config.system.logging;
- let interval = config.sysinfo_print_interval;
-
- if interval.is_zero() {
- info!("SysInfoPrinter is disabled.");
- return Ok(());
- }
- info!("SysInfoPrinter is enabled, system information will be printed every {interval}.");
- let mut interval_timer = compio::time::interval(interval.get_duration());
- loop {
- interval_timer.tick().await;
- trace!("Printing system information...");
-
- let stats = match shard.get_stats().await {
- Ok(stats) => stats,
- Err(e) => {
- error!("Failed to get system information. Error: {e}");
- continue;
- }
- };
-
- let free_memory_percent = (stats.available_memory.as_bytes_u64() as f64
- / stats.total_memory.as_bytes_u64() as f64)
- * 100f64;
-
- info!(
- "CPU: {:.2}%/{:.2}% (IggyUsage/Total), Mem: {:.2}%/{}/{}/{} (Free/IggyUsage/TotalUsed/Total), Clients: {}, Messages processed: {}, Read: {}, Written: {}, Uptime: {}",
- stats.cpu_usage,
- stats.total_cpu_usage,
- free_memory_percent,
- stats.memory_usage,
- stats.total_memory - stats.available_memory,
- stats.total_memory,
- stats.clients_count,
- stats.messages_count.human_count_bare(),
- stats.read_bytes,
- stats.written_bytes,
- stats.run_time
- );
-
- memory_pool().log_stats();
- }
-}
diff --git a/core/server/src/channels/commands/save_messages.rs b/core/server/src/channels/commands/save_messages.rs
deleted file mode 100644
index 84330563..00000000
--- a/core/server/src/channels/commands/save_messages.rs
+++ /dev/null
@@ -1,86 +0,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 crate::{shard::IggyShard, shard_info};
-use iggy_common::{Identifier, IggyError};
-use std::rc::Rc;
-use tracing::{error, info, trace};
-
-pub async fn save_messages(shard: Rc<IggyShard>) -> Result<(), IggyError> {
- let config = &shard.config.message_saver;
- if !config.enabled {
- info!("Message saver is disabled.");
- return Ok(());
- }
-
- // TODO: Maybe we should get rid of it in order to not complicate, and use the fsync settings per partition from config.
- let enforce_fsync = config.enforce_fsync;
- let interval = config.interval;
- info!(
- "Message saver is enabled, buffered messages will be automatically saved every: {interval}, enforce fsync: {enforce_fsync}."
- );
-
- let mut interval_timer = compio::time::interval(interval.get_duration());
- loop {
- interval_timer.tick().await;
- trace!("Saving buffered messages...");
-
- let namespaces = shard.get_current_shard_namespaces();
- let mut total_saved_messages = 0u32;
- let reason = "background saver triggered".to_string();
-
- for ns in namespaces {
- let stream_id = Identifier::numeric(ns.stream_id() as u32).unwrap();
- let topic_id = Identifier::numeric(ns.topic_id() as u32).unwrap();
- let partition_id = ns.partition_id();
-
- match shard
- .streams2
- .persist_messages(
- shard.id,
- &stream_id,
- &topic_id,
- partition_id,
- reason.clone(),
- &shard.config.system,
- )
- .await
- {
- Ok(batch_count) => {
- total_saved_messages += batch_count;
- }
- Err(err) => {
- error!(
- "Failed to save messages for partition {}: {}",
- partition_id, err
- );
- }
- }
- }
-
- if total_saved_messages > 0 {
- shard_info!(
- shard.id,
- "Saved {} buffered messages on disk.",
- total_saved_messages
- );
- }
-
- trace!("Finished saving buffered messages.");
- }
-}
diff --git a/core/server/src/channels/commands/verify_heartbeats.rs b/core/server/src/channels/commands/verify_heartbeats.rs
deleted file mode 100644
index 123c174f..00000000
--- a/core/server/src/channels/commands/verify_heartbeats.rs
+++ /dev/null
@@ -1,80 +0,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 crate::shard::IggyShard;
-use iggy_common::{IggyDuration, IggyError, IggyTimestamp};
-use std::rc::Rc;
-use tracing::{debug, info, trace, warn};
-
-const MAX_THRESHOLD: f64 = 1.2;
-
-pub async fn verify_heartbeats(shard: Rc<IggyShard>) -> Result<(), IggyError> {
- let config = &shard.config.heartbeat;
- if !config.enabled {
- info!("Heartbeats verification is disabled.");
- return Ok(());
- }
-
- let interval = config.interval;
- let max_interval = IggyDuration::from((MAX_THRESHOLD * interval.as_micros() as f64) as u64);
- info!("Heartbeats will be verified every: {interval}. Max allowed interval: {max_interval}.");
-
- let mut interval_timer = compio::time::interval(interval.get_duration());
-
- loop {
- interval_timer.tick().await;
- trace!("Verifying heartbeats...");
-
- let clients = {
- let client_manager = shard.client_manager.borrow();
- client_manager.get_clients()
- };
-
- let now = IggyTimestamp::now();
- let heartbeat_to = IggyTimestamp::from(now.as_micros() - max_interval.as_micros());
- debug!("Verifying heartbeats at: {now}, max allowed timestamp: {heartbeat_to}");
-
- let mut stale_clients = Vec::new();
- for client in clients {
- if client.last_heartbeat.as_micros() < heartbeat_to.as_micros() {
- warn!(
- "Stale client session: {}, last heartbeat at: {}, max allowed timestamp: {heartbeat_to}",
- client.session, client.last_heartbeat,
- );
- client.session.set_stale();
- stale_clients.push(client.session.client_id);
- } else {
- debug!(
- "Valid heartbeat at: {} for client session: {}, max allowed timestamp: {heartbeat_to}",
- client.last_heartbeat, client.session,
- );
- }
- }
-
- if stale_clients.is_empty() {
- continue;
- }
-
- let count = stale_clients.len();
- info!("Removing {count} stale clients...");
- for client_id in stale_clients {
- shard.delete_client(client_id);
- }
- info!("Removed {count} stale clients.");
- }
-}
diff --git a/core/server/src/main.rs b/core/server/src/main.rs
index 05e4ecc8..d4cb0667 100644
--- a/core/server/src/main.rs
+++ b/core/server/src/main.rs
@@ -351,26 +351,25 @@ async fn main() -> Result<(), ServerError> {
handles.push(handle);
}
- let shutdown_handles_for_signal = shutdown_handles.clone();
- /*
- ::set_handler(move || {
- info!("Received shutdown signal (SIGTERM/SIGINT), initiating graceful shutdown...");
-
- for (shard_id, stop_sender) in &shutdown_handles_for_signal {
- info!("Sending shutdown signal to shard {}", shard_id);
- if let Err(e) = stop_sender.send_blocking(()) {
- error!(
- "Failed to send shutdown signal to shard {}: {}",
- shard_id, e
- );
- }
+ // Set up signal handler to directly trigger shard shutdowns
+ ctrlc::set_handler(move || {
+ info!("Received shutdown signal (SIGTERM/SIGINT), initiating graceful shutdown...");
+ for (shard_id, stop_sender) in &shutdown_handles {
+ info!("Sending shutdown signal to shard {}", shard_id);
+ if let Err(e) = stop_sender.try_send(()) {
+ error!(
+ "Failed to send shutdown signal to shard {}: {}",
+ shard_id, e
+ );
}
- })
- .expect("Error setting Ctrl-C handler");
- */
+ }
+ })
+ .expect("Error setting Ctrl-C handler");
info!("Iggy server is running. Press Ctrl+C or send SIGTERM to shutdown.");
- for (idx, handle) in handles.into_iter().enumerate() {
+
+ // Wait for all shard threads to complete
+ for handle in handles {
handle.join().expect("Failed to join shard thread");
}
diff --git a/core/server/src/quic/quic_server.rs b/core/server/src/quic/quic_server.rs
index 6a0255ad..ef37c739 100644
--- a/core/server/src/quic/quic_server.rs
+++ b/core/server/src/quic/quic_server.rs
@@ -36,9 +36,8 @@ use crate::quic::{COMPONENT, listener, quic_socket};
use crate::server_error::QuicError;
use crate::shard::IggyShard;
-/// Starts the QUIC server.
-/// Returns the address the server is listening on.
-pub async fn span_quic_server(shard: Rc<IggyShard>) -> Result<(), iggy_common::IggyError> {
+/// Creates a QUIC endpoint for the shard
+pub async fn create_endpoint(shard: &Rc<IggyShard>) -> Result<Endpoint, iggy_common::IggyError> {
// Ensure rustls crypto provider is installed (thread-safe, idempotent)
if rustls::crypto::CryptoProvider::get_default().is_none() {
if let Err(e) = default_provider().install_default() {
@@ -92,9 +91,17 @@ pub async fn span_quic_server(shard: Rc<IggyShard>) -> Result<(), iggy_common::I
iggy_common::IggyError::QuicError
})?;
+ Ok(endpoint)
+}
+
+/// Starts the QUIC server.
+/// Returns the address the server is listening on.
+pub async fn span_quic_server(shard: Rc<IggyShard>) -> Result<(), iggy_common::IggyError> {
+ let endpoint = create_endpoint(&shard).await?;
+
let actual_addr = endpoint.local_addr().map_err(|e| {
error!("Failed to get local address: {e}");
- iggy_common::IggyError::CannotBindToSocket(addr.to_string())
+ iggy_common::IggyError::CannotBindToSocket("QUIC".to_string())
})?;
info!(
diff --git a/core/server/src/shard/background_tasks.rs b/core/server/src/shard/background_tasks.rs
new file mode 100644
index 00000000..d4b3de4e
--- /dev/null
+++ b/core/server/src/shard/background_tasks.rs
@@ -0,0 +1,157 @@
+use iggy_common::{Identifier, IggyError, IggyTimestamp};
+use std::rc::Rc;
+use tracing::{debug, error, info, trace, warn};
+
+use super::IggyShard;
+
+pub async fn save_messages(shard: Rc<IggyShard>) -> Result<(), IggyError> {
+ let config = &shard.config.message_saver;
+ if !config.enabled {
+ return Ok(());
+ }
+
+ trace!("Saving buffered messages...");
+
+ let namespaces = shard.get_current_shard_namespaces();
+ let mut total_saved_messages = 0u32;
+ let reason = "background saver triggered".to_string();
+
+ for ns in namespaces {
+ let stream_id = Identifier::numeric(ns.stream_id() as u32).unwrap();
+ let topic_id = Identifier::numeric(ns.topic_id() as u32).unwrap();
+ let partition_id = ns.partition_id();
+
+ match shard
+ .streams2
+ .persist_messages(
+ shard.id,
+ &stream_id,
+ &topic_id,
+ partition_id,
+ reason.clone(),
+ &shard.config.system,
+ )
+ .await
+ {
+ Ok(batch_count) => {
+ total_saved_messages += batch_count;
+ }
+ Err(err) => {
+ error!(
+ "Failed to save messages for partition {}: {}",
+ partition_id, err
+ );
+ }
+ }
+ }
+
+ if total_saved_messages > 0 {
+ info!("Saved {} buffered messages on disk.", total_saved_messages);
+ }
+
+ Ok(())
+}
+
+pub async fn verify_heartbeats(shard: Rc<IggyShard>) -> Result<(), IggyError> {
+ const MAX_THRESHOLD: f64 = 1.2;
+
+ let config = &shard.config.heartbeat;
+ if !config.enabled {
+ return Ok(());
+ }
+
+ let interval = config.interval;
+ let max_interval =
+ iggy_common::IggyDuration::from((MAX_THRESHOLD * interval.as_micros() as f64) as u64);
+
+ trace!("Verifying heartbeats...");
+
+ let clients = {
+ let client_manager = shard.client_manager.borrow();
+ client_manager.get_clients()
+ };
+
+ let current_timestamp = IggyTimestamp::now();
+
+ for client in clients {
+ let last_heartbeat = client.last_heartbeat;
+ let elapsed_time = current_timestamp.as_micros() - last_heartbeat.as_micros();
+ if elapsed_time > max_interval.as_micros() {
+ warn!(
+ "Client: {} heartbeat has not been received for: {} μs, the max allowed interval is: {} μs",
+ client.session.client_id,
+ elapsed_time,
+ max_interval.as_micros()
+ );
+ let mut client_manager = shard.client_manager.borrow_mut();
+ client_manager.delete_client(client.session.client_id);
+ }
+ }
+
+ Ok(())
+}
+
+pub async fn clean_personal_access_tokens(shard: Rc<IggyShard>) -> Result<(), IggyError> {
+ let config = &shard.config.personal_access_token.cleaner;
+ if !config.enabled {
+ return Ok(());
+ }
+
+ trace!("Cleaning expired personal access tokens...");
+
+ let users = shard.users.borrow();
+ let current_timestamp = IggyTimestamp::now();
+ let mut total_deleted_tokens = 0u32;
+
+ for user in users.values() {
+ let tokens_to_delete: Vec<_> = user
+ .personal_access_tokens
+ .iter()
+ .filter(|token| {
+ if let Some(expiry) = token.expiry_at {
+ current_timestamp.as_micros() > expiry.as_micros()
+ } else {
+ false
+ }
+ })
+ .map(|token| token.name.clone())
+ .collect();
+
+ for token_name in tokens_to_delete {
+ if let Err(err) = shard.delete_personal_access_token_bypass_auth(user.id, &token_name) {
+ error!(
+ "Failed to delete expired personal access token: {} for user: {}, error: {}",
+ token_name, user.id, err
+ );
+ } else {
+ debug!(
+ "Deleted expired personal access token: {} for user: {}",
+ token_name, user.id
+ );
+ total_deleted_tokens += 1;
+ }
+ }
+ }
+
+ if total_deleted_tokens > 0 {
+ info!(
+ "Deleted {} expired personal access tokens.",
+ total_deleted_tokens
+ );
+ }
+
+ Ok(())
+}
+
+pub async fn print_sysinfo(shard: Rc<IggyShard>) -> Result<(), IggyError> {
+ if shard.config.system.logging.sysinfo_print_interval.is_zero() || shard.id != 0 {
+ return Ok(());
+ }
+
+ trace!("Printing system information...");
+
+ let stats = shard.get_stats().await?;
+ info!("{:?}", stats);
+
+ Ok(())
+}
diff --git a/core/server/src/shard/builder.rs b/core/server/src/shard/builder.rs
index fc203b77..51f54af7 100644
--- a/core/server/src/shard/builder.rs
+++ b/core/server/src/shard/builder.rs
@@ -30,7 +30,7 @@ use tracing::info;
use crate::{
configs::server::ServerConfig,
io::storage::Storage,
- shard::{Shard, ShardInfo, namespace::IggyNamespace, task_registry::TaskRegistry},
+ shard::{Shard, ShardInfo, namespace::IggyNamespace},
slab::streams::Streams,
state::{StateKind, system::SystemState},
streaming::{
@@ -160,7 +160,6 @@ impl IggyShardBuilder {
stop_sender: stop_sender,
messages_receiver: Cell::new(Some(frame_receiver)),
metrics: metrics,
- task_registry: TaskRegistry::new(),
is_shutting_down: AtomicBool::new(false),
tcp_bound_address: Cell::new(None),
quic_bound_address: Cell::new(None),
diff --git a/core/server/src/shard/listener.rs b/core/server/src/shard/listener.rs
new file mode 100644
index 00000000..7805192f
--- /dev/null
+++ b/core/server/src/shard/listener.rs
@@ -0,0 +1,85 @@
+use async_channel::{Receiver, Sender, bounded};
+use iggy_common::IggyError;
+use std::cell::RefCell;
+use std::rc::Rc;
+
+use super::IggyShard;
+
+pub trait Listener {
+ fn name(&self) -> &'static str;
+
+ async fn run(
+ &mut self,
+ shard: Rc<IggyShard>,
+ shutdown: Receiver<()>,
+ connections: Rc<RefCell<Vec<ConnectionHandle>>>,
+ is_shutting_down: Rc<RefCell<bool>>,
+ ) -> Result<(), IggyError>;
+}
+
+pub struct ConnectionHandle {
+ client_id: u32,
+ shutdown_tx: Sender<()>,
+ closed_rx: Receiver<()>,
+}
+
+impl ConnectionHandle {
+ pub fn new(client_id: u32) -> (Self, ConnectionController) {
+ let (shutdown_tx, shutdown_rx) = bounded(1);
+ let (closed_tx, closed_rx) = bounded(1);
+
+ let handle = Self {
+ client_id,
+ shutdown_tx,
+ closed_rx,
+ };
+
+ let controller = ConnectionController {
+ client_id,
+ shutdown_rx,
+ closed_tx,
+ };
+
+ (handle, controller)
+ }
+
+ pub fn client_id(&self) -> u32 {
+ self.client_id
+ }
+
+ pub fn request_shutdown(&self) {
+ let _ = self.shutdown_tx.try_send(());
+ }
+
+ pub async fn wait_closed(self) -> Result<(), IggyError> {
+ self.closed_rx
+ .recv()
+ .await
+ .map_err(|_| IggyError::ConnectionClosed)?;
+ Ok(())
+ }
+}
+
+pub struct ConnectionController {
+ client_id: u32,
+ shutdown_rx: Receiver<()>,
+ closed_tx: Sender<()>,
+}
+
+impl ConnectionController {
+ pub fn client_id(&self) -> u32 {
+ self.client_id
+ }
+
+ pub fn is_shutdown_requested(&self) -> bool {
+ self.shutdown_rx.try_recv().is_ok()
+ }
+
+ pub async fn wait_shutdown(&self) {
+ let _ = self.shutdown_rx.recv().await;
+ }
+
+ pub fn notify_closed(self) {
+ let _ = self.closed_tx.try_send(());
+ }
+}
diff --git a/core/server/src/shard/listeners/mod.rs b/core/server/src/shard/listeners/mod.rs
new file mode 100644
index 00000000..3b311a97
--- /dev/null
+++ b/core/server/src/shard/listeners/mod.rs
@@ -0,0 +1,3 @@
+pub mod tcp;
+
+pub use tcp::TcpListener;
diff --git a/core/server/src/shard/listeners/tcp.rs b/core/server/src/shard/listeners/tcp.rs
new file mode 100644
index 00000000..09bea986
--- /dev/null
+++ b/core/server/src/shard/listeners/tcp.rs
@@ -0,0 +1,308 @@
+use async_channel::Receiver;
+use bytes::BytesMut;
+use compio::net::{TcpListener as CompioTcpListener, TcpOpts, TcpStream};
+use error_set::ErrContext;
+use futures::FutureExt;
+use iggy_common::{IggyError, TransportProtocol};
+use std::cell::RefCell;
+use std::net::SocketAddr;
+use std::rc::Rc;
+use std::time::Duration;
+use tracing::{debug, error, info, warn};
+
+use crate::binary::command::{ServerCommand, ServerCommandHandler};
+use crate::binary::sender::SenderKind;
+use crate::configs::tcp::TcpSocketConfig;
+use crate::shard::listener::{ConnectionController, ConnectionHandle, Listener};
+use crate::shard::transmission::event::ShardEvent;
+use crate::shard::{IggyShard, shard_error, shard_info};
+use crate::streaming::session::Session;
+
+const INITIAL_BYTES_LENGTH: usize = 4;
+
+pub struct TcpListener {
+ address: String,
+ config: TcpSocketConfig,
+}
+
+impl TcpListener {
+ pub fn new(address: String, config: TcpSocketConfig) -> Self {
+ Self { address, config }
+ }
+}
+
+impl Listener for TcpListener {
+ fn name(&self) -> &'static str {
+ "tcp_listener"
+ }
+
+ async fn run(
+ &mut self,
+ shard: Rc<IggyShard>,
+ shutdown: Receiver<()>,
+ connections: Rc<RefCell<Vec<ConnectionHandle>>>,
+ is_shutting_down: Rc<RefCell<bool>>,
+ ) -> Result<(), IggyError> {
+ let mut addr: SocketAddr = self
+ .address
+ .parse()
+ .map_err(|_| IggyError::InvalidServerAddress)?;
+
+ if shard.id != 0 && addr.port() == 0 {
+ shard_info!(shard.id, "Waiting for TCP address from shard 0...");
+ loop {
+ futures::select! {
+ _ = shutdown.recv().fuse() => {
+ return Ok(());
+ }
+ _ = compio::time::sleep(Duration::from_millis(10)).fuse() => {
+ if let Some(bound_addr) = shard.tcp_bound_address.get() {
+ addr = bound_addr;
+ shard_info!(shard.id, "Received TCP address: {}", addr);
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ let listener = create_tcp_listener(addr, &self.config)
+ .await
+ .map_err(|_| IggyError::CannotBindToSocket(addr.to_string()))
+ .with_error_context(|err| {
+ format!("Failed to bind TCP server to address: {addr}, {err}")
+ })?;
+
+ let actual_addr = listener.local_addr().map_err(|e| {
+ shard_error!(shard.id, "Failed to get local address: {e}");
+ IggyError::CannotBindToSocket(addr.to_string())
+ })?;
+
+ shard_info!(shard.id, "TCP server has started on: {:?}", actual_addr);
+
+ if shard.id == 0 {
+ if addr.port() == 0 {
+ let event = ShardEvent::TcpBound {
+ address: actual_addr,
+ };
+ shard.broadcast_event_to_all_shards(event).await;
+ }
+
+ let mut current_config = shard.config.clone();
+ current_config.tcp.address = actual_addr.to_string();
+
+ let runtime_path = current_config.system.get_runtime_path();
+ let current_config_path = format!("{runtime_path}/current_config.toml");
+ let current_config_content =
+ toml::to_string(¤t_config).expect("Cannot serialize current_config");
+
+ let buf_result = compio::fs::write(¤t_config_path, current_config_content).await;
+ match buf_result.0 {
+ Ok(_) => shard_info!(
+ shard.id,
+ "Current config written to: {}",
+ current_config_path
+ ),
+ Err(e) => shard_error!(
+ shard.id,
+ "Failed to write current config to {}: {}",
+ current_config_path,
+ e
+ ),
+ }
+ }
+
+ loop {
+ futures::select! {
+ _ = shutdown.recv().fuse() => {
+ shard_info!(shard.id, "TCP server shutting down, no longer accepting connections");
+ break;
+ }
+
+ result = listener.accept().fuse() => {
+ if *is_shutting_down.borrow() {
+ if let Ok((stream, address)) = result {
+ shard_info!(shard.id, "Rejecting new connection from {} during shutdown", address);
+ drop(stream);
+ }
+ continue;
+ }
+
+ match result {
+ Ok((stream, address)) => {
+ shard_info!(shard.id, "Accepted new TCP connection: {}", address);
+ self.handle_new_connection(stream, address, &shard, &connections).await;
+ }
+ Err(error) => shard_error!(shard.id, "Unable to accept TCP socket: {}", error),
+ }
+ }
+ }
+ }
+
+ Ok(())
+ }
+}
+
+impl TcpListener {
+ async fn handle_new_connection(
+ &self,
+ stream: TcpStream,
+ addr: SocketAddr,
+ shard: &Rc<IggyShard>,
+ connections: &Rc<RefCell<Vec<ConnectionHandle>>>,
+ ) {
+ let transport = TransportProtocol::Tcp;
+ let session = shard.add_client(&addr, transport);
+ let client_id = session.client_id;
+
+ shard_info!(
+ shard.id,
+ "Added {} client with session: {} for IP address: {}",
+ transport,
+ session,
+ addr
+ );
+
+ shard.add_active_session(session.clone());
+ let event = ShardEvent::NewSession {
+ address: addr,
+ transport,
+ };
+ let _responses = shard.broadcast_event_to_all_shards(event).await;
+
+ let (handle, controller) = ConnectionHandle::new(client_id);
+ connections.borrow_mut().push(handle);
+
+ let sender = SenderKind::get_tcp_sender(stream);
+ let shard_clone = Rc::clone(shard);
+ let connections_clone = connections.clone();
+
+ compio::runtime::spawn(async move {
+ info!("TCP connection {} started", client_id);
+
+ handle_tcp_connection(session, sender, shard_clone, controller).await;
+
+ // Clean up connection
+ connections_clone
+ .borrow_mut()
+ .retain(|h| h.client_id() != client_id);
+ info!("TCP connection {} ended", client_id);
+ });
+ }
+}
+
+async fn handle_tcp_connection(
+ session: Rc<Session>,
+ mut sender: SenderKind,
+ shard: Rc<IggyShard>,
+ controller: ConnectionController,
+) {
+ let client_id = controller.client_id();
+ let mut length_buffer = BytesMut::with_capacity(INITIAL_BYTES_LENGTH);
+ let mut code_buffer = BytesMut::with_capacity(INITIAL_BYTES_LENGTH);
+
+ loop {
+ if controller.is_shutdown_requested() {
+ info!("Connection {} received shutdown signal", client_id);
+ let _ = sender.send_error_response(IggyError::Disconnected).await;
+ break;
+ }
+
+ let read_future = sender.read(length_buffer.clone());
+ let (result, mut initial_buffer) = futures::select! {
+ _ = controller.wait_shutdown().fuse() => {
+ info!("Connection stop signal received for session: {}", session);
+ let _ = sender.send_error_response(IggyError::Disconnected).await;
+ break;
+ }
+ result = read_future.fuse() => {
+ result
+ }
+ };
+
+ match result {
+ Ok(_) => {}
+ Err(error) => {
+ if error.as_code() == IggyError::ConnectionClosed.as_code() {
+ break;
+ } else {
+ error!("Connection error: {:?}", error);
+ let _ = sender.send_error_response(error).await;
+ continue;
+ }
+ }
+ }
+
+ let length =
+ u32::from_le_bytes(initial_buffer[0..INITIAL_BYTES_LENGTH].try_into().unwrap());
+ let (res, mut code_buffer_out) = sender.read(code_buffer.clone()).await;
+ if let Err(e) = res {
+ error!("Failed to read command code: {}", e);
+ break;
+ }
+ let code = u32::from_le_bytes(code_buffer_out[0..INITIAL_BYTES_LENGTH].try_into().unwrap());
+
+ initial_buffer.clear();
+ code_buffer_out.clear();
+ length_buffer = initial_buffer;
+ code_buffer = code_buffer_out;
+
+ debug!("Received TCP request, length: {}, code: {}", length, code);
+
+ let command = match ServerCommand::from_code_and_reader(code, &mut sender, length - 4).await
+ {
+ Ok(cmd) => cmd,
+ Err(e) => {
+ error!("Failed to parse command: {}", e);
+ let _ = sender.send_error_response(e).await;
+ continue;
+ }
+ };
+
+ match command.handle(&mut sender, length, &session, &shard).await {
+ Ok(_) => {
+ debug!("Command handled successfully for session: {}", session);
+ }
+ Err(error) => {
+ error!("Command failed for session: {}, error: {}", session, error);
+ if let IggyError::ClientNotFound(_) = error {
+ let _ = sender.send_error_response(error).await;
+ break;
+ } else {
+ let _ = sender.send_error_response(error).await;
+ }
+ }
+ }
+ }
+
+ controller.notify_closed();
+}
+
+async fn create_tcp_listener(
+ addr: SocketAddr,
+ config: &TcpSocketConfig,
+) -> Result<CompioTcpListener, std::io::Error> {
+ let opts = TcpOpts::new().reuse_port(true).reuse_port(true);
+ let opts = if config.override_defaults {
+ let recv_buffer_size = config
+ .recv_buffer_size
+ .as_bytes_u64()
+ .try_into()
+ .expect("Failed to parse recv_buffer_size for TCP socket");
+
+ let send_buffer_size = config
+ .send_buffer_size
+ .as_bytes_u64()
+ .try_into()
+ .expect("Failed to parse send_buffer_size for TCP socket");
+
+ opts.recv_buffer_size(recv_buffer_size)
+ .send_buffer_size(send_buffer_size)
+ .keepalive(config.keepalive)
+ .linger(config.linger.get_duration())
+ .nodelay(config.nodelay)
+ } else {
+ opts
+ };
+ CompioTcpListener::bind_with_options(addr, opts).await
+}
diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs
index 4771d238..9c05884c 100644
--- a/core/server/src/shard/mod.rs
+++ b/core/server/src/shard/mod.rs
@@ -16,12 +16,15 @@ inner() * or more contributor license agreements. See the NOTICE file
* under the License.
*/
+pub mod background_tasks;
pub mod builder;
+pub mod listener;
+pub mod listeners;
pub mod logging;
pub mod namespace;
pub mod stats;
pub mod system;
-pub mod task_registry;
+pub mod system_tasks;
pub mod tasks;
pub mod transmission;
@@ -29,7 +32,6 @@ use ahash::{AHashMap, AHashSet, HashMap};
use builder::IggyShardBuilder;
use dashmap::DashMap;
use error_set::ErrContext;
-use futures::future::try_join_all;
use hash32::{Hasher, Murmur3Hasher};
use iggy_common::{
EncryptorKind, Identifier, IggyError, IggyTimestamp, Permissions, PollingKind, UserId,
@@ -57,12 +59,9 @@ use transmission::connector::{Receiver, ShardConnector, StopReceiver, StopSender
use crate::{
binary::handlers::messages::poll_messages_handler::IggyPollMetadata,
configs::server::ServerConfig,
- http::http_server,
io::fs_utils,
shard::{
namespace::{IggyFullNamespace, IggyNamespace},
- task_registry::TaskRegistry,
- tasks::messages::spawn_shard_message_task,
transmission::{
event::ShardEvent,
frame::{ShardFrame, ShardResponse},
@@ -92,7 +91,6 @@ use crate::{
users::{permissioner::Permissioner, user::User},
utils::ptr::EternalPtr,
},
- tcp::tcp_server::spawn_tcp_server,
versioning::SemanticVersion,
};
@@ -101,8 +99,6 @@ pub const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
static USER_ID: AtomicU32 = AtomicU32::new(1);
-type Task = Pin<Box<dyn Future<Output = Result<(), IggyError>>>>;
-
pub(crate) struct Shard {
id: u16,
connection: ShardConnector<ShardFrame>,
@@ -170,7 +166,6 @@ pub struct IggyShard {
pub messages_receiver: Cell<Option<Receiver<ShardFrame>>>,
pub(crate) stop_receiver: StopReceiver,
pub(crate) stop_sender: StopSender,
- pub(crate) task_registry: TaskRegistry,
pub(crate) is_shutting_down: AtomicBool,
pub(crate) tcp_bound_address: Cell<Option<SocketAddr>>,
pub(crate) quic_bound_address: Cell<Option<SocketAddr>>,
@@ -220,7 +215,6 @@ impl IggyShard {
messages_receiver: Cell::new(None),
stop_receiver,
stop_sender,
- task_registry: TaskRegistry::new(),
is_shutting_down: AtomicBool::new(false),
tcp_bound_address: Cell::new(None),
quic_bound_address: Cell::new(None),
@@ -244,76 +238,86 @@ impl IggyShard {
}
pub async fn run(self: &Rc<Self>) -> Result<(), IggyError> {
- // Workaround to ensure that the statistics are initialized before the server
- // loads streams and starts accepting connections. This is necessary to
- // have the correct statistics when the server starts.
let now = Instant::now();
self.get_stats().await?;
shard_info!(self.id, "Starting...");
self.init().await?;
- // TODO: Fixme
- //self.assert_init();
- // Create all tasks (tcp listener, http listener, command processor, in the future also the background jobs).
- let mut tasks: Vec<Task> = vec![Box::pin(spawn_shard_message_task(self.clone()))];
- if self.config.tcp.enabled {
- tasks.push(Box::pin(spawn_tcp_server(self.clone())));
+ let task_manager = tasks::TaskManager::new(self.clone());
+
+ // System tasks (shutdown last)
+ task_manager.spawn_system("shard_receiver", |shard, shutdown| {
+ system_tasks::receive_shard_messages(shard, shutdown)
+ });
+
+ // Background tasks
+ if self.config.message_saver.enabled {
+ task_manager.spawn_background(
+ "message_saver",
+ self.config.message_saver.interval.get_duration(),
+ |shard| background_tasks::save_messages(shard),
+ );
}
- if self.config.http.enabled && self.id == 0 {
- println!("Starting HTTP server on shard: {}", self.id);
- tasks.push(Box::pin(http_server::start(
- self.config.http.clone(),
- self.clone(),
- )));
+ if self.config.heartbeat.enabled {
+ task_manager.spawn_background(
+ "heartbeat_verifier",
+ self.config.heartbeat.interval.get_duration(),
+ |shard| background_tasks::verify_heartbeats(shard),
+ );
}
- if self.config.quic.enabled {
- tasks.push(Box::pin(crate::quic::quic_server::span_quic_server(
- self.clone(),
- )));
+ if self.config.personal_access_token.cleaner.enabled {
+ task_manager.spawn_background(
+ "token_cleaner",
+ self.config
+ .personal_access_token
+ .cleaner
+ .interval
+ .get_duration(),
+ |shard| background_tasks::clean_personal_access_tokens(shard),
+ );
}
- tasks.push(Box::pin(
- crate::channels::commands::clean_personal_access_tokens::clear_personal_access_tokens(
- self.clone(),
- ),
- ));
- // TOOD: Fixme, not always id 0 is the first shard.
- if self.id == 0 {
- tasks.push(Box::pin(
- crate::channels::commands::print_sysinfo::print_sys_info(self.clone()),
+ if !self.config.system.logging.sysinfo_print_interval.is_zero() && self.id == 0 {
+ task_manager.spawn_background(
+ "sysinfo_printer",
+ self.config
+ .system
+ .logging
+ .sysinfo_print_interval
+ .get_duration(),
+ |shard| background_tasks::print_sysinfo(shard),
+ );
+ }
+
+ // Listeners (shutdown first)
+ if self.config.tcp.enabled {
+ task_manager.spawn_listener(listeners::TcpListener::new(
+ self.config.tcp.address.clone(),
+ self.config.tcp.socket.clone(),
));
}
- tasks.push(Box::pin(
- crate::channels::commands::verify_heartbeats::verify_heartbeats(self.clone()),
- ));
- tasks.push(Box::pin(
- crate::channels::commands::save_messages::save_messages(self.clone()),
- ));
+ let elapsed = now.elapsed();
+ shard_info!(self.id, "Initialized in {} ms.", elapsed.as_millis());
- let stop_receiver = self.get_stop_receiver();
- let shard_for_shutdown = self.clone();
+ // Wait for shutdown signal
+ self.get_stop_receiver().recv().await;
- /*
- compio::runtime::spawn(async move {
- let _ = stop_receiver.recv().await;
- info!("Shard {} received shutdown signal", shard_for_shutdown.id);
+ shard_info!(self.id, "Received shutdown signal");
- let shutdown_success = shard_for_shutdown.trigger_shutdown().await;
- if !shutdown_success {
- shard_error!(shard_for_shutdown.id, "shutdown timed out");
- } else {
- shard_info!(shard_for_shutdown.id, "shutdown completed successfully");
- }
- });
- */
+ // Orchestrated shutdown
+ let shutdown_success = task_manager
+ .shutdown_all(Duration::from_secs(10))
+ .await
+ .is_ok();
+ if !shutdown_success {
+ shard_error!(self.id, "shutdown timed out");
+ } else {
+ shard_info!(self.id, "shutdown completed successfully");
+ }
- let elapsed = now.elapsed();
- shard_info!(self.id, "Initialized in {} ms.", elapsed.as_millis());
- let result = try_join_all(tasks).await;
- result?;
Ok(())
}
@@ -421,7 +425,8 @@ impl IggyShard {
pub async fn trigger_shutdown(&self) -> bool {
self.is_shutting_down.store(true, Ordering::SeqCst);
info!("Shard {} shutdown state set", self.id);
- self.task_registry.shutdown_all(SHUTDOWN_TIMEOUT).await
+ // TODO: Shutdown is now handled by TaskManager in run() method
+ true
}
pub fn get_available_shards_count(&self) -> u32 {
diff --git a/core/server/src/shard/system_tasks.rs b/core/server/src/shard/system_tasks.rs
new file mode 100644
index 00000000..027edcce
--- /dev/null
+++ b/core/server/src/shard/system_tasks.rs
@@ -0,0 +1,46 @@
+use async_channel::Receiver;
+use futures::{FutureExt, StreamExt};
+use iggy_common::IggyError;
+use std::rc::Rc;
+use tracing::info;
+
+use super::IggyShard;
+use super::transmission::frame::ShardFrame;
+
+pub async fn receive_shard_messages(
+ shard: Rc<IggyShard>,
+ shutdown: Receiver<()>,
+) -> Result<(), IggyError> {
+ let messages_receiver = shard.messages_receiver.take();
+ if messages_receiver.is_none() {
+ return Ok(());
+ }
+
+ let mut messages_receiver = messages_receiver.unwrap();
+
+ loop {
+ futures::select! {
+ _ = shutdown.recv().fuse() => {
+ info!("Shard message receiver shutting down");
+ shard.messages_receiver.replace(Some(messages_receiver));
+ break;
+ }
+ frame = messages_receiver.next().fuse() => {
+ if let Some(frame) = frame {
+ let ShardFrame { message, response_sender } = frame;
+ match (shard.handle_shard_message(message).await, response_sender) {
+ (Some(response), Some(response_sender)) => {
+ response_sender
+ .send(response)
+ .await
+ .expect("Failed to send response back to origin shard.");
+ }
+ _ => {}
+ };
+ }
+ }
+ }
+ }
+
+ Ok(())
+}
diff --git a/core/server/src/shard/task_registry.rs b/core/server/src/shard/task_registry.rs
deleted file mode 100644
index ff674a8e..00000000
--- a/core/server/src/shard/task_registry.rs
+++ /dev/null
@@ -1,108 +0,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 async_channel::{Receiver, Sender, bounded};
-use compio::runtime::JoinHandle;
-use futures::future::join_all;
-use std::cell::RefCell;
-use std::collections::HashMap;
-use std::future::Future;
-use std::time::Duration;
-use tracing::{error, info, warn};
-
-pub struct TaskRegistry {
- tasks: RefCell<Vec<JoinHandle<()>>>,
- active_connections: RefCell<HashMap<u32, Sender<()>>>,
-}
-
-impl TaskRegistry {
- pub fn new() -> Self {
- Self {
- tasks: RefCell::new(Vec::new()),
- active_connections: RefCell::new(HashMap::new()),
- }
- }
-
- pub fn spawn_tracked<F>(&self, future: F)
- where
- F: Future<Output = ()> + 'static,
- {
- let handle = compio::runtime::spawn(future);
- self.tasks.borrow_mut().push(handle);
- }
-
- pub fn add_connection(&self, client_id: u32) -> Receiver<()> {
- let (stop_sender, stop_receiver) = bounded(1);
- self.active_connections
- .borrow_mut()
- .insert(client_id, stop_sender);
- stop_receiver
- }
-
- pub fn remove_connection(&self, client_id: &u32) {
- self.active_connections.borrow_mut().remove(client_id);
- }
-
- pub async fn shutdown_all(&self, timeout: Duration) -> bool {
- info!("Initiating task registry shutdown");
-
- let connections = self.active_connections.borrow();
- for (client_id, stop_sender) in connections.iter() {
- info!("Sending shutdown signal to client {}", client_id);
- if let Err(e) = stop_sender.send(()).await {
- warn!(
- "Failed to send shutdown signal to client {}: {}",
- client_id, e
- );
- }
- }
- drop(connections);
-
- let tasks = self.tasks.take();
- let total = tasks.len();
-
- if total == 0 {
- info!("No tasks to shut down");
- return true;
- }
-
- let timeout_futures: Vec<_> = tasks
- .into_iter()
- .enumerate()
- .map(|(idx, handle)| async move {
- match compio::time::timeout(timeout, handle).await {
- Ok(_) => (idx, true),
- Err(_) => {
- warn!("Task {} did not complete within timeout", idx);
- (idx, false)
- }
- }
- })
- .collect();
-
- let results = join_all(timeout_futures).await;
- let completed = results.iter().filter(|(_, success)| *success).count();
-
- info!(
- "Task registry shutdown complete. {} of {} tasks completed",
- completed, total
- );
-
- completed == total
- }
-}
diff --git a/core/server/src/shard/tasks.rs b/core/server/src/shard/tasks.rs
new file mode 100644
index 00000000..75a7b5d8
--- /dev/null
+++ b/core/server/src/shard/tasks.rs
@@ -0,0 +1,285 @@
+use async_channel::{Receiver, Sender, bounded};
+use compio::runtime::JoinHandle;
+use futures::FutureExt;
+use iggy_common::IggyError;
+use std::cell::RefCell;
+use std::future::Future;
+use std::rc::Rc;
+use std::time::{Duration, Instant};
+use tracing::{debug, error, info, warn};
+
+use super::IggyShard;
+use super::listener::{ConnectionHandle, Listener};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum ShutdownPhase {
+ Listeners = 1,
+ Connections = 2,
+ DataPersistence = 3,
+ Background = 4,
+ System = 5,
+}
+
+struct Task {
+ name: &'static str,
+ phase: ShutdownPhase,
+ shutdown_tx: Sender<()>,
+ handle: JoinHandle<()>,
+}
+
+pub struct TaskManager {
+ shard: Rc<IggyShard>,
+ tasks: RefCell<Vec<Task>>,
+ connections: Rc<RefCell<Vec<ConnectionHandle>>>,
+ is_shutting_down: Rc<RefCell<bool>>,
+}
+
+impl TaskManager {
+ pub fn new(shard: Rc<IggyShard>) -> Self {
+ Self {
+ shard,
+ tasks: RefCell::new(Vec::new()),
+ connections: Rc::new(RefCell::new(Vec::new())),
+ is_shutting_down: Rc::new(RefCell::new(false)),
+ }
+ }
+
+ pub fn spawn_listener<L: Listener + 'static>(&self, mut listener: L) {
+ let name = listener.name();
+ let (shutdown_tx, shutdown_rx) = bounded(1);
+
+ let shard = self.shard.clone();
+ let connections = self.connections.clone();
+ let is_shutting_down = self.is_shutting_down.clone();
+
+ let handle = compio::runtime::spawn(async move {
+ info!("Starting listener: {}", name);
+
+ match listener
+ .run(shard, shutdown_rx, connections, is_shutting_down)
+ .await
+ {
+ Ok(()) => info!("Listener {} stopped", name),
+ Err(e) => error!("Listener {} failed: {}", name, e),
+ }
+ });
+
+ self.add_task(name, ShutdownPhase::Listeners, shutdown_tx, handle);
+ }
+
+ pub fn spawn_background<F, Fut>(&self, name: &'static str, interval: Duration, work: F)
+ where
+ F: Fn(Rc<IggyShard>) -> Fut + 'static,
+ Fut: Future<Output = Result<(), IggyError>> + 'static,
+ {
+ let (shutdown_tx, shutdown_rx) = bounded(1);
+ let shard = self.shard.clone();
+
+ let handle = compio::runtime::spawn(async move {
+ info!("Starting background task: {}", name);
+
+ loop {
+ futures::select! {
+ _ = shutdown_rx.recv().fuse() => {
+ info!("Background task {} stopping", name);
+ break;
+ }
+ _ = compio::time::sleep(interval).fuse() => {
+ if let Err(e) = work(shard.clone()).await {
+ error!("Background task {} error: {}", name, e);
+ }
+ }
+ }
+ }
+
+ info!("Background task {} stopped", name);
+ });
+
+ self.add_task(name, ShutdownPhase::Background, shutdown_tx, handle);
+ }
+
+ pub fn spawn_system<F, Fut>(&self, name: &'static str, work: F)
+ where
+ F: FnOnce(Rc<IggyShard>, Receiver<()>) -> Fut + 'static,
+ Fut: Future<Output = Result<(), IggyError>> + 'static,
+ {
+ let (shutdown_tx, shutdown_rx) = bounded(1);
+ let shard = self.shard.clone();
+
+ let handle = compio::runtime::spawn(async move {
+ info!("Starting system task: {}", name);
+
+ if let Err(e) = work(shard, shutdown_rx).await {
+ error!("System task {} failed: {}", name, e);
+ } else {
+ info!("System task {} stopped", name);
+ }
+ });
+
+ self.add_task(name, ShutdownPhase::System, shutdown_tx, handle);
+ }
+
+ pub fn register_connection(&self, handle: ConnectionHandle) {
+ let client_id = handle.client_id();
+ self.connections.borrow_mut().push(handle);
+ debug!("Registered connection {}", client_id);
+ }
+
+ pub fn unregister_connection(&self, client_id: u32) {
+ self.connections
+ .borrow_mut()
+ .retain(|h| h.client_id() != client_id);
+ debug!("Unregistered connection {}", client_id);
+ }
+
+ pub fn is_shutting_down(&self) -> bool {
+ *self.is_shutting_down.borrow()
+ }
+
+ pub async fn shutdown_all(&self, timeout_per_phase: Duration) -> Result<(), IggyError> {
+ info!("Starting orchestrated shutdown");
+ *self.is_shutting_down.borrow_mut() = true;
+
+ info!("Phase 1: Stopping listeners");
+ self.shutdown_phase(ShutdownPhase::Listeners, timeout_per_phase)
+ .await?;
+
+ info!(
+ "Phase 2: Closing {} connections",
+ self.connections.borrow().len()
+ );
+ self.shutdown_connections(timeout_per_phase).await?;
+
+ info!("Phase 3: Final data save");
+ self.save_all_data().await?;
+
+ info!("Phase 4: Stopping background tasks");
+ self.shutdown_phase(ShutdownPhase::Background, timeout_per_phase)
+ .await?;
+
+ info!("Phase 5: Stopping system tasks");
+ self.shutdown_phase(ShutdownPhase::System, timeout_per_phase)
+ .await?;
+
+ info!("Shutdown complete");
+ Ok(())
+ }
+
+ fn add_task(
+ &self,
+ name: &'static str,
+ phase: ShutdownPhase,
+ shutdown_tx: Sender<()>,
+ handle: JoinHandle<()>,
+ ) {
+ let task = Task {
+ name,
+ phase,
+ shutdown_tx,
+ handle,
+ };
+
+ self.tasks.borrow_mut().push(task);
+ info!("Spawned {} task: {}", phase_name(phase), name);
+ }
+
+ async fn shutdown_phase(
+ &self,
+ phase: ShutdownPhase,
+ timeout: Duration,
+ ) -> Result<(), IggyError> {
+ let mut phase_tasks = Vec::new();
+ {
+ let mut tasks = self.tasks.borrow_mut();
+ let mut i = 0;
+ while i < tasks.len() {
+ if tasks[i].phase == phase {
+ phase_tasks.push(tasks.remove(i));
+ } else {
+ i += 1;
+ }
+ }
+ }
+
+ for task in &phase_tasks {
+ let _ = task.shutdown_tx.send(()).await;
+ }
+
+ let deadline = Instant::now() + timeout;
+ for task in phase_tasks {
+ let remaining = deadline.saturating_duration_since(Instant::now());
+ match compio::time::timeout(remaining, task.handle).await {
+ Ok(Ok(())) => debug!("Task {} completed", task.name),
+ Ok(Err(_)) => warn!("Task {} panicked", task.name),
+ Err(_) => warn!("Task {} timed out", task.name),
+ }
+ }
+
+ Ok(())
+ }
+
+ async fn shutdown_connections(&self, timeout: Duration) -> Result<(), IggyError> {
+ let connections = std::mem::take(&mut *self.connections.borrow_mut());
+
+ for conn in &connections {
+ conn.request_shutdown();
+ }
+
+ let deadline = Instant::now() + timeout;
+ for conn in connections {
+ let client_id = conn.client_id();
+ let remaining = deadline.saturating_duration_since(Instant::now());
+ if let Err(_) = compio::time::timeout(remaining, conn.wait_closed()).await {
+ warn!("Connection {} shutdown timeout", client_id);
+ }
+ }
+
+ Ok(())
+ }
+
+ async fn save_all_data(&self) -> Result<(), IggyError> {
+ info!("Final data persistence before shutdown");
+
+ let namespaces = self.shard.get_current_shard_namespaces();
+ let mut total_saved = 0u32;
+
+ for ns in namespaces {
+ match self
+ .shard
+ .streams2
+ .persist_messages(
+ self.shard.id,
+ &iggy_common::Identifier::numeric(ns.stream_id() as u32).unwrap(),
+ &iggy_common::Identifier::numeric(ns.topic_id() as u32).unwrap(),
+ ns.partition_id(),
+ "shutdown save".to_string(),
+ &self.shard.config.system,
+ )
+ .await
+ {
+ Ok(count) => total_saved += count,
+ Err(e) => error!(
+ "Failed to save messages for partition {}: {}",
+ ns.partition_id(),
+ e
+ ),
+ }
+ }
+
+ if total_saved > 0 {
+ info!("Saved {} messages during shutdown", total_saved);
+ }
+
+ Ok(())
+ }
+}
+
+fn phase_name(phase: ShutdownPhase) -> &'static str {
+ match phase {
+ ShutdownPhase::Listeners => "listener",
+ ShutdownPhase::Connections => "connection",
+ ShutdownPhase::DataPersistence => "data",
+ ShutdownPhase::Background => "background",
+ ShutdownPhase::System => "system",
+ }
+}
diff --git a/core/server/src/shard/tasks/messages.rs b/core/server/src/shard/tasks/messages.rs
deleted file mode 100644
index 2feadd1c..00000000
--- a/core/server/src/shard/tasks/messages.rs
+++ /dev/null
@@ -1,59 +0,0 @@
-use futures::{FutureExt, StreamExt};
-use iggy_common::IggyError;
-use std::{rc::Rc, time::Duration};
-
-use crate::{
- shard::{IggyShard, transmission::frame::ShardFrame},
- shard_error, shard_info,
-};
-
-async fn run_shard_messages_receiver(shard: Rc<IggyShard>) -> Result<(), IggyError> {
- let mut messages_receiver = shard.messages_receiver.take().unwrap();
-
- shard_info!(shard.id, "Starting message passing task");
- loop {
- let shutdown_check = async {
- loop {
- if shard.is_shutting_down() {
- return;
- }
- compio::time::sleep(Duration::from_millis(100)).await;
- }
- };
-
- futures::select! {
- _ = shutdown_check.fuse() => {
- shard_info!(shard.id, "Message receiver shutting down");
- break;
- }
- frame = messages_receiver.next().fuse() => {
- if let Some(frame) = frame {
- let ShardFrame {
- message,
- response_sender,
- } = frame;
- match (shard.handle_shard_message(message).await, response_sender) {
- (Some(response), Some(response_sender)) => {
- response_sender
- .send(response)
- .await
- .expect("Failed to send response back to origin shard.");
- }
- _ => {}
- };
- }
- }
- }
- }
-
- Ok(())
-}
-
-pub async fn spawn_shard_message_task(shard: Rc<IggyShard>) -> Result<(), IggyError> {
- let result = run_shard_messages_receiver(shard.clone()).await;
- if let Err(err) = result {
- shard_error!(shard.id, "Error running shard message receiver: {err}");
- return Err(err);
- }
- Ok(())
-}
diff --git a/core/server/src/shard/tasks/mod.rs b/core/server/src/shard/tasks/mod.rs
deleted file mode 100644
index ba63992f..00000000
--- a/core/server/src/shard/tasks/mod.rs
+++ /dev/null
@@ -1 +0,0 @@
-pub mod messages;
diff --git a/core/server/src/tcp/tcp_listener.rs b/core/server/src/tcp/tcp_listener.rs
index 18565ae6..ddc099c1 100644
--- a/core/server/src/tcp/tcp_listener.rs
+++ b/core/server/src/tcp/tcp_listener.rs
@@ -20,7 +20,6 @@ use crate::binary::sender::SenderKind;
use crate::configs::tcp::TcpSocketConfig;
use crate::shard::IggyShard;
use crate::shard::transmission::event::ShardEvent;
-use crate::tcp::connection_handler::{handle_connection, handle_error};
use crate::{shard_error, shard_info};
use compio::net::{TcpListener, TcpOpts};
use error_set::ErrContext;
@@ -178,23 +177,16 @@ async fn accept_loop(
let client_id = session.client_id;
shard_info!(shard.id, "Created new session: {}", session);
- let mut sender = SenderKind::get_tcp_sender(stream);
-
- let conn_stop_receiver = shard_clone.task_registry.add_connection(client_id);
-
- let shard_for_conn = shard_clone.clone();
- shard_clone.task_registry.spawn_tracked(async move {
- if let Err(error) = handle_connection(&session, &mut sender, &shard_for_conn, conn_stop_receiver).await {
- handle_error(error);
- }
- shard_for_conn.task_registry.remove_connection(&client_id);
-
- if let Err(error) = sender.shutdown().await {
- shard_error!(shard.id, "Failed to shutdown TCP stream for client: {}, address: {}. {}", client_id, address, error);
- } else {
- shard_info!(shard.id, "Successfully closed TCP stream for client: {}, address: {}.", client_id, address);
- }
- });
+ let sender = SenderKind::get_tcp_sender(stream);
+
+ // TODO: Update to use new TaskManager system
+ // use crate::shard::task::connection::tcp_connection_task;
+ // let task = tcp_connection_task(session, sender, shard_clone.clone());
+ // shard_clone.task_registry.spawn_connection(
+ // client_id,
+ // transport,
+ // task,
+ // );
}
Err(error) => shard_error!(shard.id, "Unable to accept TCP socket. {}", error),
}
diff --git a/core/server/src/tcp/tcp_tls_listener.rs b/core/server/src/tcp/tcp_tls_listener.rs
index ec69dbcd..7e2cd54d 100644
--- a/core/server/src/tcp/tcp_tls_listener.rs
+++ b/core/server/src/tcp/tcp_tls_listener.rs
@@ -20,7 +20,6 @@ use crate::binary::sender::SenderKind;
use crate::configs::tcp::TcpSocketConfig;
use crate::shard::IggyShard;
use crate::shard::transmission::event::ShardEvent;
-use crate::tcp::connection_handler::{handle_connection, handle_error};
use crate::{shard_error, shard_info, shard_warn};
use compio::net::{TcpListener, TcpOpts};
use compio_tls::TlsAcceptor;
@@ -219,44 +218,41 @@ async fn accept_loop(
// Perform TLS handshake in a separate task to avoid blocking the accept loop
let task_shard = shard_clone.clone();
- task_shard.task_registry.spawn_tracked(async move {
+ compio::runtime::spawn(async move {
match acceptor.accept(stream).await {
Ok(tls_stream) => {
// TLS handshake successful, now create session
- shard_info!(shard_clone.id, "TLS handshake successful, adding TCP client: {}", address);
+ shard_info!(task_shard.id, "TLS handshake successful, adding TCP client: {}", address);
let transport = TransportProtocol::Tcp;
- let session = shard_clone.add_client(&address, transport);
- shard_info!(shard_clone.id, "Added {} client with session: {} for IP address: {}", transport, session, address);
+ let session = task_shard.add_client(&address, transport);
+ shard_info!(task_shard.id, "Added {} client with session: {} for IP address: {}", transport, session, address);
//TODO: Those can be shared with other shards.
- shard_clone.add_active_session(session.clone());
+ task_shard.add_active_session(session.clone());
// Broadcast session to all shards.
let event = ShardEvent::NewSession { address, transport };
// TODO: Fixme look inside of broadcast_event_to_all_shards method.
- let _responses = shard_clone.broadcast_event_to_all_shards(event.into()).await;
+ let _responses = task_shard.broadcast_event_to_all_shards(event.into()).await;
let client_id = session.client_id;
- shard_info!(shard_clone.id, "Created new session: {}", session);
+ shard_info!(task_shard.id, "Created new session: {}", session);
- let conn_stop_receiver = shard_clone.task_registry.add_connection(client_id);
- let shard_for_conn = shard_clone.clone();
- let mut sender = SenderKind::get_tcp_tls_sender(tls_stream);
- if let Err(error) = handle_connection(&session, &mut sender, &shard_for_conn, conn_stop_receiver).await {
- handle_error(error);
- }
- shard_for_conn.task_registry.remove_connection(&client_id);
+ let sender = SenderKind::get_tcp_tls_sender(tls_stream);
- if let Err(error) = sender.shutdown().await {
- shard_error!(shard.id, "Failed to shutdown TCP TLS stream for client: {}, address: {}. {}", client_id, address, error);
- } else {
- shard_info!(shard.id, "Successfully closed TCP TLS stream for client: {}, address: {}.", client_id, address);
- }
+ // TODO: Update to use new TaskManager system
+ // use crate::shard::task::connection::tcp_connection_task;
+ // let task = tcp_connection_task(session, sender, task_shard.clone());
+ // task_shard.task_registry.spawn_connection(
+ // client_id,
+ // transport,
+ // task,
+ // );
}
Err(e) => {
- shard_error!(shard_clone.id, "Failed to accept TLS connection from '{}': {}", address, e);
+ shard_error!(task_shard.id, "Failed to accept TLS connection from '{}': {}", address, e);
// No session was created, so no cleanup needed
}
}
- });
+ }).detach();
}
Err(error) => shard_error!(shard.id, "Unable to accept TCP TLS socket. {}", error),
}