use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt::{Debug, Display, Formatter};
use std::sync::RwLock;
use std::sync::atomic::{AtomicBool, Ordering};
use crate::envelope::SharedEnvelope;
use crate::error::RedisError;
use crate::pubsub::PubSubMode;
#[derive(Clone)]
pub(crate) enum Route {
Stream,
List { envelope: Option<SharedEnvelope> },
Channel {
mode: PubSubMode,
envelope: Option<SharedEnvelope>,
},
}
impl Route {
const fn noun(&self) -> &'static str {
match self {
Self::Stream => "stream",
Self::List { .. } => "list",
Self::Channel { .. } => "Pub/Sub channel",
}
}
fn writes_like(&self, other: &Self) -> bool {
match (self, other) {
(Self::Stream, Self::Stream) => true,
(Self::List { envelope: a }, Self::List { envelope: b }) => {
framing(a.as_ref()) == framing(b.as_ref())
}
(
Self::Channel {
mode: mode_a,
envelope: a,
},
Self::Channel {
mode: mode_b,
envelope: b,
},
) => mode_a == mode_b && framing(a.as_ref()) == framing(b.as_ref()),
_ => false,
}
}
}
fn framing(envelope: Option<&SharedEnvelope>) -> Option<&'static str> {
envelope.map(|codec| codec.framing())
}
impl Debug for Route {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Stream => f.write_str(self.noun()),
Self::List { envelope } => write!(f, "list framed {}", Framing(envelope)),
Self::Channel { mode, envelope } => {
let mode = match mode {
PubSubMode::Classic => "classic",
PubSubMode::Sharded => "sharded",
};
write!(f, "{mode} {} framed {}", self.noun(), Framing(envelope))
}
}
}
}
struct Framing<'a>(&'a Option<SharedEnvelope>);
impl Display for Framing<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match framing(self.0.as_ref()) {
None => f.write_str("as raw bytes"),
Some(codec) => write!(f, "by `{codec}`"),
}
}
}
#[derive(Debug, Default)]
pub(crate) struct Routes {
table: RwLock<HashMap<String, Held>>,
foreign: AtomicBool,
}
#[derive(Debug)]
struct Held {
route: Route,
holders: usize,
}
impl Routes {
pub(crate) fn record_subscription(
&self,
name: &str,
dead_letter: Option<&str>,
route: &Route,
) -> Result<Recorded<'_>, RedisError> {
let mut routes = self.table.write().expect("redis route table poisoned");
let names = std::iter::once(name).chain(dead_letter);
for written in names.clone() {
if let Some(Held { route: held, .. }) = routes.get(written)
&& !held.writes_like(route)
{
return Err(RedisError::InvalidOptions(format!(
"the subscription on `{name}` writes `{written}` as a {route:?}, but this \
service already reads or dead-letters to `{written}` as a {held:?}: one of \
the two would receive writes it cannot read, so give one of them another \
name or the same settings",
)));
}
}
let mut recorded = Recorded {
routes: self,
added: Vec::new(),
};
for written in names {
match routes.entry(written.to_owned()) {
Entry::Vacant(vacant) => {
if !matches!(route, Route::Stream) {
self.foreign.store(true, Ordering::Release);
}
vacant.insert(Held {
route: route.clone(),
holders: 1,
});
}
Entry::Occupied(mut occupied) => occupied.get_mut().holders += 1,
}
recorded.added.push(written.to_owned());
}
drop(routes);
Ok(recorded)
}
fn forget(&self, names: &[String]) {
if let Ok(mut routes) = self.table.write() {
for name in names {
if let Entry::Occupied(mut held) = routes.entry(name.clone()) {
held.get_mut().holders -= 1;
if held.get().holders == 0 {
held.remove();
}
}
}
}
}
pub(crate) fn route(&self, name: &str) -> Route {
if !self.foreign.load(Ordering::Relaxed) {
return Route::Stream;
}
self.table
.read()
.expect("redis route table poisoned")
.get(name)
.map_or(Route::Stream, |held| held.route.clone())
}
}
#[must_use = "the routes are forgotten again unless kept"]
pub(crate) struct Recorded<'a> {
routes: &'a Routes,
added: Vec<String>,
}
impl Recorded<'_> {
pub(crate) fn keep(mut self) {
self.added.clear();
}
}
impl Drop for Recorded<'_> {
fn drop(&mut self) {
if !self.added.is_empty() {
self.routes.forget(&self.added);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_name_nothing_reads_is_a_stream() {
assert!(matches!(Routes::default().route("orders"), Route::Stream));
}
#[test]
fn a_failed_open_keeps_a_route_another_subscription_shares() {
let routes = Routes::default();
let channel = Route::Channel {
mode: PubSubMode::Classic,
envelope: None,
};
let failing = routes
.record_subscription("events", Some("events.dead"), &channel)
.expect("record");
routes
.record_subscription("events", Some("events.dead"), &channel)
.expect("a second subscription shares the channel")
.keep();
drop(failing);
assert!(matches!(routes.route("events"), Route::Channel { .. }));
assert!(matches!(routes.route("events.dead"), Route::Channel { .. }));
}
#[test]
fn a_recorded_name_keeps_its_family() {
let routes = Routes::default();
routes
.record_subscription("jobs", None, &Route::List { envelope: None })
.expect("record")
.keep();
routes
.record_subscription("jobs", None, &Route::List { envelope: None })
.expect("a second list on the same key shares it")
.keep();
assert!(matches!(routes.route("jobs"), Route::List { .. }));
}
#[test]
fn a_name_read_in_two_families_is_refused() {
let routes = Routes::default();
routes
.record_subscription("jobs", None, &Route::List { envelope: None })
.expect("record")
.keep();
let err = routes
.record_subscription("orders", Some("jobs"), &Route::Stream)
.err()
.expect("a list key cannot be a stream too");
let text = err.to_string();
assert!(text.contains("`orders`") && text.contains("stream") && text.contains("list"));
assert!(
matches!(routes.route("orders"), Route::Stream),
"a refused subscription records neither of its names"
);
}
#[test]
fn a_channel_read_in_two_modes_is_refused() {
let routes = Routes::default();
let channel = |mode| Route::Channel {
mode,
envelope: None,
};
routes
.record_subscription("events", None, &channel(PubSubMode::Classic))
.expect("record")
.keep();
let err = routes
.record_subscription("events", None, &channel(PubSubMode::Sharded))
.err()
.expect("a sharded subscription never hears a classic publish");
assert!(err.to_string().contains("sharded"), "got {err}");
}
#[test]
fn a_subscription_that_does_not_open_forgets_its_routes() {
let routes = Routes::default();
drop(
routes
.record_subscription("jobs", Some("jobs.dead"), &Route::List { envelope: None })
.expect("record"),
);
assert!(matches!(routes.route("jobs"), Route::Stream));
assert!(matches!(routes.route("jobs.dead"), Route::Stream));
routes
.record_subscription("jobs", None, &Route::Stream)
.expect("the name is free again")
.keep();
}
}