gitdata.ai/imks
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(config): integrate etcd for service discovery and config management

Browse Source 
- Add etcd-client dependency for distributed configuration storage
- Implement EtcdConfig with priority: etcd > environment variables > defaults
- Add ServiceRegistry for service registration with lease keep-alive
- Integrate etcd-based service discovery for appks gRPC connections
- Add service watcher for real-time service instance updates
- Migrate Redis configuration from single URL to cluster node list
- Update Dockerfile with default IMKS_HOST and IMKS_PORT environment variables
- Add etcd bootstrap configuration through environment variables
- Implement Redis cluster URL building with optional authentication
This commit is contained in:
zhenyi
2026-06-11 22:50:38 +08:00
parent e72866db8d
commit c794b818ff
24 changed files with 984 additions and 952 deletions
Download Patch File Download Diff File Expand all files Collapse all files
socket/adapter/redis.rs
+13 -10
View File
@@ -3,13 +3,12 @@ use std::sync::Arc;
use async_trait::async_trait;
use dashmap::DashMap;
use fred::clients::Client;
use fred::interfaces::{KeysInterface, SetsInterface};
use tokio::sync::mpsc;
use crate::socket::adapter::{
Adapter, AdapterError, BroadcastOptions, BusMessage, LocalBroadcastFn, SocketInfo,
};
use crate::socket::message_bus::redis::RedisCommandClient;
use crate::socket::message_bus::MessageBus;
use crate::socket::packet::Packet;
use crate::socket::parser;
@@ -68,7 +67,7 @@ async fn handle_bus_message(
pub struct RedisAdapter {
message_bus: Arc<dyn MessageBus>,
redis_client: Client,
redis_client: RedisCommandClient,
room_subscribers: DashMap<String, mpsc::Receiver<Vec<u8>>>,
socket_rooms: DashMap<String, HashSet<String>>,
rooms: DashMap<String, HashSet<String>>,
@@ -83,7 +82,7 @@ pub struct RedisAdapter {
impl RedisAdapter {
pub fn new(
message_bus: Arc<dyn MessageBus>,
redis_client: Client,
redis_client: RedisCommandClient,
server_id: String,
namespace: String,
on_local_broadcast: LocalBroadcastFn,
@@ -195,12 +194,12 @@ impl Adapter for RedisAdapter {
let srk = socket_rooms_key(ns, sid);
self.redis_client
.sadd::<(), _, _>(&rk, sid)
.query::<()>(redis::cmd("SADD").arg(&rk).arg(sid))
.await
.map_err(|e| AdapterError::Redis(e.to_string()))?;
self.redis_client
.sadd::<(), _, _>(&srk, room)
.query::<()>(redis::cmd("SADD").arg(&srk).arg(room))
.await
.map_err(|e| AdapterError::Redis(e.to_string()))?;
@@ -241,12 +240,12 @@ impl Adapter for RedisAdapter {
let srk = socket_rooms_key(ns, sid);
self.redis_client
.srem::<(), _, _>(&rk, sid)
.query::<()>(redis::cmd("SREM").arg(&rk).arg(sid))
.await
.map_err(|e| AdapterError::Redis(e.to_string()))?;
self.redis_client
.srem::<(), _, _>(&srk, room)
.query::<()>(redis::cmd("SREM").arg(&srk).arg(room))
.await
.map_err(|e| AdapterError::Redis(e.to_string()))?;
@@ -308,7 +307,11 @@ impl Adapter for RedisAdapter {
}
let rk = room_key(ns, room);
if let Err(e) = self.redis_client.srem::<(), _, _>(&rk, sid).await {
if let Err(e) = self
.redis_client
.query::<()>(redis::cmd("SREM").arg(&rk).arg(sid))
.await
{
tracing::warn!("Redis SREM room error: {}", e);
}
}
@@ -316,7 +319,7 @@ impl Adapter for RedisAdapter {
let srk = socket_rooms_key(ns, sid);
self.redis_client
.del::<(), _>(&srk)
.query::<()>(redis::cmd("DEL").arg(&srk))
.await
.map_err(|e| AdapterError::Redis(e.to_string()))?;
socket/message_bus/redis.rs
+178 -53
View File
@@ -1,54 +1,106 @@
use async_trait::async_trait;
use fred::clients::{Client, SubscriberClient};
use fred::interfaces::{ClientLike, EventInterface, PubsubInterface};
use fred::prelude::*;
use futures_util::StreamExt;
use redis::aio::ConnectionManager;
use redis::cluster::ClusterClient;
use redis::cluster_async::ClusterConnection;
use redis::{Client, FromRedisValue};
use tokio::sync::mpsc;
use tokio::time::{Duration, timeout};
use crate::socket::message_bus::{MessageBus, MessageBusError};
const REDIS_CONNECT_TIMEOUT_SECS: u64 = 5;
#[derive(Clone)]
pub enum RedisCommandClient {
Single(ConnectionManager),
Cluster(ClusterConnection),
}
impl RedisCommandClient {
pub async fn query<T: FromRedisValue>(&self, cmd: &mut redis::Cmd) -> redis::RedisResult<T> {
match self {
Self::Single(conn) => {
let mut conn = conn.clone();
cmd.query_async(&mut conn).await
}
Self::Cluster(conn) => {
let mut conn = conn.clone();
cmd.query_async(&mut conn).await
}
}
}
}
pub struct RedisMessageBus {
client: Client,
subscriber: SubscriberClient,
command_client: RedisCommandClient,
pubsub_client: Client,
}
impl RedisMessageBus {
/// Connect to a Redis cluster.
/// Connect to Redis using the same `redis` crate as appks.
///
/// `cluster_url` should be in `redis-cluster://` format, e.g.:
/// `redis-cluster://host1:6379,host2:6379,host3:6379`
pub async fn new(cluster_url: &str) -> Result<Self, MessageBusError> {
let config =
Config::from_url(cluster_url).map_err(|e| MessageBusError::Redis(e.to_string()))?;
/// Supports both single-node `redis://host:port` and cluster
/// `redis-cluster://host1:6379?node=host2:6379` URLs.
pub async fn new(redis_url: &str) -> Result<Self, MessageBusError> {
tracing::info!("Connecting to Redis");
let client = Client::new(config.clone(), None, None, None);
let subscriber = SubscriberClient::new(config, None, None, None);
let connect_timeout = Duration::from_secs(REDIS_CONNECT_TIMEOUT_SECS);
let parsed = parse_redis_url(redis_url)?;
let _ = client.connect().await;
let _ = subscriber.connect().await;
let (command_client, pubsub_url) = match parsed {
ParsedRedisConfig::Single(url) => {
let client = Client::open(url.as_str())
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
let conn = timeout(connect_timeout, client.get_connection_manager())
.await
.map_err(|_| {
MessageBusError::Redis(format!(
"Redis connection timeout after {REDIS_CONNECT_TIMEOUT_SECS}s"
))
})?
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
(RedisCommandClient::Single(conn), url)
}
ParsedRedisConfig::Cluster(nodes) => {
let cluster_client = ClusterClient::new(nodes.iter().map(String::as_str).collect::<Vec<_>>())
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
let conn = timeout(connect_timeout, cluster_client.get_async_connection())
.await
.map_err(|_| {
MessageBusError::Redis(format!(
"Redis cluster connection timeout after {REDIS_CONNECT_TIMEOUT_SECS}s"
))
})?
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
let pubsub_url = nodes
.first()
.cloned()
.ok_or_else(|| MessageBusError::Redis("Redis cluster nodes are empty".into()))?;
(RedisCommandClient::Cluster(conn), pubsub_url)
}
};
client
.wait_for_connect()
.await
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
subscriber
.wait_for_connect()
.await
let pubsub_client = Client::open(pubsub_url.as_str())
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
tracing::info!(cluster_url, "Redis cluster connected");
Ok(Self { client, subscriber })
tracing::info!("Redis connected");
Ok(Self {
command_client,
pubsub_client,
})
}
pub fn client(&self) -> &Client {
&self.client
pub fn client(&self) -> RedisCommandClient {
self.command_client.clone()
}
}
#[async_trait]
impl MessageBus for RedisMessageBus {
async fn publish(&self, channel: &str, message: &[u8]) -> Result<(), MessageBusError> {
self.client
.publish::<(), _, Vec<u8>>(channel, message.to_vec())
self.command_client
.query::<()>(redis::cmd("PUBLISH").arg(channel).arg(message))
.await
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
Ok(())
@@ -56,23 +108,27 @@ impl MessageBus for RedisMessageBus {
async fn subscribe(&self, channel: &str) -> Result<mpsc::Receiver<Vec<u8>>, MessageBusError> {
let (tx, rx) = mpsc::channel::<Vec<u8>>(256);
self.subscriber
.subscribe(channel.to_string())
let mut pubsub = self
.pubsub_client
.get_async_pubsub()
.await
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
let subscriber = self.subscriber.clone();
let channel_owned = channel.to_string();
let mut message_rx = subscriber.message_rx();
pubsub
.subscribe(channel)
.await
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
let channel_owned = channel.to_string();
tokio::spawn(async move {
while let Ok(message) = message_rx.recv().await {
if message.channel == channel_owned {
let data: Vec<u8> = FromValue::from_value(message.value).unwrap_or_default();
if tx.send(data).await.is_err() {
break;
}
let mut stream = pubsub.on_message();
while let Some(message) = stream.next().await {
if message.get_channel_name() != channel_owned {
continue;
}
let payload = message.get_payload::<Vec<u8>>().unwrap_or_default();
if tx.send(payload).await.is_err() {
break;
}
}
});
@@ -80,23 +136,92 @@ impl MessageBus for RedisMessageBus {
Ok(rx)
}
async fn unsubscribe(&self, channel: &str) -> Result<(), MessageBusError> {
self.subscriber
.unsubscribe(channel.to_string())
.await
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
async fn unsubscribe(&self, _channel: &str) -> Result<(), MessageBusError> {
// Each subscription owns its dedicated async PubSub connection inside
// the spawned listener task. Dropping the receiver stops local delivery.
Ok(())
}
async fn close(&self) -> Result<(), MessageBusError> {
self.client
.quit()
.await
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
self.subscriber
.quit()
.await
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
Ok(())
}
}
enum ParsedRedisConfig {
Single(String),
Cluster(Vec<String>),
}
fn parse_redis_url(redis_url: &str) -> Result<ParsedRedisConfig, MessageBusError> {
let Some(rest) = redis_url.strip_prefix("redis-cluster://") else {
return Ok(ParsedRedisConfig::Single(redis_url.to_string()));
};
let (first, query) = rest.split_once('?').unwrap_or((rest, ""));
let (auth, first_node) = split_auth(first);
let mut nodes = Vec::new();
if !first_node.is_empty() {
nodes.push(to_redis_node_url(auth, first_node));
}
for part in query.split('&') {
let Some(node) = part.strip_prefix("node=") else {
continue;
};
if !node.is_empty() {
nodes.push(to_redis_node_url(auth, node));
}
}
if nodes.is_empty() {
return Err(MessageBusError::Redis("Redis cluster URL has no nodes".into()));
}
Ok(ParsedRedisConfig::Cluster(nodes))
}
fn split_auth(value: &str) -> (&str, &str) {
value
.rsplit_once('@')
.map(|(auth, host)| (auth, host))
.unwrap_or(("", value))
}
fn to_redis_node_url(auth: &str, node: &str) -> String {
if node.starts_with("redis://") || node.starts_with("rediss://") {
return node.to_string();
}
if auth.is_empty() {
format!("redis://{node}")
} else {
format!("redis://{auth}@{node}")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_single_redis_url() {
match parse_redis_url("redis://127.0.0.1:6379").unwrap() {
ParsedRedisConfig::Single(url) => assert_eq!(url, "redis://127.0.0.1:6379"),
ParsedRedisConfig::Cluster(_) => panic!("expected single redis config"),
}
}
#[test]
fn parses_cluster_url_for_redis_rs() {
match parse_redis_url("redis-cluster://:pass@127.0.0.1:6380?node=127.0.0.1:6381").unwrap() {
ParsedRedisConfig::Cluster(nodes) => assert_eq!(
nodes,
vec![
"redis://:pass@127.0.0.1:6380".to_string(),
"redis://:pass@127.0.0.1:6381".to_string()
]
),
ParsedRedisConfig::Single(_) => panic!("expected cluster redis config"),
}
}
}
socket/session_store/redis.rs
+40 -28
View File
@@ -1,8 +1,8 @@
use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use fred::prelude::*;
use crate::socket::message_bus::redis::RedisCommandClient;
use crate::socket::message_bus::redis::RedisMessageBus;
use crate::socket::session_store::{SessionError, SessionInfo, SessionStoreTrait};
@@ -17,14 +17,14 @@ const DEFAULT_TTL_SECS: u64 = 60;
const KEY_PREFIX: &str = "socket.io:session";
pub struct RedisSessionStore {
client: Client,
client: RedisCommandClient,
ttl_secs: u64,
}
impl RedisSessionStore {
pub fn new(bus: &RedisMessageBus, ttl_secs: Option<u64>) -> Self {
Self {
client: bus.client().clone(),
client: bus.client(),
ttl_secs: ttl_secs.unwrap_or(DEFAULT_TTL_SECS),
}
}
@@ -45,22 +45,29 @@ impl SessionStoreTrait for RedisSessionStore {
let key = self.key(sid);
let now = now_millis();
// Batch all fields in a single HSET call for efficiency
let fields: Vec<(&str, String)> = vec![
("sid", sid.to_string()),
("transport", transport.to_string()),
("state", "connecting".to_string()),
("server_id", server_id.to_string()),
("created_at", now.to_string()),
("last_ping", now.to_string()),
];
// Batch all fields in a single HMSET-style call
self.client
.hset::<(), _, _>(&key, fields)
.query::<()>(
redis::cmd("HSET")
.arg(&key)
.arg("sid")
.arg(sid)
.arg("transport")
.arg(transport)
.arg("state")
.arg("connecting")
.arg("server_id")
.arg(server_id)
.arg("created_at")
.arg(now.to_string())
.arg("last_ping")
.arg(now.to_string()),
)
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
self.client
.expire::<(), _>(&key, self.ttl_secs as i64, None)
.query::<()>(redis::cmd("EXPIRE").arg(&key).arg(self.ttl_secs as i64))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
@@ -70,11 +77,9 @@ impl SessionStoreTrait for RedisSessionStore {
async fn get(&self, sid: &str) -> Result<Option<SessionInfo>, SessionError> {
let key = self.key(sid);
// Use hgetall directly — if the key doesn't exist Redis returns an empty map.
// This avoids the TOCTOU race between EXISTS and HGETALL.
let values: std::collections::HashMap<String, String> = self
.client
.hgetall::<std::collections::HashMap<String, String>, _>(&key)
.query(redis::cmd("HGETALL").arg(&key))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
@@ -103,14 +108,13 @@ impl SessionStoreTrait for RedisSessionStore {
async fn set_state(&self, sid: &str, state: &str) -> Result<(), SessionError> {
let key = self.key(sid);
// Use HSET (not HSETNX) to overwrite existing fields
self.client
.hset::<(), _, _>(&key, ("state", state))
.query::<()>(redis::cmd("HSET").arg(&key).arg("state").arg(state))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
self.client
.expire::<(), _>(&key, self.ttl_secs as i64, None)
.query::<()>(redis::cmd("EXPIRE").arg(&key).arg(self.ttl_secs as i64))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
@@ -120,14 +124,18 @@ impl SessionStoreTrait for RedisSessionStore {
async fn set_transport(&self, sid: &str, transport: &str) -> Result<(), SessionError> {
let key = self.key(sid);
// Use HSET (not HSETNX) to overwrite existing fields
self.client
.hset::<(), _, _>(&key, ("transport", transport))
.query::<()>(
redis::cmd("HSET")
.arg(&key)
.arg("transport")
.arg(transport),
)
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
self.client
.expire::<(), _>(&key, self.ttl_secs as i64, None)
.query::<()>(redis::cmd("EXPIRE").arg(&key).arg(self.ttl_secs as i64))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
@@ -138,14 +146,18 @@ impl SessionStoreTrait for RedisSessionStore {
let key = self.key(sid);
let now = now_millis();
// Use HSET (not HSETNX) to overwrite existing fields
self.client
.hset::<(), _, _>(&key, ("last_ping", now.to_string()))
.query::<()>(
redis::cmd("HSET")
.arg(&key)
.arg("last_ping")
.arg(now.to_string()),
)
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
self.client
.expire::<(), _>(&key, self.ttl_secs as i64, None)
.query::<()>(redis::cmd("EXPIRE").arg(&key).arg(self.ttl_secs as i64))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
@@ -156,7 +168,7 @@ impl SessionStoreTrait for RedisSessionStore {
let key = self.key(sid);
self.client
.del::<(), _>(&key)
.query::<()>(redis::cmd("DEL").arg(&key))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
@@ -168,7 +180,7 @@ impl SessionStoreTrait for RedisSessionStore {
let exists: bool = self
.client
.exists::<bool, _>(&key)
.query(redis::cmd("EXISTS").arg(&key))
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;