Count the number of subscriptions through containers

rmqtt · Sep 23, 2023 · 5854651 · 5854651
1 parent 6c0a46a
commit 5854651
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 0 deletions.
diff --git a/rmqtt-plugins/rmqtt-cluster-broadcast/src/shared.rs b/rmqtt-plugins/rmqtt-cluster-broadcast/src/shared.rs
@@ -512,6 +512,11 @@ impl Shared for &'static ClusterShared {
         replys
     }
 
+    #[inline]
+    async fn subscriptions_count(&self) -> usize {
+        self.inner.subscriptions_count().await
+    }
+
     #[inline]
     fn get_grpc_clients(&self) -> GrpcClients {
         self.grpc_clients.clone()

diff --git a/rmqtt-plugins/rmqtt-cluster-raft/src/shared.rs b/rmqtt-plugins/rmqtt-cluster-raft/src/shared.rs
@@ -450,6 +450,10 @@ impl Shared for &'static ClusterShared {
         self.inner.query_subscriptions(q).await
     }
 
+    #[inline]
+    async fn subscriptions_count(&self) -> usize {
+        self.inner.subscriptions_count().await
+    }
     #[inline]
     fn get_grpc_clients(&self) -> GrpcClients {
         self.grpc_clients.clone()

diff --git a/rmqtt/src/broker/default.rs b/rmqtt/src/broker/default.rs
@@ -554,6 +554,16 @@ impl Shared for &'static DefaultShared {
     async fn query_subscriptions(&self, q: SubsSearchParams) -> Vec<SubsSearchResult> {
         self._query_subscriptions(&q).await
     }
+
+    #[inline]
+    async fn subscriptions_count(&self) -> usize {
+        futures::future::join_all(
+            self.peers.iter().map(|entry| async move { entry.s.subscriptions.len().await }),
+        )
+        .await
+        .iter()
+        .sum()
+    }
 }
 
 pub struct DefaultIter<'a> {

diff --git a/rmqtt/src/broker/mod.rs b/rmqtt/src/broker/mod.rs
@@ -101,6 +101,8 @@ pub trait Shared: Sync + Send {
     ///
     async fn query_subscriptions(&self, q: SubsSearchParams) -> Vec<SubsSearchResult>;
 
+    async fn subscriptions_count(&self) -> usize;
+
     ///This node is not included
     #[inline]
     fn get_grpc_clients(&self) -> GrpcClients {