zus_common/

endpoint.rs

use {
  bytes::Bytes,
  futures::{StreamExt, stream::SplitSink},
  std::{
    sync::{
      Arc,
      atomic::{AtomicU64, Ordering},
    },
    time::Duration,
  },
  tokio::{
    net::TcpStream,
    sync::{mpsc, oneshot},
    time::timeout,
  },
  tokio_util::codec::Framed,
  tracing::{info, instrument},
};

use crate::{
  codec::RpcCodec,
  error::{Result, ZusError},
  protocol::RpcMessage,
};

/// Request sent to the writer task (includes response callback for sync calls)
struct WriterRequest {
  msg: RpcMessage,
  response_tx: Option<oneshot::Sender<Bytes>>,
}

/// Response routing info stored by sequence number
struct PendingRequest {
  response_tx: oneshot::Sender<Bytes>,
}

/// Writer task that owns the SplitSink and receives messages via channel
async fn writer_task(
  mut writer: SplitSink<Framed<TcpStream, RpcCodec>, RpcMessage>,
  mut rx: mpsc::UnboundedReceiver<WriterRequest>,
  pending_requests: Arc<dashmap::DashMap<u64, PendingRequest>>,
) {
  use futures::SinkExt;

  while let Some(request) = rx.recv().await {
    let sequence = request.msg.header.sequence;

    // Register response callback if this is a sync call
    if let Some(response_tx) = request.response_tx {
      pending_requests.insert(sequence, PendingRequest { response_tx });
    }

    // Send the message
    if let Err(e) = writer.send(request.msg).await {
      tracing::error!("Writer task send error: {:?}", e);
      // Clean up pending request on error
      pending_requests.remove(&sequence);
      break;
    }
  }

  tracing::info!("Writer task shutting down");
}

/// Reader task that routes responses to the appropriate oneshot channels
async fn reader_task(
  mut reader: futures::stream::SplitStream<Framed<TcpStream, RpcCodec>>,
  pending_requests: Arc<dashmap::DashMap<u64, PendingRequest>>,
) {
  while let Some(result) = reader.next().await {
    match result {
      | Ok(msg) => {
        // Handle both regular responses (RSP) and system responses (SYSRSP)
        // C++ ZooServer may send SYSRSP (type 3) for certain operations like sync/heartbeat
        if msg.header.msg_type == zus_proto::constants::MSG_TYPE_RSP
          || msg.header.msg_type == zus_proto::constants::MSG_TYPE_SYSRSP
        {
          let sequence = msg.header.sequence;
          if let Some((_, pending)) = pending_requests.remove(&sequence) {
            if pending.response_tx.send(msg.body).is_err() {
              tracing::warn!("Failed to send response for sequence {}", sequence);
            }
          } else {
            tracing::warn!("No pending request for sequence {}", sequence);
          }
        } else {
          tracing::warn!("Unexpected message type: {}", msg.header.msg_type);
        }
      }
      | Err(e) => {
        tracing::error!("Reader task error: {:?}", e);
        break;
      }
    }
  }

  tracing::info!("Reader task shutting down");
}

/// RPC Endpoint (replacing Java's RpcEndPoint)
#[derive(Clone)]
pub struct RpcEndpoint {
  host: String,
  port: u16,
  writer_tx: mpsc::UnboundedSender<WriterRequest>,
  sequence: Arc<AtomicU64>,
}

impl RpcEndpoint {
  /// Create and connect to the remote endpoint
  #[instrument(name = "rpc_connect", skip_all, fields(host = %host, port = port))]
  pub async fn connect(host: String, port: u16) -> Result<Self> {
    let addr = format!("{host}:{port}");
    info!("Connecting to {}", addr);

    let stream = TcpStream::connect(&addr).await?;
    let framed = Framed::new(stream, RpcCodec::new());

    // Split the stream into read and write halves
    let (writer, reader) = framed.split();

    // Shared pending requests map for correlating responses
    let pending_requests = Arc::new(dashmap::DashMap::new());

    // Create channel for writer task
    let (writer_tx, writer_rx) = mpsc::unbounded_channel();

    // Spawn writer task
    tokio::spawn(writer_task(writer, writer_rx, pending_requests.clone()));

    // Spawn reader task
    tokio::spawn(reader_task(reader, pending_requests));

    Ok(Self {
      host,
      port,
      writer_tx,
      sequence: Arc::new(AtomicU64::new(1)),
    })
  }

  /// Synchronous RPC call
  #[instrument(name = "rpc_sync_call", skip(self, body), fields(endpoint = %self.address(), timeout_ms = timeout_ms))]
  pub async fn sync_call(&self, method: Bytes, body: Bytes, timeout_ms: u64) -> Result<Bytes> {
    // Generate sequence number
    let sequence = self.sequence.fetch_add(1, Ordering::SeqCst);

    // Create the request message
    let msg = RpcMessage::new_request(sequence, method, body);

    // Create oneshot channel for response
    let (response_tx, response_rx) = oneshot::channel();

    // Send via channel with response callback
    self
      .writer_tx
      .send(WriterRequest {
        msg,
        response_tx: Some(response_tx),
      })
      .map_err(|_| ZusError::Connection("Writer channel closed".to_string()))?;

    // Wait for response with timeout
    match timeout(Duration::from_millis(timeout_ms), response_rx).await {
      | Ok(Ok(response)) => Ok(response),
      | Ok(Err(_)) => Err(ZusError::Connection("Response channel closed".to_string())),
      | Err(_) => Err(ZusError::Timeout),
    }
  }

  /// Asynchronous RPC call (fire and forget)
  #[instrument(name = "rpc_notify_call", skip(self, body), fields(endpoint = %self.address()))]
  pub async fn notify_call(&self, method: Bytes, body: Bytes) -> Result<()> {
    // Generate sequence number
    let sequence = self.sequence.fetch_add(1, Ordering::SeqCst);

    // Create message
    let msg = RpcMessage::new_notify(sequence, method, body);

    // Send via channel (no response callback)
    self
      .writer_tx
      .send(WriterRequest { msg, response_tx: None })
      .map_err(|_| ZusError::Connection("Writer channel closed".to_string()))?;

    Ok(())
  }

  /// Get endpoint address
  pub fn address(&self) -> String {
    format!("{}:{}", self.host, self.port)
  }

  pub fn host(&self) -> &str {
    &self.host
  }

  pub fn port(&self) -> u16 {
    self.port
  }
}

#[cfg(test)]
mod tests {

  #[tokio::test]
  async fn test_endpoint_address() {
    // Can't actually connect in unit tests without a server,
    // but we can test the address formatting
    let host = "localhost".to_string();
    let port = 9527u16;
    let addr = format!("{host}:{port}");
    assert_eq!(addr, "localhost:9527");
  }
}