Skip to main content

flare_core/server/connection/manager/
ops.rs

1use super::*;
2
3impl ConnectionManager {
4    /// 获取用户的所有连接
5    ///
6    /// # 参数
7    /// - `user_id`: 用户 ID
8    ///
9    /// # 返回
10    /// 该用户的所有连接 ID 列表
11    pub fn get_user_connections(&self, user_id: &str) -> Vec<String> {
12        self.user_connection_shard(user_id)
13            .read()
14            .ok()
15            .and_then(|user_connections| user_connections.get(user_id).cloned())
16            .unwrap_or_default()
17    }
18
19    /// 更新连接的用户 ID（用于认证后绑定用户）
20    ///
21    /// # 参数
22    /// - `connection_id`: 连接 ID
23    /// - `user_id`: 新的用户 ID
24    pub fn bind_user(&self, connection_id: &str, user_id: String) -> Result<()> {
25        let mut shard = self
26            .connection_shard(connection_id)
27            .write()
28            .map_err(|_| FlareError::general_error("Failed to lock connection shard"))?;
29
30        let (_, _, info) = shard.get_mut(connection_id).ok_or_else(|| {
31            FlareError::protocol_error(format!("Connection {} not found", connection_id))
32        })?;
33
34        // 如果之前有用户 ID，先移除旧映射
35        if let Some(old_user_id) = &info.user_id {
36            self.remove_user_connection(old_user_id, connection_id)?;
37        }
38
39        // 更新用户 ID
40        info.user_id = Some(user_id.clone());
41
42        // 添加到新用户映射
43        self.insert_user_connection(user_id, connection_id)?;
44
45        Ok(())
46    }
47
48    /// 更新连接的最后活跃时间
49    pub fn update_connection_active(&self, connection_id: &str) -> Result<()> {
50        let mut shard = self
51            .connection_shard(connection_id)
52            .write()
53            .map_err(|_| FlareError::general_error("Failed to lock connection shard"))?;
54
55        if let Some((_, _, info)) = shard.get_mut(connection_id) {
56            info.update_active();
57            drop(shard);
58            Ok(())
59        } else {
60            drop(shard);
61            Err(FlareError::protocol_error(format!(
62                "Connection {} not found",
63                connection_id
64            )))
65        }
66    }
67
68    pub(super) fn update_connections_active<I>(&self, connection_ids: I) -> usize
69    where
70        I: IntoIterator<Item = String>,
71    {
72        let now = Instant::now();
73        let mut ids_by_shard = vec![Vec::new(); self.connection_shards.len()];
74        for connection_id in connection_ids {
75            let shard_index = self.connection_shard_index(&connection_id);
76            ids_by_shard[shard_index].push(connection_id);
77        }
78
79        ids_by_shard
80            .into_iter()
81            .enumerate()
82            .map(|(shard_index, ids)| {
83                if ids.is_empty() {
84                    return 0;
85                }
86
87                let Ok(mut shard) = self.connection_shards[shard_index].write() else {
88                    return 0;
89                };
90
91                let mut updated = 0;
92                for connection_id in ids {
93                    if let Some((_, _, info)) = shard.get_mut(&connection_id) {
94                        info.last_active = now;
95                        updated += 1;
96                    }
97                }
98                updated
99            })
100            .sum()
101    }
102
103    pub(super) async fn fanout_serialized_frame_to_auth_snapshots(
104        &self,
105        snapshots: Vec<ConnectionAuthSnapshot>,
106        frame: &crate::common::protocol::Frame,
107        data: &[u8],
108        fanout: &'static str,
109    ) -> Vec<String> {
110        stream::iter(snapshots)
111            .map(|snapshot| async move {
112                let connection_id = snapshot.0.clone();
113                match self
114                    .send_serialized_frame_to_auth_snapshot_without_active(snapshot, frame, data)
115                    .await
116                {
117                    Ok(connection_id) => Some(connection_id),
118                    Err(error) => {
119                        tracing::warn!(
120                            connection_id = %connection_id,
121                            error = ?error,
122                            fanout,
123                            "Connection frame fanout failed"
124                        );
125                        None
126                    }
127                }
128            })
129            .buffer_unordered(self.fanout_concurrency)
130            .filter_map(|connection_id| async move { connection_id })
131            .collect()
132            .await
133    }
134
135    /// 多连接同帧扇出：无加密连接按（序列化格式, 压缩）分组，**每组只序列化一次**共享字节；
136    /// 带加密连接回退逐连接序列化（密文可能与连接态绑定，不可跨连接共享）。
137    /// 群聊单节点 N 订阅者的下行从 N 次 serialize+compress → 组数次（通常 1）。
138    /// 返回（成功数, 失败数），成功连接统一批量刷新活跃时间。
139    pub(super) async fn fanout_frame_grouped_by_encoding(
140        &self,
141        connection_ids: Vec<String>,
142        frame: &crate::common::protocol::Frame,
143    ) -> (i32, i32) {
144        use std::collections::HashMap;
145
146        let snapshots = self.connection_snapshots_for_ids(connection_ids);
147        let total = snapshots.len();
148        if total == 0 {
149            return (0, 0);
150        }
151
152        let mut plain_groups: HashMap<
153            (
154                crate::common::protocol::SerializationFormat,
155                crate::common::compression::CompressionAlgorithm,
156            ),
157            Vec<ConnectionSnapshot>,
158        > = HashMap::new();
159        let mut encrypted: Vec<ConnectionSnapshot> = Vec::new();
160        for snapshot in snapshots {
161            let info = &snapshot.2;
162            if info.encryption == crate::common::encryption::EncryptionAlgorithm::None {
163                plain_groups
164                    .entry((info.serialization_format, info.compression.clone()))
165                    .or_default()
166                    .push(snapshot);
167            } else {
168                encrypted.push(snapshot);
169            }
170        }
171
172        let mut successful: Vec<String> = Vec::with_capacity(total);
173        for ((format, compression), group) in plain_groups {
174            let parser = crate::common::MessageParser::new(
175                format,
176                compression,
177                crate::common::encryption::EncryptionAlgorithm::None,
178            );
179            let data = match parser.serialize(frame) {
180                Ok(data) => data,
181                Err(error) => {
182                    tracing::warn!(
183                        error = ?error,
184                        group_size = group.len(),
185                        "grouped frame serialization failed; group counted as failures"
186                    );
187                    continue;
188                }
189            };
190            let sent: Vec<String> = stream::iter(group)
191                .map(|snapshot| {
192                    let data = &data;
193                    async move {
194                        let (connection_id, connection, info) = snapshot;
195                        if let Err(error) =
196                            Self::ensure_frame_allowed_for_connection(&connection_id, &info, frame)
197                        {
198                            tracing::warn!(connection_id = %connection_id, error = ?error, "grouped fanout skipped");
199                            return None;
200                        }
201                        match self
202                            .send_to_connection_handle(&connection_id, connection, data)
203                            .await
204                        {
205                            Ok(()) => Some(connection_id),
206                            Err(error) => {
207                                tracing::warn!(connection_id = %connection_id, error = ?error, "grouped fanout write failed");
208                                None
209                            }
210                        }
211                    }
212                })
213                .buffer_unordered(self.fanout_concurrency)
214                .filter_map(|connection_id| async move { connection_id })
215                .collect()
216                .await;
217            successful.extend(sent);
218        }
219        if !encrypted.is_empty() {
220            let sent = self
221                .fanout_frame_to_snapshots(encrypted, frame, None, "grouped_fanout_encrypted")
222                .await;
223            successful.extend(sent);
224        }
225
226        let success = successful.len() as i32;
227        self.update_connections_active(successful);
228        (success, (total as i32).saturating_sub(success))
229    }
230
231    pub(super) async fn fanout_frame_to_snapshots(
232        &self,
233        snapshots: Vec<ConnectionSnapshot>,
234        frame: &crate::common::protocol::Frame,
235        parser: Option<&crate::common::MessageParser>,
236        fanout: &'static str,
237    ) -> Vec<String> {
238        stream::iter(snapshots)
239            .map(|snapshot| async move {
240                let connection_id = snapshot.0.clone();
241                match self
242                    .send_frame_to_snapshot_without_active(snapshot, frame, parser)
243                    .await
244                {
245                    Ok(connection_id) => Some(connection_id),
246                    Err(error) => {
247                        tracing::warn!(
248                            connection_id = %connection_id,
249                            error = ?error,
250                            fanout,
251                            "Connection frame fanout failed"
252                        );
253                        None
254                    }
255                }
256            })
257            .buffer_unordered(self.fanout_concurrency)
258            .filter_map(|connection_id| async move { connection_id })
259            .collect()
260            .await
261    }
262
263    /// 设置连接为已验证状态
264    pub fn set_connection_authenticated(
265        &self,
266        connection_id: &str,
267        user_id: Option<String>,
268    ) -> Result<()> {
269        let mut shard = self
270            .connection_shard(connection_id)
271            .write()
272            .map_err(|_| FlareError::general_error("Failed to lock connection shard"))?;
273
274        let (_, _, info) = shard.get_mut(connection_id).ok_or_else(|| {
275            FlareError::protocol_error(format!("Connection {} not found", connection_id))
276        })?;
277
278        // 保存旧的 user_id（在调用 set_authenticated 之前）
279        let old_user_id = info.user_id.clone();
280
281        let final_user_id = user_id.or(old_user_id.clone());
282
283        // 设置认证状态（如果 final_user_id 是 Some，会设置 user_id）
284        info.set_authenticated(final_user_id.clone());
285
286        // 如果有 user_id（传入的或已存在的），确保用户连接映射正确
287        if let Some(user_id) = final_user_id {
288            // 如果 user_id 发生变化，需要更新映射
289            let user_id_changed = old_user_id
290                .as_ref()
291                .map(|old| old != &user_id)
292                .unwrap_or(true);
293
294            if user_id_changed {
295                // 如果之前有旧用户 ID，先移除旧映射
296                if let Some(old_user_id) = old_user_id {
297                    self.remove_user_connection(&old_user_id, connection_id)?;
298                }
299
300                // 添加新映射（检查是否已存在，避免重复）
301                self.insert_user_connection(user_id, connection_id)?;
302            } else {
303                // user_id 没有变化，只需确保映射存在
304                self.insert_user_connection(user_id, connection_id)?;
305            }
306        }
307
308        Ok(())
309    }
310
311    /// 更新连接的协商信息（设备信息、序列化格式、压缩算法）
312    #[allow(clippy::too_many_arguments)]
313    pub fn update_connection_negotiation(
314        &self,
315        connection_id: &str,
316        device_info: Option<crate::common::device::DeviceInfo>,
317        serialization_format: crate::common::protocol::SerializationFormat,
318        compression: crate::common::compression::CompressionAlgorithm,
319        encryption: crate::common::encryption::EncryptionAlgorithm,
320        user_id: Option<String>,
321        metadata: Option<HashMap<String, String>>,
322    ) -> Result<()> {
323        let mut shard = self
324            .connection_shard(connection_id)
325            .write()
326            .map_err(|_| FlareError::general_error("Failed to lock connection shard"))?;
327
328        let (_, _, info) = shard.get_mut(connection_id).ok_or_else(|| {
329            FlareError::protocol_error(format!("Connection {} not found", connection_id))
330        })?;
331
332        // 更新协商信息（但不标记协商完成）
333        // 协商完成将在 CONNECT_ACK 发送完成后由 update_connection_negotiation_with_pipeline 标记
334        info.device_info = device_info;
335        info.serialization_format = serialization_format;
336        info.compression = compression;
337        info.encryption = encryption;
338        // 若有传入 metadata，将其所有键值合并到 ConnectionInfo.metadata（同 key 覆盖）
339        if let Some(meta) = metadata {
340            for (k, v) in meta {
341                info.metadata.insert(k, v);
342            }
343        }
344        // 注意：这里不设置 negotiation_completed = true
345        // 也不创建 cached_parser，这些将在 CONNECT_ACK 发送完成后设置
346
347        // 保存旧的 user_id（在修改之前）
348        let old_user_id = info.user_id.clone();
349
350        let user_id_to_set = user_id.clone().or(old_user_id.clone());
351
352        // 添加调试日志
353        if user_id_to_set.is_none() {
354            tracing::trace!(connection_id = %connection_id,incoming_user_id = ?user_id,old_user_id = ?old_user_id,"update_connection_negotiation: user_id_to_set is None, user_id will not be set");
355        }
356
357        if let Some(user_id_val) = user_id_to_set {
358            // 如果之前有用户 ID 且与新 user_id 不同，先移除旧映射
359            if let Some(old_user_id) = old_user_id
360                && old_user_id != user_id_val
361            {
362                self.remove_user_connection(&old_user_id, connection_id)?;
363            }
364
365            // 更新用户 ID
366            info.user_id = Some(user_id_val.clone());
367
368            // 添加到新用户映射
369            self.insert_user_connection(user_id_val, connection_id)?;
370        }
371
372        Ok(())
373    }
374
375    /// 更新连接的协商信息（设备信息、序列化格式、压缩算法、加密方式）并设置 pipeline
376    #[allow(clippy::too_many_arguments)]
377    pub fn update_connection_negotiation_with_pipeline(
378        &self,
379        connection_id: &str,
380        device_info: Option<crate::common::device::DeviceInfo>,
381        serialization_format: crate::common::protocol::SerializationFormat,
382        compression: crate::common::compression::CompressionAlgorithm,
383        encryption: crate::common::encryption::EncryptionAlgorithm,
384        user_id: Option<String>,
385        parser: crate::common::MessageParser,
386        pipeline: Option<std::sync::Arc<crate::common::message::pipeline::MessagePipeline>>,
387    ) -> Result<()> {
388        let mut shard = self
389            .connection_shard(connection_id)
390            .write()
391            .map_err(|_| FlareError::general_error("Failed to lock connection shard"))?;
392
393        let (_, _, info) = shard.get_mut(connection_id).ok_or_else(|| {
394            FlareError::protocol_error(format!("Connection {} not found", connection_id))
395        })?;
396
397        // 更新协商信息
398        info.device_info = device_info;
399        info.serialization_format = serialization_format;
400        info.compression = compression;
401        info.encryption = encryption;
402        // 标记协商已完成
403        info.negotiation_completed = true;
404        // 缓存 parser 和 pipeline
405        info.cached_parser = Some(std::sync::Arc::new(parser));
406        info.cached_pipeline = pipeline;
407
408        // 保存旧的 user_id（在修改之前）
409        let old_user_id = info.user_id.clone();
410
411        let user_id_to_set = user_id.clone().or(old_user_id.clone());
412
413        // 添加调试日志
414        if user_id_to_set.is_none() {
415            tracing::trace!(
416                connection_id = %connection_id,
417                incoming_user_id = ?user_id,
418                old_user_id = ?old_user_id,
419                "update_connection_negotiation_with_pipeline: user_id_to_set is None, user_id will not be set"
420            );
421        }
422
423        if let Some(user_id_val) = user_id_to_set {
424            // 如果之前有用户 ID 且与新 user_id 不同，先移除旧映射
425            if let Some(old_user_id) = old_user_id
426                && old_user_id != user_id_val
427            {
428                self.remove_user_connection(&old_user_id, connection_id)?;
429            }
430
431            // 更新用户 ID
432            info.user_id = Some(user_id_val.clone());
433
434            // 添加到新用户映射
435            self.insert_user_connection(user_id_val, connection_id)?;
436        }
437
438        Ok(())
439    }
440
441    /// 标记协商已确认（客户端收到 CONNECT_ACK 后发送确认）
442    pub fn mark_negotiation_confirmed(&self, connection_id: &str) -> Result<()> {
443        let mut shard = self
444            .connection_shard(connection_id)
445            .write()
446            .map_err(|_| FlareError::general_error("Failed to lock connection shard"))?;
447
448        let (_, _, info) = shard.get_mut(connection_id).ok_or_else(|| {
449            FlareError::protocol_error(format!("Connection {} not found", connection_id))
450        })?;
451
452        if !info.negotiation_completed {
453            return Err(FlareError::protocol_error(format!(
454                "Cannot confirm negotiation for connection {}: negotiation not completed",
455                connection_id
456            )));
457        }
458
459        info.negotiation_confirmed = true;
460        tracing::trace!(
461            "[ConnectionManager] 协商已确认: connection_id={}",
462            connection_id
463        );
464
465        Ok(())
466    }
467
468    /// 获取所有连接 ID
469    pub fn list_connections(&self) -> Vec<String> {
470        self.connection_shards
471            .iter()
472            .flat_map(|shard| {
473                shard
474                    .read()
475                    .ok()
476                    .map(|connections| connections.keys().cloned().collect::<Vec<_>>())
477                    .unwrap_or_default()
478            })
479            .collect()
480    }
481
482    /// 获取连接总数
483    pub fn connection_count(&self) -> usize {
484        self.connection_count.load(Ordering::Relaxed)
485    }
486
487    /// 获取当前绑定了连接的用户数
488    pub fn user_count(&self) -> usize {
489        self.user_count.load(Ordering::Relaxed)
490    }
491
492    /// 清理超时连接
493    ///
494    /// # 参数
495    /// - `timeout`: 超时时间
496    ///
497    /// # 返回
498    /// 被清理的连接 ID 列表
499    pub fn cleanup_timeout_connections(&self, timeout: Duration) -> Vec<String> {
500        let timeout_connections = self.timeout_connection_snapshots(timeout);
501        self.remove_connection_snapshots(
502            timeout_connections
503                .iter()
504                .map(|(connection_id, _, _)| connection_id.clone()),
505        )
506    }
507
508    /// 获取连接统计信息
509    pub fn stats(&self) -> TraitConnectionStats {
510        let total_connections = self.connection_count();
511        let total_users = self.user_count();
512
513        TraitConnectionStats {
514            total_connections,
515            total_users,
516        }
517    }
518
519    pub(super) fn frame_allowed_before_auth(frame: &crate::common::protocol::Frame) -> bool {
520        frame
521            .command
522            .as_ref()
523            .and_then(|cmd| {
524                if let Some(crate::common::protocol::flare::core::commands::command::Type::System(
525                    sys_cmd,
526                )) = &cmd.r#type
527                {
528                    Some(
529                        sys_cmd.r#type
530                            == crate::common::protocol::flare::core::commands::system_command::Type::ConnectAck
531                                as i32
532                            || sys_cmd.r#type
533                                == crate::common::protocol::flare::core::commands::system_command::Type::Ping
534                                    as i32
535                            || sys_cmd.r#type
536                                == crate::common::protocol::flare::core::commands::system_command::Type::Pong
537                                    as i32
538                            || sys_cmd.r#type
539                                == crate::common::protocol::flare::core::commands::system_command::Type::Error
540                                    as i32
541                            || sys_cmd.r#type
542                                == crate::common::protocol::flare::core::commands::system_command::Type::Close
543                                    as i32,
544                    )
545                } else {
546                    None
547                }
548            })
549            .unwrap_or(false)
550    }
551
552    pub(super) fn serialize_frame_for_connection(
553        connection_id: &str,
554        info: &ConnectionInfo,
555        frame: &crate::common::protocol::Frame,
556        parser: Option<&crate::common::MessageParser>,
557    ) -> Result<Vec<u8>> {
558        Self::ensure_frame_allowed_for_connection(connection_id, info, frame)?;
559
560        if let Some(parser) = parser {
561            return parser.serialize(frame);
562        }
563
564        if let Some(parser) = &info.cached_parser {
565            return parser.serialize(frame);
566        }
567
568        crate::common::MessageParser::new(
569            info.serialization_format,
570            info.compression.clone(),
571            info.encryption.clone(),
572        )
573        .serialize(frame)
574    }
575
576    pub(super) fn ensure_frame_allowed_for_connection(
577        connection_id: &str,
578        info: &ConnectionInfo,
579        frame: &crate::common::protocol::Frame,
580    ) -> Result<()> {
581        if info.authenticated || Self::frame_allowed_before_auth(frame) {
582            return Ok(());
583        }
584
585        Err(FlareError::authentication_failed(format!(
586            "连接 {} 未验证，无法发送消息",
587            connection_id
588        )))
589    }
590
591    pub(super) async fn send_to_connection_handle(
592        &self,
593        connection_id: &str,
594        connection: ConnectionWriteHandle,
595        data: &[u8],
596    ) -> Result<()> {
597        match connection.try_enqueue(data) {
598            Ok(()) => Ok(()),
599            Err(err) => {
600                connection.close_underlying_in_background();
601                let _ = ConnectionManager::remove_connection(self, connection_id);
602                Err(err)
603            }
604        }
605    }
606
607    pub(super) async fn send_frame_to_snapshot(
608        &self,
609        snapshot: ConnectionSnapshot,
610        frame: &crate::common::protocol::Frame,
611        parser: Option<&crate::common::MessageParser>,
612    ) -> Result<()> {
613        let connection_id = self
614            .send_frame_to_snapshot_without_active(snapshot, frame, parser)
615            .await?;
616        ConnectionManager::update_connection_active(self, &connection_id)?;
617        Ok(())
618    }
619
620    pub(super) async fn send_frame_to_snapshot_without_active(
621        &self,
622        snapshot: ConnectionSnapshot,
623        frame: &crate::common::protocol::Frame,
624        parser: Option<&crate::common::MessageParser>,
625    ) -> Result<String> {
626        let (connection_id, connection, info) = snapshot;
627        let data = Self::serialize_frame_for_connection(&connection_id, &info, frame, parser)?;
628
629        self.send_to_connection_handle(&connection_id, connection, &data)
630            .await?;
631        Ok(connection_id)
632    }
633
634    pub(super) async fn send_serialized_frame_to_auth_snapshot_without_active(
635        &self,
636        snapshot: ConnectionAuthSnapshot,
637        frame: &crate::common::protocol::Frame,
638        data: &[u8],
639    ) -> Result<String> {
640        let (connection_id, connection, authenticated) = snapshot;
641        if !authenticated && !Self::frame_allowed_before_auth(frame) {
642            return Err(FlareError::authentication_failed(format!(
643                "连接 {} 未验证，无法发送消息",
644                connection_id
645            )));
646        }
647
648        self.send_to_connection_handle(&connection_id, connection, data)
649            .await?;
650        Ok(connection_id)
651    }
652}