zhenyi

Fork 0

Clone

perf(socket): shard hub registries to reduce lock contention

ZhenYi committed 3 days ago

parent ab0cb4d commit 62f5365 Browse files

File tree 2

2

core.rs Modified
registry.rs Modified
No files match.

lib/infra/socket/src/hub/core.rs +40 −9 View

@@ -1,5 +1,7 @@

11use std::collections::HashMap;

2use std::hash::{Hash, Hasher};

23use std::sync::Arc;

4use std::sync::atomic::{AtomicUsize, Ordering};

35use std::time::Duration;

46

57use serde_json::Value;

@@ -12,6 +14,7 @@

1214use super::super::{EmitPacket, SocketEvent, SocketOptions, SocketPayload};

1315

1416const EVENT_CAPACITY: usize = 1024;

17const SHARDS: usize = 64;

1518

1619#[derive(Clone)]

1720pub struct SocketHub {

@@ -22,10 +25,11 @@

2225 pub(super) options: SocketOptions,

2326 pub(super) backend: Arc<dyn SocketBackend>,

2427 pub(super) events: mpsc::Sender<SocketEvent>,

25 pub(super) sessions: Mutex<HashMap<String, mpsc::Sender<Input>>>,

26 pub(super) sockets: Mutex<HashMap<i64, UserSockets>>,

28 pub(super) sessions: Vec<Mutex<HashMap<String, mpsc::Sender<Input>>>>,

29 pub(super) sockets: Vec<Mutex<HashMap<i64, UserSockets>>>,

2730 pub(crate) shutdown: watch::Sender<bool>,

28 pub(super) transports: Mutex<JoinSet<()>>,

31 pub(super) transports: Vec<Mutex<JoinSet<()>>>,

32 pub(super) transport_next: AtomicUsize,

2933}

3034

3135#[derive(Default)]

@@ -34,6 +38,21 @@

3438 pub(super) changed: Arc<Notify>,

3539}

3640

41impl Inner {

42 pub(super) fn sessions_shard(&self, sid: &str) -> &Mutex<HashMap<String, mpsc::Sender<Input>>> {

43 &self.sessions[sid_shard(sid)]

44 }

45

46 pub(super) fn sockets_shard(&self, user_id: i64) -> &Mutex<HashMap<i64, UserSockets>> {

47 &self.sockets[(user_id as u64 as usize) & (SHARDS - 1)]

48 }

49

50 pub(super) fn transport_shard(&self) -> &Mutex<JoinSet<()>> {

51 let index = self.transport_next.fetch_add(1, Ordering::Relaxed) & (SHARDS - 1);

52 &self.transports[index]

53 }

54}

55

3756impl SocketHub {

3857 pub fn new(

3958 options: SocketOptions,

@@ -46,10 +65,11 @@

4665 options,

4766 backend,

4867 events,

49 sessions: Mutex::new(HashMap::new()),

50 sockets: Mutex::new(HashMap::new()),

68 sessions: shards(),

69 sockets: shards(),

5170 shutdown: watch::channel(false).0,

52 transports: Mutex::new(JoinSet::new()),

71 transports: shards(),

72 transport_next: AtomicUsize::new(0),

5373 }),

5474 };

5575 Ok((hub, receiver))

@@ -140,14 +160,15 @@

140160 }

141161

