072c704
use std::sync::Arc;
use std::time::Duration;
use actor::ActorHost;
use serde_json::Value;
use socket::{SocketEvent, SocketHub, SocketPayload};
use tokio::net::TcpListener;
use tokio::sync::{mpsc, watch};
use crate::gateway::action::EventProperties;
use crate::gateway::platform::PlatformClient;
use crate::gateway::socket_backend::RealtimeBackend;
use crate::gateway::subscriptions::{SubscriptionRequest, Subscriptions};
use crate::gateway::{ConsumeResult, EventsPipe, action, is_personal};
const SHUTDOWN_DRAIN: Duration = Duration::from_secs(10);
pub struct EventActor;
impl actor::VirtualActor for EventActor {
type Message = Value;
type Response = Value;
fn receive<'a>(
&'a mut self,
context: &'a actor::ActorContext,
message: Self::Message,
) -> impl std::future::Future<Output = Result<Self::Response, actor::ActorError>> + Send + 'a
{
async move {
context
.host()
.publish_event(context.id().user_id(), message.clone())
.await?;
Ok(message)
}
}
}
pub fn backend(
host: ActorHost,
platform: PlatformClient,
subscriptions: Subscriptions,
push_store: actor::PushStore,
) -> Arc<RealtimeBackend> {
Arc::new(RealtimeBackend::new(
host,
platform,
subscriptions,
push_store,
))
}
async fn handle_socket_event(
event: &SocketEvent,
host: &ActorHost,
platform: &PlatformClient,
subscriptions: &Subscriptions,
) -> Result<Value, String> {
match event.event.as_str() {
"forget-device" => {
if event.data.len() != 1 || !event.binaries.is_empty() {
return Err("forget-device requires one JSON argument".to_string());
}
let device = event.data[0]
.get("device_id")
.and_then(Value::as_str)
.ok_or("device_id is required")?;
let device = uuid::Uuid::parse_str(device)
.map_err(|_| "device_id must be a UUID")?
.to_string();
host.push_store()
.forget_device(event.user_id, &device)
.await?;
Ok(Value::Null)
}
"subscribe" | "unsubscribe" => {
if !event.binaries.is_empty() || event.data.len() != 1 {
return Err("subscription requires one JSON argument".to_string());
}
let request: SubscriptionRequest =
serde_json::from_value(event.data[0].clone()).map_err(|error| error.to_string())?;
if event.event == "unsubscribe" {
let id = request.id.as_deref().ok_or("subscription id is required")?;
subscriptions
.unsubscribe(
event.user_id,
&event.device_id,
&event.namespace,
&request.topic,
&request.tag,
id,
)
.await?;
return Ok(Value::Null);
}
let id = subscriptions
.subscribe(
platform,
event.user_id,
&event.device_id,
&event.namespace,
request,
)
.await?;
Ok(serde_json::json!({"id":id}))
}
_ => {
let payload = serde_json::json!({
"source": "socket",
"event_id": uuid::Uuid::new_v4().to_string(),
"namespace": event.namespace,
"session_id": event.session_id,
"user_id": event.user_id,
"event": event.event,
"data": event.data,
"binaries": event.binaries,
});
host.actor::<EventActor>(event.user_id)
.call(payload)
.await
.map(|_| Value::Null)
.map_err(|error| error.to_string())
}
}
}
/// 启动 socket 节点：浏览器上行事件 + actor 唤醒投递。
pub async fn run_socket_node(
hub: SocketHub,
mut events: mpsc::Receiver<SocketEvent>,
host: ActorHost,
platform: PlatformClient,
subscriptions: Subscriptions,
listen_address: &str,
shutdown: watch::Receiver<bool>,
) -> Result<(), String> {
let mut deliveries = host
.take_events()
.await
.ok_or_else(|| "actor events already attached".to_string())?;
let listener = TcpListener::bind(listen_address)
.await
.map_err(|error| format!("binding socket listener: {error}"))?;
let events_host = host.clone();
let events_platform = platform.clone();
let events_hub = hub.clone();
let events_task = tokio::spawn(async move {
let mut lanes = Vec::new();
let mut workers = tokio::task::JoinSet::new();
for _ in 0..16 {
let (sender, mut receiver) = mpsc::channel::<SocketEvent>(64);
lanes.push(sender);
let host = events_host.clone();
let platform = events_platform.clone();
let hub = events_hub.clone();
let subscriptions = subscriptions.clone();
workers.spawn(async move {
while let Some(event) = receiver.recv().await {
let result =
handle_socket_event(&event, &host, &platform, &subscriptions).await;
match (event.ack, result) {
(Some(ack), Ok(result)) => {
let _ = ack.send(SocketPayload {
data: vec![serde_json::json!({"ok":true,"data":result})],
binaries: Vec::new(),
});
}
(Some(ack), Err(error)) => {
let _ = ack.send(SocketPayload {
data: vec![serde_json::json!({"ok":false,"error":error})],
binaries: Vec::new(),
});
}
(None, Err(error)) => {
hub.emit(
event.user_id,
&event.namespace,
"error",
vec![serde_json::json!({"message":error})],
)
.await;
}
_ => {}
}
}
});
}
while let Some(event) = events.recv().await {
let lane = (event.user_id as u64 % 16) as usize;
if let Err(error) = lanes[lane].try_send(event) {
let event = error.into_inner();
if let Some(ack) = event.ack {
let _ = ack.send(SocketPayload {
data: vec![
serde_json::json!({"ok":false,"error":"user event queue is full"}),
],
binaries: Vec::new(),
});
} else {
events_hub
.emit(
event.user_id,
&event.namespace,
"error",
vec![serde_json::json!({"message":"user event queue is full"})],
)
.await;
}
}
}
drop(lanes);
while workers.join_next().await.is_some() {}
});
let delivery_host = host.clone();
let delivery_hub = hub.clone();
let delivery_task = tokio::spawn(async move {
while let Some(event) = deliveries.recv().await {
let ready = if event.payload.get("wake").and_then(Value::as_bool) == Some(true) {
true
} else {
let mut payload = event.payload.clone();
if payload.get("source").and_then(Value::as_str) == Some("socket")
&& payload.get("event_id").is_none()
{
match delivery_host
.push_store()
.legacy_event_id(event.user_id, delivery_host.node_id(), event.entry_id())
.await
{
Ok(id) => {
payload["event_id"] = Value::String(id);
}
Err(_) => {
delivery_host.retry_event(&event).await;
continue;
}
}
}
delivery_host
.push_store()
.append(event.user_id, payload)
.await
.is_ok()
};
if ready {
delivery_hub.wake_user(event.user_id).await;
if let Err(error) = delivery_host.complete_event(&event).await {
tracing::warn!(%error, "failed to acknowledge push wake");
}
} else {
delivery_host.retry_event(&event).await;
}
}
});
let result = hub.serve(listener, shutdown).await;
hub.shutdown();
events_task.abort();
delivery_task.abort();
let _ = events_task.await;
let _ = delivery_task.await;
result.map_err(|error| error.to_string())
}
/// RocketMQ 消费节点：16 lane 分片 + action 扇出。
pub async fn run_consumer(
config: &config::event::consumer::RocketMqConsumerConfig,
runtime: &queue::QueueRuntime,
host: ActorHost,
platform: PlatformClient,
subscriptions: Subscriptions,
mut shutdown: watch::Receiver<bool>,
) -> Result<(), String> {
let mut pipe = EventsPipe::new(config, runtime)?;
pipe.subscribe_all()
.await
.map_err(|error| error.to_string())?;
pipe.start().await.map_err(|error| error.to_string())?;
let (results, mut completed) = mpsc::channel(128);
let mut lanes = Vec::new();
let mut workers = tokio::task::JoinSet::new();
for _ in 0..16 {
let (sender, mut receiver) = mpsc::channel(64);
lanes.push(sender);
let host = host.clone();
let platform = platform.clone();
let subscriptions = subscriptions.clone();
let results = results.clone();
workers.spawn(async move {
let context = action::EventContext {
host: &host,
platform: &platform,
subscriptions: &subscriptions,
};
while let Some(message) = receiver.recv().await {
let result = action::action(message, &context).await;
if results.send(result).await.is_err() {
break;
}
}
});
}
drop(results);
if !*shutdown.borrow() {
loop {
tokio::select! {
_ = shutdown.changed() => break,
incoming = pipe.next_event() => {
let Some(message) = incoming else { break };
let lane = event_lane(&message);
if let Err(error) = lanes[lane].try_send(message) {
pipe.apply_result(ConsumeResult::Retry {message: error.into_inner(),delay_level:0}).await.map_err(|error| error.to_string())?;
}
}
result = completed.recv() => {
let Some(result) = result else { break };
pipe.apply_result(result).await.map_err(|error| error.to_string())?;
}
}
}
}
pipe.shutdown().await;
drop(lanes);
let drained = tokio::time::timeout(SHUTDOWN_DRAIN, async {
loop {
tokio::select! {
result = completed.recv() => {
let Some(result) = result else { break };
let _ = pipe.apply_result(result).await;
}
joined = workers.join_next() => {
if joined.is_none() { break }
}
}
}
})
.await;
if drained.is_err() {
workers.abort_all();
while workers.join_next().await.is_some() {}
}
Ok(())
}
fn event_lane(message: &rocketmq_model::common::message::message_ext::MessageExt) -> usize {
use rocketmq_model::common::message::MessageTrait;
use std::hash::{Hash, Hasher};
let properties = message.properties();
let topic = message.topic();
let tag = message
.tags_ref()
.map(ToString::to_string)
.unwrap_or_default();
let fields: &[&str] = if is_personal(topic.as_str(), &tag) {
&["recipient", "user"]
} else {
&["room_id", "issue_id", "pull_id", "repo_id", "org", "user"]
};
let key = fields.iter().find_map(|field| properties.get(*field));
let mut hasher = std::collections::hash_map::DefaultHasher::new();
topic.hash(&mut hasher);
key.map(|key| key.as_str()).hash(&mut hasher);
(hasher.finish() % 16) as usize
}
pub fn _event_properties_touch(_: &EventProperties) {}