feat(server): socket transfer for message polling

Author: tungtose

core/integration/tests/server/scenarios/consumer_group_with_single_client_polling_messages_scenario.rs

@@ -175,6 +175,7 @@ async fn execute_using_none_key(client: &IggyClient) {
         let message = &polled_messages.messages[0];
         assert_eq!(message.header.offset, offset);
         let payload = from_utf8(&message.payload).unwrap();
+
         assert_eq!(
             payload,
             &create_extended_message_payload(partition_id, entity_id)

core/server/src/binary/handlers/messages/poll_messages_handler.rs

@@ -22,12 +22,18 @@ use crate::binary::command::{
 use crate::binary::handlers::utils::receive_and_validate;
 use crate::shard::IggyShard;
 use crate::shard::system::messages::PollingArgs;
+use crate::shard::transmission::message::{
+    ShardMessage, ShardRequest, ShardRequestPayload, SocketTransferPayload,
+};
+use crate::streaming::segments::IggyMessagesBatchSet;
 use crate::streaming::session::Session;
+use crate::streaming::{streams, topics};
 use anyhow::Result;
-use iggy_common::SenderKind;
+use iggy_common::sharding::IggyNamespace;
+use iggy_common::{ConsumerKind, SenderKind};
 use iggy_common::{IggyError, PollMessages, PooledBuffer};
 use std::rc::Rc;
-use tracing::{debug, trace};
+use tracing::{debug, error, trace};
 
 #[derive(Debug)]
 pub struct IggyPollMetadata {
@@ -57,6 +63,7 @@ impl ServerCommandHandler for PollMessages {
         shard: &Rc<IggyShard>,
     ) -> Result<HandlerResult, IggyError> {
         debug!("session: {session}, command: {self}");
+
         let PollMessages {
             consumer,
             partition_id,
@@ -66,10 +73,95 @@ impl ServerCommandHandler for PollMessages {
             stream_id,
             topic_id,
         } = self;
-        let args = PollingArgs::new(strategy, count, auto_commit);
+
+        shard.ensure_topic_exists(&stream_id, &topic_id)?;
 
         let user_id = session.get_user_id();
         let client_id = session.client_id;
+
+        let args = PollingArgs::new(strategy, count, auto_commit);
+
+        if consumer.kind == ConsumerKind::Consumer
+            && partition_id.is_some()
+            && let Some((polling_consumer, resolved_partition_id)) = shard
+                .resolve_consumer_with_partition_id(
+                    &stream_id,
+                    &topic_id,
+                    &consumer,
+                    client_id,
+                    partition_id,
+                    true,
+                )?
+        {
+            let numeric_stream_id = shard
+                .streams
+                .with_stream_by_id(&stream_id, streams::helpers::get_stream_id());
+            let numeric_topic_id = shard.streams.with_topic_by_id(
+                &stream_id,
+                &topic_id,
+                topics::helpers::get_topic_id(),
+            );
+            let namespace =
+                IggyNamespace::new(numeric_stream_id, numeric_topic_id, resolved_partition_id);
+
+            let enabled_socket_migration = shard.config.tcp.socket_migration;
+
+            if enabled_socket_migration
+                && !session.is_migrated()
+                && let Some(target_shard) = shard.find_shard(&namespace)
+                && target_shard.id != shard.id
+            {
+                debug!(
+                    "TCP wrong shard detected: migrating from_shard {}, to_shard {}",
+                    shard.id, target_shard.id
+                );
+
+                shard.ensure_partition_exists(&stream_id, &topic_id, resolved_partition_id)?;
+
+                if let Some(fd) = sender.take_and_migrate_tcp() {
+                    let data = SocketTransferPayload::PollMessages {
+                        consumer: polling_consumer,
+                        args: args.clone(),
+                    };
+
+                    let payload = ShardRequestPayload::SocketTransfer {
+                        fd,
+                        from_shard: shard.id,
+                        client_id: session.client_id,
+                        user_id,
+                        address: session.ip_address,
+                        initial_data: data,
+                    };
+
+                    let request = ShardRequest::new(
+                        stream_id.clone(),
+                        topic_id.clone(),
+                        resolved_partition_id,
+                        payload,
+                    );
+
+                    let socket_transfer_msg = ShardMessage::Request(request);
+
+                    debug!(
+                        "Sending migrate socket to another shard {:?}",
+                        socket_transfer_msg
+                    );
+
+                    if let Err(e) = shard
+                        .send_request_to_shard_or_recoil(Some(&namespace), socket_transfer_msg)
+                        .await
+                    {
+                        error!("tranfer socket to another shard failed, drop connection. {e:?}");
+                        return Ok(HandlerResult::Finished);
+                    }
+
+                    return Ok(HandlerResult::Migrated {
+                        to_shard: target_shard.id,
+                    });
+                }
+            }
+        }
+
         let (metadata, mut batch) = shard
             .poll_messages(
                 client_id,
@@ -82,41 +174,52 @@ impl ServerCommandHandler for PollMessages {
             )
             .await?;
 
-        // Collect all chunks first into a Vec to extend their lifetimes.
-        // This ensures the Bytes (in reality Arc<[u8]>) references from each IggyMessagesBatch stay alive
-        // throughout the async vectored I/O operation, preventing "borrowed value does not live
-        // long enough" errors while optimizing transmission by using larger chunks.
-
-        // 4 bytes for partition_id + 8 bytes for current_offset + 4 bytes for messages_count + size of all batches.
-        let response_length = 4 + 8 + 4 + batch.size();
-        let response_length_bytes = response_length.to_le_bytes();
-
-        let mut bufs = Vec::with_capacity(batch.containers_count() + 5);
-        let mut partition_id_buf = PooledBuffer::with_capacity(4);
-        let mut current_offset_buf = PooledBuffer::with_capacity(8);
-        let mut count_buf = PooledBuffer::with_capacity(4);
-        partition_id_buf.put_u32_le(metadata.partition_id);
-        current_offset_buf.put_u64_le(metadata.current_offset);
-        count_buf.put_u32_le(batch.count());
-
-        bufs.push(partition_id_buf);
-        bufs.push(current_offset_buf);
-        bufs.push(count_buf);
-
-        batch.iter_mut().for_each(|m| bufs.push(m.take_messages()));
-        trace!(
-            "Sending {} messages to client ({} bytes) to client",
-            batch.count(),
-            response_length
-        );
+        let (response_length_bytes, bufs) = prepare_message_batch_buffers(&metadata, &mut batch);
 
         sender
             .send_ok_response_vectored(&response_length_bytes, bufs)
             .await?;
+
         Ok(HandlerResult::Finished)
     }
 }
 
+pub fn prepare_message_batch_buffers(
+    metadata: &IggyPollMetadata,
+    batch: &mut IggyMessagesBatchSet,
+) -> ([u8; 4], Vec<PooledBuffer>) {
+    // Collect all chunks first into a Vec to extend their lifetimes.
+    // This ensures the Bytes (in reality Arc<[u8]>) references from each IggyMessagesBatch stay alive
+    // throughout the async vectored I/O operation, preventing "borrowed value does not live
+    // long enough" errors while optimizing transmission by using larger chunks.
+
+    // 4 bytes for partition_id + 8 bytes for current_offset + 4 bytes for messages_count + size of all batches.
+    let response_length = 4 + 8 + 4 + batch.size();
+    let response_length_bytes = response_length.to_le_bytes();
+
+    let mut bufs = Vec::with_capacity(batch.containers_count() + 5);
+    let mut partition_id_buf = PooledBuffer::with_capacity(4);
+    let mut current_offset_buf = PooledBuffer::with_capacity(8);
+    let mut count_buf = PooledBuffer::with_capacity(4);
+    partition_id_buf.put_u32_le(metadata.partition_id);
+    current_offset_buf.put_u64_le(metadata.current_offset);
+    count_buf.put_u32_le(batch.count());
+
+    bufs.push(partition_id_buf);
+    bufs.push(current_offset_buf);
+    bufs.push(count_buf);
+
+    batch.iter_mut().for_each(|m| bufs.push(m.take_messages()));
+
+    trace!(
+        "Sending {} messages to client ({} bytes) to client",
+        batch.count(),
+        response_length
+    );
+
+    (response_length_bytes, bufs)
+}
+
 impl BinaryServerCommand for PollMessages {
     async fn from_sender(
         sender: &mut SenderKind,

core/server/src/binary/handlers/messages/send_messages_handler.rs

@@ -18,7 +18,9 @@
 
 use crate::binary::command::{BinaryServerCommand, HandlerResult, ServerCommandHandler};
 use crate::shard::IggyShard;
-use crate::shard::transmission::message::{ShardMessage, ShardRequest, ShardRequestPayload};
+use crate::shard::transmission::message::{
+    ShardMessage, ShardRequest, ShardRequestPayload, SocketTransferPayload,
+};
 use crate::streaming::segments::{IggyIndexesMut, IggyMessagesBatchMut};
 use crate::streaming::session::Session;
 use crate::streaming::{streams, topics};
@@ -32,7 +34,7 @@ use iggy_common::sharding::IggyNamespace;
 use iggy_common::{INDEX_SIZE, PartitioningKind};
 use iggy_common::{IggyError, Partitioning, SendMessages, Validatable};
 use std::rc::Rc;
-use tracing::{debug, error, info, instrument};
+use tracing::{debug, error, instrument};
 
 impl ServerCommandHandler for SendMessages {
     fn code(&self) -> u32 {
@@ -163,7 +165,7 @@ impl ServerCommandHandler for SendMessages {
             && target_shard.id != shard.id
         {
             debug!(
-                "TCP wrong shared detected: migrating from_shard {}, to_shard {}",
+                "TCP wrong shard detected: migrating from_shard {}, to_shard {}",
                 shard.id, target_shard.id
             );
 
@@ -174,7 +176,7 @@ impl ServerCommandHandler for SendMessages {
                     client_id: session.client_id,
                     user_id,
                     address: session.ip_address,
-                    initial_data: batch,
+                    initial_data: SocketTransferPayload::SendMessages { batch },
                 };
 
                 let request = ShardRequest::new(
@@ -194,7 +196,7 @@ impl ServerCommandHandler for SendMessages {
                     return Ok(HandlerResult::Finished);
                 }
 
-                info!("Sending socket transfer to shard {}", target_shard.id);
+                debug!("Sending socket transfer to shard {}", target_shard.id);
                 return Ok(HandlerResult::Migrated {
                     to_shard: target_shard.id,
                 });