Adds Redis Pub/Sub integration

Cargo.toml

@@ -41,8 +41,9 @@ mysql = { version = "20.0", optional = true }
 dashmap = { version = "4.0", optional = true }
 zip = { version = "0.5.8", optional = true }
 rand = {version = "0.8", optional = true}
-dmsort = {version = "1.0.0", optional = true}
+dmsort = {version = "1.0.0", optional = true }
 toml-dep = { version = "0.5.8", package="toml", optional = true }
+redis = { version = "0.21.4", optional = true }
 
 [features]
 default = ["cellularnoise", "dmi", "file", "git", "http", "json", "log", "noise", "sql", "time", "toml", "url"]
@@ -64,6 +65,7 @@ url = ["url-dep", "percent-encoding"]
 hash = ["base64", "const-random", "md-5", "hex", "sha-1", "sha2",  "twox-hash", "serde", "serde_json"]
 unzip = ["zip", "jobs"]
 worleynoise = ["rand","dmsort"]
+redis_pubsub = ["flume", "redis", "serde", "serde_json"]
 
 # internal feature-like things
 jobs = ["flume"]

dmsrc/redis_pubsub.dm

@@ -0,0 +1,7 @@
+#define RUSTG_REDIS_ERROR_CHANNEL "RUSTG_REDIS_ERROR_CHANNEL"
+
+#define rustg_redis_connect(addr) call(RUST_G, "redis_connect")(addr)
+/proc/rustg_redis_disconnect() return call(RUST_G, "redis_disconnect")()
+#define rustg_redis_subscribe(channel) call(RUST_G, "redis_subscribe")(channel)
+/proc/rustg_redis_get_messages() return call(RUST_G, "redis_get_messages")()
+#define rustg_redis_publish(channel, message) call(RUST_G, "redis_publish")(channel, message)

src/lib.rs

@@ -24,6 +24,8 @@ pub mod json;
 pub mod log;
 #[cfg(feature = "noise")]
 pub mod noise_gen;
+#[cfg(feature = "redis_pubsub")]
+pub mod redis_pubsub;
 #[cfg(feature = "sql")]
 pub mod sql;
 #[cfg(feature = "time")]

src/redis_pubsub.rs

@@ -0,0 +1,164 @@
+use redis::{Client, Commands, RedisError};
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::thread;
+use std::time::Duration;
+
+static ERROR_CHANNEL: &'static str = "RUSTG_REDIS_ERROR_CHANNEL";
+
+thread_local! {
+    static REQUEST_SENDER: RefCell<Option<flume::Sender<PubSubRequest>>> = RefCell::new(None);
+    static RESPONSE_RECEIVER: RefCell<Option<flume::Receiver<PubSubResponse>>> = RefCell::new(None);
+}
+
+enum PubSubRequest {
+    Quit,
+    Subscribe(String),
+    Publish(String, String),
+}
+
+// response might not be a good name, since those are not sent in response to requests
+enum PubSubResponse {
+    Disconnected(String),
+    Message(String, String),
+}
+
+fn handle_redis_inner(
+    client: Client,
+    control: flume::Receiver<PubSubRequest>,
+    out: flume::Sender<PubSubResponse>,
+) -> Result<(), RedisError> {
+    let mut conn = client.get_connection()?;
+    let mut pub_conn = client.get_connection()?;
+    let mut pubsub = conn.as_pubsub();
+    let _ = pubsub.set_read_timeout(Some(Duration::from_secs(1)));
+
+    'outer: loop {
+        for req in control.try_iter() {
+            match req {
+                PubSubRequest::Quit => break 'outer,
+                PubSubRequest::Subscribe(chan) => {
+                    pubsub.subscribe(&chan)?;
+                }
+                PubSubRequest::Publish(chan, msg) => {
+                    // kinda lame how PubSub doesn't have the Pub
+                    pub_conn.publish(&chan, &msg)?
+                }
+            }
+        }
+
+        if let Some(msg) = match pubsub.get_message() {
+            Ok(msg) => Some(msg),
+            Err(e) => {
+                if e.is_timeout() {
+                    None
+                } else {
+                    return Err(e);
+                }
+            }
+        } {
+            let chan = msg.get_channel_name().to_owned();
+            let data: String = msg.get_payload().unwrap_or_default();
+            let _ = out.send(PubSubResponse::Message(chan, data));
+        }
+    }
+
+    Ok(())
+}
+
+fn handle_redis(
+    client: Client,
+    control: flume::Receiver<PubSubRequest>,
+    out: flume::Sender<PubSubResponse>,
+) {
+    let out_copy = out.clone();
+    if let Err(e) = handle_redis_inner(client, control, out) {
+        let _ = out_copy.send(PubSubResponse::Disconnected(e.to_string()));
+    }
+}
+
+fn connect(addr: &str) -> Result<(), RedisError> {
+    let client = redis::Client::open(addr)?;
+    let (c_sender, c_receiver) = flume::unbounded();
+    let (o_sender, o_receiver) = flume::unbounded();
+    REQUEST_SENDER.with(|cell| cell.replace(Some(c_sender)));
+    RESPONSE_RECEIVER.with(|cell| cell.replace(Some(o_receiver)));
+    thread::spawn(|| handle_redis(client, c_receiver, o_sender));
+    Ok(())
+}
+
+fn disconnect() {
+    REQUEST_SENDER.with(|cell| {
+        if let Some(chan) = cell.borrow_mut().as_ref() {
+            let _ = chan.send(PubSubRequest::Quit);
+        }
+        cell.replace(None);
+    });
+    RESPONSE_RECEIVER.with(|cell| {
+        cell.replace(None);
+    });
+}
+
+fn subscribe(channel: &str) {
+    REQUEST_SENDER.with(|cell| {
+        if let Some(chan) = cell.borrow_mut().as_ref() {
+            let _ = chan.send(PubSubRequest::Subscribe(channel.to_owned()));
+        };
+    });
+}
+
+fn publish(channel: &str, msg: &str) {
+    REQUEST_SENDER.with(|cell| {
+        if let Some(chan) = cell.borrow_mut().as_ref() {
+            let _ = chan.send(PubSubRequest::Publish(channel.to_owned(), msg.to_owned()));
+        };
+    });
+}
+
+fn get_messages() -> String {
+    let mut result: HashMap<String, Vec<String>> = HashMap::new();
+    RESPONSE_RECEIVER.with(|cell| {
+        let opt = cell.borrow_mut();
+        if let Some(recv) = opt.as_ref() {
+            for resp in recv.try_iter() {
+                match resp {
+                    PubSubResponse::Message(chan, msg) => {
+                        result.entry(chan).or_default().push(msg);
+                    }
+                    PubSubResponse::Disconnected(error) => {
+                        // Pardon the in-band signaling but it's probably the best way to do this
+                        result
+                            .entry(ERROR_CHANNEL.to_owned())
+                            .or_default()
+                            .push(error);
+                    }
+                }
+            }
+        }
+    });
+
+    serde_json::to_string(&result).unwrap_or("{}".to_owned())
+}
+
+byond_fn! { redis_connect(addr) {
+    connect(addr).err().map(|e| e.to_string())
+} }
+
+byond_fn! { redis_disconnect() {
+    disconnect();
+    Some("")
+} }
+
+byond_fn! { redis_subscribe(channel) {
+    subscribe(channel);
+    Some("")
+} }
+
+byond_fn! { redis_get_messages() {
+    Some(get_messages())
+} }
+
+byond_fn! { redis_publish(channel, message) {
+    publish(channel, message);
+    Some("")
+} }