7
7
use std:: { collections:: HashMap , net:: IpAddr } ;
8
8
9
9
use chrono:: { DateTime , Utc } ;
10
- use mas_storage:: { RepositoryAccess , user:: BrowserSessionRepository } ;
10
+ use mas_storage:: { RepositoryAccess , RepositoryError , user:: BrowserSessionRepository } ;
11
11
use opentelemetry:: {
12
12
Key , KeyValue ,
13
- metrics:: { Counter , Histogram } ,
13
+ metrics:: { Counter , Gauge , Histogram } ,
14
14
} ;
15
15
use sqlx:: PgPool ;
16
16
use tokio_util:: sync:: CancellationToken ;
@@ -25,8 +25,8 @@ use crate::{
25
25
/// database automatically.
26
26
///
27
27
/// The [`ActivityRecord`] structure plus the key in the [`HashMap`] takes less
28
- /// than 100 bytes, so this should allocate around a megabyte of memory.
29
- static MAX_PENDING_RECORDS : usize = 10_000 ;
28
+ /// than 100 bytes, so this should allocate around 100kB of memory.
29
+ static MAX_PENDING_RECORDS : usize = 1000 ;
30
30
31
31
const TYPE : Key = Key :: from_static_str ( "type" ) ;
32
32
const SESSION_KIND : Key = Key :: from_static_str ( "session_kind" ) ;
@@ -45,6 +45,7 @@ struct ActivityRecord {
45
45
pub struct Worker {
46
46
pool : PgPool ,
47
47
pending_records : HashMap < ( SessionKind , Ulid ) , ActivityRecord > ,
48
+ pending_records_gauge : Gauge < u64 > ,
48
49
message_counter : Counter < u64 > ,
49
50
flush_time_histogram : Histogram < u64 > ,
50
51
}
@@ -80,9 +81,17 @@ impl Worker {
80
81
. with_unit ( "ms" )
81
82
. build ( ) ;
82
83
84
+ let pending_records_gauge = METER
85
+ . u64_gauge ( "mas.activity_tracker.pending_records" )
86
+ . with_description ( "The number of pending activity records" )
87
+ . with_unit ( "{records}" )
88
+ . build ( ) ;
89
+ pending_records_gauge. record ( 0 , & [ ] ) ;
90
+
83
91
Self {
84
92
pool,
85
93
pending_records : HashMap :: with_capacity ( MAX_PENDING_RECORDS ) ,
94
+ pending_records_gauge,
86
95
message_counter,
87
96
flush_time_histogram,
88
97
}
@@ -165,6 +174,10 @@ impl Worker {
165
174
let _ = tx. send ( ( ) ) ;
166
175
}
167
176
}
177
+
178
+ // Update the gauge
179
+ self . pending_records_gauge
180
+ . record ( self . pending_records . len ( ) as u64 , & [ ] ) ;
168
181
}
169
182
170
183
// Flush one last time
@@ -193,18 +206,22 @@ impl Worker {
193
206
Err ( e) => {
194
207
self . flush_time_histogram
195
208
. record ( duration_ms, & [ KeyValue :: new ( RESULT , "failure" ) ] ) ;
196
- tracing:: error!( "Failed to flush activity tracker: {}" , e) ;
209
+ tracing:: error!(
210
+ error = & e as & dyn std:: error:: Error ,
211
+ "Failed to flush activity tracker"
212
+ ) ;
197
213
}
198
214
}
199
215
}
200
216
201
217
/// Fallible part of [`Self::flush`].
202
218
#[ tracing:: instrument( name = "activity_tracker.flush" , skip( self ) ) ]
203
- async fn try_flush ( & mut self ) -> Result < ( ) , anyhow :: Error > {
219
+ async fn try_flush ( & mut self ) -> Result < ( ) , RepositoryError > {
204
220
let pending_records = & self . pending_records ;
205
221
206
222
let mut repo = mas_storage_pg:: PgRepository :: from_pool ( & self . pool )
207
- . await ?
223
+ . await
224
+ . map_err ( RepositoryError :: from_error) ?
208
225
. boxed ( ) ;
209
226
210
227
let mut browser_sessions = Vec :: new ( ) ;
0 commit comments