use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use tokio::sync::RwLock;
use crate::dao::SaTokenDao;
use crate::error::{SaTokenError, SaTokenResult};
use crate::keys::LOGIN_TYPE_DEFAULT;
mod push;
mod store;
pub use push::dispatch_to_pushers;
pub use store::{DistributedOnlineStore, LocalOnlineStore, OnlineStore, StoredOnlineUser};
#[derive(Debug, Clone)]
pub struct OnlineUser {
pub login_type: String,
pub login_id: String,
pub token: String,
pub device: String,
pub connect_time: DateTime<Utc>,
pub last_activity: DateTime<Utc>,
pub metadata: HashMap<String, String>,
}
impl OnlineUser {
pub fn new(
login_id: impl Into<String>,
token: impl Into<String>,
device: impl Into<String>,
) -> Self {
let now = Utc::now();
Self {
login_type: LOGIN_TYPE_DEFAULT.to_string(),
login_id: login_id.into(),
token: token.into(),
device: device.into(),
connect_time: now,
last_activity: now,
metadata: HashMap::new(),
}
}
}
#[derive(Debug, Clone)]
pub struct PushMessage {
pub message_id: String,
pub content: String,
pub message_type: MessageType,
pub timestamp: DateTime<Utc>,
pub metadata: HashMap<String, String>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum MessageType {
Text,
Binary,
KickOut,
Notification,
Custom(String),
}
#[async_trait]
pub trait MessagePusher: Send + Sync {
async fn push(&self, login_id: &str, message: PushMessage) -> Result<(), SaTokenError>;
}
pub struct OnlineManager {
store: Arc<dyn OnlineStore>,
pushers: Arc<RwLock<Vec<Arc<dyn MessagePusher>>>>,
}
impl std::fmt::Debug for OnlineManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("OnlineManager { .. }")
}
}
impl OnlineManager {
pub fn new() -> Self {
Self::local()
}
pub fn local() -> Self {
Self {
store: Arc::new(LocalOnlineStore::new()),
pushers: Arc::new(RwLock::new(Vec::new())),
}
}
pub fn distributed(dao: Arc<SaTokenDao>) -> Self {
Self {
store: Arc::new(DistributedOnlineStore::new(
dao,
Some(Duration::from_secs(86400)),
)),
pushers: Arc::new(RwLock::new(Vec::new())),
}
}
pub fn with_store(store: Arc<dyn OnlineStore>) -> Self {
Self {
store,
pushers: Arc::new(RwLock::new(Vec::new())),
}
}
pub async fn register_pusher(&self, pusher: Arc<dyn MessagePusher>) {
self.pushers.write().await.push(pusher);
}
pub async fn mark_online(&self, user: OnlineUser) -> SaTokenResult<()> {
self.store.mark_online(user).await
}
pub async fn mark_offline(&self, login_id: &str, token: &str) -> SaTokenResult<()> {
self.store
.mark_offline(LOGIN_TYPE_DEFAULT, login_id, token)
.await
}
pub async fn mark_offline_with_type(
&self,
login_type: &str,
login_id: &str,
token: &str,
) -> SaTokenResult<()> {
self.store.mark_offline(login_type, login_id, token).await
}
pub async fn mark_offline_all(&self, login_id: &str) -> SaTokenResult<()> {
self.store
.mark_offline_all(LOGIN_TYPE_DEFAULT, login_id)
.await
}
pub async fn mark_offline_all_with_type(
&self,
login_type: &str,
login_id: &str,
) -> SaTokenResult<()> {
self.store.mark_offline_all(login_type, login_id).await
}
pub async fn is_online(&self, login_id: &str) -> SaTokenResult<bool> {
self.store.is_online(LOGIN_TYPE_DEFAULT, login_id).await
}
pub async fn get_online_count(&self) -> SaTokenResult<usize> {
self.store.get_online_count().await
}
pub async fn get_online_users(&self) -> SaTokenResult<Vec<String>> {
self.store.get_online_users().await
}
pub async fn get_user_sessions(&self, login_id: &str) -> SaTokenResult<Vec<OnlineUser>> {
self.store
.get_user_sessions(LOGIN_TYPE_DEFAULT, login_id)
.await
}
pub async fn update_activity(&self, login_id: &str, token: &str) -> SaTokenResult<()> {
self.store
.update_activity(LOGIN_TYPE_DEFAULT, login_id, token)
.await
}
pub async fn update_activity_with_type(
&self,
login_type: &str,
login_id: &str,
token: &str,
) -> SaTokenResult<()> {
self.store
.update_activity(login_type, login_id, token)
.await
}
async fn cloned_pushers(&self) -> Vec<Arc<dyn MessagePusher>> {
self.pushers.read().await.clone()
}
pub async fn push_to_user(&self, login_id: &str, content: String) -> SaTokenResult<()> {
let message = PushMessage {
message_id: uuid::Uuid::new_v4().to_string(),
content,
message_type: MessageType::Text,
timestamp: Utc::now(),
metadata: HashMap::new(),
};
let pushers = self.cloned_pushers().await;
dispatch_to_pushers(&pushers, login_id, message).await
}
pub async fn push_to_users(
&self,
login_ids: Vec<String>,
content: String,
) -> SaTokenResult<()> {
for login_id in login_ids {
self.push_to_user(&login_id, content.clone()).await?;
}
Ok(())
}
pub async fn broadcast(&self, content: String) -> SaTokenResult<()> {
let login_ids = self.get_online_users().await?;
self.push_to_users(login_ids, content).await
}
pub async fn push_message_to_user(
&self,
login_id: &str,
message: PushMessage,
) -> SaTokenResult<()> {
let pushers = self.cloned_pushers().await;
dispatch_to_pushers(&pushers, login_id, message).await
}
pub async fn kick_out_notify(&self, login_id: &str, reason: String) -> SaTokenResult<()> {
let message = PushMessage {
message_id: uuid::Uuid::new_v4().to_string(),
content: reason,
message_type: MessageType::KickOut,
timestamp: Utc::now(),
metadata: HashMap::new(),
};
self.push_message_to_user(login_id, message).await?;
self.mark_offline_all(login_id).await
}
}
impl Default for OnlineManager {
fn default() -> Self {
Self::new()
}
}
pub struct InMemoryPusher {
messages: Arc<RwLock<HashMap<String, Vec<PushMessage>>>>,
}
impl std::fmt::Debug for InMemoryPusher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("InMemoryPusher { .. }")
}
}
impl InMemoryPusher {
pub fn new() -> Self {
Self {
messages: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn get_messages(&self, login_id: &str) -> Vec<PushMessage> {
self.messages
.read()
.await
.get(login_id)
.cloned()
.unwrap_or_default()
}
pub async fn clear_messages(&self, login_id: &str) {
self.messages.write().await.remove(login_id);
}
}
impl Default for InMemoryPusher {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl MessagePusher for InMemoryPusher {
async fn push(&self, login_id: &str, message: PushMessage) -> Result<(), SaTokenError> {
self.messages
.write()
.await
.entry(login_id.to_string())
.or_default()
.push(message);
Ok(())
}
}