feat: add ws stats task id (#294)

openstream-fm · Mar 12, 2024 · ed68b8a · ed68b8a
2 parents f1fb384 + e985b78
commit ed68b8a
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 2 deletions.
diff --git a/defs/db/WsStatsConnection.ts b/defs/db/WsStatsConnection.ts
@@ -5,6 +5,7 @@ import type { DateTime } from "../DateTime.js";
 export type WsStatsConnection = {
   _id: string;
   st: string;
+  task_id: string;
   sd: string | null | undefined;
   dp: string;
   du: number | null | undefined;

diff --git a/rs/packages/api/src/ws_stats/routes/connection.rs b/rs/packages/api/src/ws_stats/routes/connection.rs
@@ -164,6 +164,7 @@ impl WsConnectionHandler {
         }
       };
 
+      let task_id = WsStatsConnection::random_task_id();
       let reconnections: u16;
       let connection_id: String;
       let created_at: DateTime;
@@ -176,6 +177,7 @@ impl WsConnectionHandler {
 
           let connection = WsStatsConnection {
             id: connection_id.clone(),
+            task_id: task_id.clone(),
             station_id: station_id.clone(),
             start_deployment_id: Some(deployment_id.clone()),
             current_deployment_id: deployment_id,
@@ -209,6 +211,7 @@ impl WsConnectionHandler {
 
           let update = doc! {
             "$set": {
+              WsStatsConnection::KEY_TASK_ID: &task_id,
               WsStatsConnection::KEY_CURRENT_DEPLOYMENT_ID: &deployment_id,
               WsStatsConnection::KEY_IS_OPEN: true,
               WsStatsConnection::KEY_CLOSED_AT: null,
@@ -351,6 +354,12 @@ impl WsConnectionHandler {
 
       let duration_ms = ((*DateTime::now() - *created_at).as_seconds_f64() * 1000.0).round();
 
+      // only update if the task id is the same (not taken from another connection)
+      let filter = doc! {
+        WsStatsConnection::KEY_ID: &connection_id,
+        WsStatsConnection::KEY_TASK_ID: &task_id,
+      };
+
       let update = doc! {
         "$set": {
           WsStatsConnection::KEY_IS_OPEN: false,
@@ -359,7 +368,9 @@ impl WsConnectionHandler {
         }
       };
 
-      let _ = WsStatsConnection::update_by_id(&connection_id, update).await;
+      let _ = WsStatsConnection::cl()
+        .update_one(filter, update, None)
+        .await;
 
       log::info!(
         target: "ws-stats",

diff --git a/rs/packages/db/src/lib.rs b/rs/packages/db/src/lib.rs
@@ -20,7 +20,6 @@ use schemars::JsonSchema;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use serde_util::DateTime;
 use std::borrow::Borrow;
-use std::collections::HashSet;
 use ts_rs::TS;
 
 pub mod error;

diff --git a/rs/packages/db/src/models/ws_stats_connection/mod.rs b/rs/packages/db/src/models/ws_stats_connection/mod.rs
@@ -25,6 +25,9 @@ pub struct WsStatsConnection {
   #[serde(rename = "st")]
   pub station_id: String,
 
+  #[serde(default)]
+  pub task_id: String,
+
   #[serde(rename = "sd")]
   pub start_deployment_id: Option<String>,
 
@@ -74,6 +77,13 @@ pub struct WsStatsConnection {
   pub abnormally_closed: bool,
 }
 
+impl WsStatsConnection {
+  /// create a random task id, this task id is always used with the id field, so its only unique within the same document
+  pub fn random_task_id() -> String {
+    uid::uid(6)
+  }
+}
+
 impl Model for WsStatsConnection {
   const CL_NAME: &'static str = "ws_stats_connection";
   const UID_LEN: usize = 12;