142162 pub async fn wake_user(&self, user_id: i64) {

143 if let Some(user) = self.inner.sockets.lock().await.get(&user_id) {

163 let sockets = self.inner.sockets_shard(user_id).lock().await;

164 if let Some(user) = sockets.get(&user_id) {

144165 user.changed.notify_waiters();

145166 }

146167 }

147168

148169 pub async fn notification(&self, user_id: i64) -> Arc<Notify> {

149170 self.inner

150 .sockets

171 .sockets_shard(user_id)

151172 .lock()

152173 .await

153174 .entry(user_id)

@@ -171,7 +192,7 @@

171192

172193 /// 关闭某用户的全部在线会话（踢下线 / 权限撤销）。

173194 pub async fn close_user(&self, user_id: i64) -> usize {

174 let sockets = self.inner.sockets.lock().await;

195 let sockets = self.inner.sockets_shard(user_id).lock().await;

175196 let mut closed = 0;

176197 if let Some(user) = sockets.get(&user_id) {

177198 for entries in user.namespaces.values() {

@@ -211,3 +232,13 @@

211232 format!("/{namespace}")

212233 }

213234}

235

236fn shards<T: Default>() -> Vec<Mutex<T>> {

237 (0..SHARDS).map(|_| Mutex::new(T::default())).collect()

238}

239

240fn sid_shard(sid: &str) -> usize {

241 let mut hasher = std::collections::hash_map::DefaultHasher::new();

242 sid.hash(&mut hasher);

243 (hasher.finish() as usize) & (SHARDS - 1)

244}

lib/infra/socket/src/hub/registry.rs +25 −11 View

@@ -13,7 +13,7 @@

1313 &self,

1414 future: impl std::future::Future<Output = ()> + Send + 'static,

1515 ) {

16 let mut transports = self.inner.transports.lock().await;

16 let mut transports = self.inner.transport_shard().lock().await;

1717 while let Some(result) = transports.try_join_next() {

1818 if let Err(error) = result {

1919 eprintln!("kenevent: WebSocket task failed: {error}");

@@ -23,15 +23,24 @@

2323 }

2424

2525 pub(crate) async fn finish_transports(&self) {

26 let mut transports = self.inner.transports.lock().await;

26 let mut transports = Vec::with_capacity(self.inner.transports.len());

27 for shard in &self.inner.transports {

28 transports.push(shard.lock().await);

29 }

2730 if tokio::time::timeout(std::time::Duration::from_secs(5), async {

28 while transports.join_next().await.is_some() {}

31 for set in &mut transports {

32 while set.join_next().await.is_some() {}

33 }

2934 })

3035 .await

3136 .is_err()

3237 {

33 transports.abort_all();

34 while transports.join_next().await.is_some() {}

38 for set in &mut transports {

39 set.abort_all();

40 }

41 for set in &mut transports {

42 while set.join_next().await.is_some() {}

43 }

3544 }

3645 }

3746

@@ -54,7 +63,7 @@

5463 ) -> mpsc::Sender<Input> {

5564 let (input, receiver) = mpsc::channel(INPUT_CAPACITY);

5665 self.inner

57 .sessions

66 .sessions_shard(&sid)

5867 .lock()

5968 .await

6069 .insert(sid.clone(), input.clone());

@@ -79,11 +88,16 @@

7988 }

8089

8190 pub(crate) async fn remove_session(&self, sid: &str) {

82 self.inner.sessions.lock().await.remove(sid);

91 self.inner.sessions_shard(sid).lock().await.remove(sid);

8392 }

8493

8594 pub(crate) async fn input(&self, sid: &str) -> Option<mpsc::Sender<Input>> {

86 self.inner.sessions.lock().await.get(sid).cloned()

95 self.inner

96 .sessions_shard(sid)

97 .lock()

98 .await

99 .get(sid)

100 .cloned()

87101 }

88102

89103 pub(crate) async fn register(

@@ -94,7 +108,7 @@

94108 input: mpsc::Sender<Input>,

95109 ) {

96110 self.inner

97 .sockets

111 .sockets_shard(user_id)

98112 .lock()

99113 .await

100114 .entry(user_id)

@@ -106,7 +120,7 @@

106120 }

107121

108122 pub(crate) async fn unregister(&self, user_id: i64, namespace: &str, session_id: &str) {

109 let mut sockets = self.inner.sockets.lock().await;

123 let mut sockets = self.inner.sockets_shard(user_id).lock().await;

110124 if let Some(user) = sockets.get_mut(&user_id) {

111125 let key = namespace_key(namespace);

112126 if let Some(entries) = user.namespaces.get_mut(&key) {

@@ -123,7 +137,7 @@

123137

124138 pub(super) async fn inputs(&self, user_id: i64, namespace: &str) -> Vec<mpsc::Sender<Input>> {

125139 self.inner

126 .sockets

140 .sockets_shard(user_id)

127141 .lock()

128142 .await

129143 .get(&user_id)