use crate::{KevyError, KevyResult};
use std::collections::VecDeque;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::Duration;
use kevy_embedded::Subscription;
use kevy_resp::{Reply, encode_command};
use kevy_resp_client::ReplyReadBuf;
use crate::subscribe_io::{await_acks, frame_to_event, invalid, recv_remote, send_to, shape};
use crate::{Target, parse_url, resolve_store};
#[derive(Debug)]
pub struct Subscriber {
inner: Inner,
}
#[derive(Debug)]
enum Inner {
Remote {
stream: TcpStream,
buf: ReplyReadBuf,
pending: VecDeque<PubsubEvent>,
},
Embedded {
subscription: Subscription,
timeout: Option<Duration>,
},
}
pub use kevy_resp_client::PubsubEvent;
impl Subscriber {
pub fn connect(url: &str) -> KevyResult<Self> {
let target = parse_url(url)?;
let inner = match target {
Target::EmbedMemoryAnonymous => {
return Err(KevyError::Unsupported("anonymous mem:// has no other producer; use mem://<name> for a shared bus".into()));
}
Target::EmbedMemoryNamed(_) | Target::EmbedPersist(_) => Inner::Embedded {
subscription: resolve_store(&target)?.subscribe(&[]),
timeout: None,
},
Target::Remote(remote_url) => {
let (host, port) = remote_host_port(&remote_url)?;
let stream = TcpStream::connect((host.as_str(), port))?;
stream.set_nodelay(true).ok();
Inner::Remote {
stream,
buf: ReplyReadBuf::with_capacity(8192),
pending: VecDeque::new(),
}
}
};
Ok(Self { inner })
}
pub fn connect_channels(url: &str, channels: &[&[u8]]) -> KevyResult<Self> {
if channels.is_empty() {
return Err(KevyError::InvalidInput("Subscriber::connect_channels needs ≥ 1 channel — use Subscriber::connect() for empty start".into()));
}
let mut s = Self::connect(url)?;
s.subscribe(channels)?;
Ok(s)
}
pub fn subscribe(&mut self, channels: &[&[u8]]) -> KevyResult<()> {
if channels.is_empty() {
return Err(KevyError::InvalidInput("SUBSCRIBE needs ≥ 1 channel".into()));
}
match &mut self.inner {
Inner::Remote { stream, buf, pending } => {
send_to(stream, b"SUBSCRIBE", channels)?;
await_acks(stream, buf, pending, channels.len(), false)
}
Inner::Embedded { subscription, .. } => {
subscription.subscribe(channels);
Ok(())
}
}
}
pub fn psubscribe(&mut self, patterns: &[&[u8]]) -> KevyResult<()> {
if patterns.is_empty() {
return Err(KevyError::InvalidInput("PSUBSCRIBE needs ≥ 1 pattern".into()));
}
match &mut self.inner {
Inner::Remote { stream, buf, pending } => {
send_to(stream, b"PSUBSCRIBE", patterns)?;
await_acks(stream, buf, pending, patterns.len(), true)
}
Inner::Embedded { subscription, .. } => {
subscription.psubscribe(patterns);
Ok(())
}
}
}
pub fn unsubscribe(&mut self, channels: &[&[u8]]) -> KevyResult<()> {
match &mut self.inner {
Inner::Remote { stream, .. } => send_to(stream, b"UNSUBSCRIBE", channels),
Inner::Embedded { subscription, .. } => {
subscription.unsubscribe(channels);
Ok(())
}
}
}
pub fn punsubscribe(&mut self, patterns: &[&[u8]]) -> KevyResult<()> {
match &mut self.inner {
Inner::Remote { stream, .. } => send_to(stream, b"PUNSUBSCRIBE", patterns),
Inner::Embedded { subscription, .. } => {
subscription.punsubscribe(patterns);
Ok(())
}
}
}
pub fn recv(&mut self) -> KevyResult<PubsubEvent> {
match &mut self.inner {
Inner::Remote { stream, buf, pending } => match pending.pop_front() {
Some(ev) => Ok(ev),
None => recv_remote(stream, buf),
},
Inner::Embedded {
subscription,
timeout,
} => {
let frame = match *timeout {
Some(d) => subscription.recv_timeout(d)?,
None => subscription.recv()?,
};
Ok(frame_to_event(frame))
}
}
}
pub fn recv_message(&mut self) -> KevyResult<(Vec<u8>, Vec<u8>)> {
loop {
match self.recv()? {
PubsubEvent::Message { channel, payload } => return Ok((channel, payload)),
PubsubEvent::Pmessage { channel, payload, .. } => {
return Ok((channel, payload));
}
_ => {}
}
}
}
pub fn hello3(&mut self) -> KevyResult<PubsubEvent> {
match &mut self.inner {
Inner::Embedded { .. } => Err(KevyError::Unsupported("HELLO 3 is a remote/TCP-only operation; embedded backend has no proto switch".into())),
Inner::Remote { stream, buf, .. } => {
let mut frame = Vec::new();
encode_command(&mut frame, &[b"HELLO".to_vec(), b"3".to_vec()]);
stream.write_all(&frame)?;
let mut chunk = [0u8; 4096];
loop {
match buf.parse_next() {
Ok(Some(reply)) => return classify_hello3_reply(reply),
Ok(None) => {}
Err(_) => {
return Err(KevyError::Protocol("malformed HELLO 3 reply".into()));
}
}
let n = stream.read(&mut chunk)?;
if n == 0 {
return Err(KevyError::Closed);
}
buf.extend(&chunk[..n]);
}
}
}
}
pub fn set_read_timeout(&mut self, dur: Option<Duration>) -> KevyResult<()> {
match &mut self.inner {
Inner::Remote { stream, .. } => Ok(stream.set_read_timeout(dur)?),
Inner::Embedded { timeout, .. } => {
*timeout = dur;
Ok(())
}
}
}
pub fn events(&mut self) -> SubscriberEvents<'_> {
SubscriberEvents { sub: self }
}
pub fn messages(&mut self) -> SubscriberMessages<'_> {
SubscriberMessages { sub: self }
}
}
#[derive(Debug)]
pub struct SubscriberEvents<'a> {
sub: &'a mut Subscriber,
}
impl Iterator for SubscriberEvents<'_> {
type Item = KevyResult<PubsubEvent>;
fn next(&mut self) -> Option<Self::Item> {
match self.sub.recv() {
Err(KevyError::Closed) => None,
other => Some(other),
}
}
}
#[derive(Debug)]
pub struct SubscriberMessages<'a> {
sub: &'a mut Subscriber,
}
impl Iterator for SubscriberMessages<'_> {
type Item = KevyResult<(Vec<u8>, Vec<u8>)>;
fn next(&mut self) -> Option<Self::Item> {
match self.sub.recv_message() {
Err(KevyError::Closed) => None,
other => Some(other),
}
}
}
fn classify_hello3_reply(reply: Reply) -> KevyResult<PubsubEvent> {
match reply {
Reply::Map(_) | Reply::Array(_) => Ok(PubsubEvent::Subscribe {
channel: b"HELLO".to_vec(),
count: 3,
}),
Reply::Error(e) => Err(KevyError::Protocol(String::from_utf8_lossy(&e).into_owned())),
other => Err(invalid(format!(
"unexpected HELLO 3 reply shape: {}",
shape(&other)
))),
}
}
fn remote_host_port(url: &str) -> KevyResult<(String, u16)> {
let (_scheme, rest) = url.split_once("://").ok_or_else(|| {
KevyError::InvalidInput("URL missing '://'".into())
})?;
if rest.contains('@') {
return Err(KevyError::Unsupported("userinfo (user:pass@host) is unsupported — kevy has no AUTH".into()));
}
let authority = rest.split('/').next().unwrap_or("");
let (host, port) = match authority.rsplit_once(':') {
Some((h, p)) => {
let port: u16 = p.parse().map_err(|_| {
KevyError::InvalidInput(format!("bad port: {p}"))
})?;
(h.to_string(), port)
}
None => (authority.to_string(), 6379),
};
if host.is_empty() {
return Err(KevyError::InvalidInput("empty host".into()));
}
Ok((host, port))
}
#[cfg(test)]
#[path = "subscribe_tests.rs"]
mod tests;