redis_objects/lib.rs
1//! An object oriented wrapper around certain redis objects.
2
3#![warn(missing_docs, non_ascii_idents, trivial_numeric_casts,
4 unused_crate_dependencies, noop_method_call, single_use_lifetimes, trivial_casts,
5 unused_lifetimes, nonstandard_style, variant_size_differences)]
6#![deny(keyword_idents)]
7// #![warn(clippy::missing_docs_in_private_items)]
8// #![allow(clippy::needless_return)]
9// #![allow(clippy::while_let_on_iterator, clippy::collapsible_else_if)]
10
11use std::sync::Arc;
12use std::time::Duration;
13
14use log::debug;
15use queue::MultiQueue;
16use quota::UserQuotaTracker;
17use redis::AsyncTypedCommands;
18use redis::IntoConnectionInfo;
19pub use redis::Msg;
20use serde::Serialize;
21use serde::de::DeserializeOwned;
22use tokio::sync::mpsc;
23use tracing::instrument;
24
25pub use self::queue::PriorityQueue;
26pub use self::queue::Queue;
27// pub use self::quota::QuotaGuard;
28pub use self::hashmap::Hashmap;
29pub use self::counters::{AutoExportingMetrics, AutoExportingMetricsBuilder, MetricMessage};
30pub use self::pubsub::{JsonListenerBuilder, ListenerBuilder, Publisher};
31pub use self::set::Set;
32
33pub mod queue;
34pub mod quota;
35pub mod hashmap;
36pub mod counters;
37pub mod pubsub;
38pub mod set;
39
40/// maximum exponent over 2 used in the retry loop
41pub const BACKOFF_MAXIMUM: f64 = 3.0;
42
43/// Handle for a pool of connections to a redis server.
44pub struct RedisObjects {
45 pool: bb8::Pool<redis::Client>,
46 client: redis::Client,
47 hostname: String,
48 pubsub_prefix: String,
49}
50
51impl std::fmt::Debug for RedisObjects {
52 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53 f.debug_struct("Redis").field("host", &self.hostname).finish()
54 }
55}
56
57impl RedisObjects {
58 /// Open given more limited connection info
59 pub fn open_host(host: &str, port: u16, db: i64) -> Result<Arc<Self>, ErrorTypes> {
60 let info = redis::ConnectionAddr::Tcp(host.to_string(), port).into_connection_info()?;
61 let settings = info.redis_settings().clone();
62 let info = info.set_redis_settings(settings.set_db(db));
63 Self::open(info)
64 }
65
66 /// Open a connection using the local tls configuration
67 pub fn open_host_native_tls(host: &str, port: u16, db: i64) -> Result<Arc<Self>, ErrorTypes> {
68 let info = redis::ConnectionAddr::TcpTls {
69 host: host.to_string(),
70 port,
71 insecure: false,
72 tls_params: None,
73 }.into_connection_info()?;
74 let settings = info.redis_settings().clone();
75 let info = info.set_redis_settings(settings.set_db(db));
76 Self::open(info)
77 }
78
79 /// Construct a builder to construct a redis connection
80 pub fn builder() -> Builder {
81 Builder::default()
82 }
83
84 /// Open a connection pool
85 pub fn open(config: redis::ConnectionInfo) -> Result<Arc<Self>, ErrorTypes> {
86 Self::_open(config, Default::default())
87 }
88
89 /// Open a connection pool
90 fn _open(config: redis::ConnectionInfo, pubsub_prefix: String) -> Result<Arc<Self>, ErrorTypes> {
91 debug!("Create redis connection pool.");
92
93 let hostname = config.addr().to_string();
94
95 let settings = config.tcp_settings().clone()
96 .set_keepalive(redis::io::tcp::socket2::TcpKeepalive::new().with_interval(Duration::from_secs(5)))
97 .set_nodelay(true);
98 let config = config.set_tcp_settings(settings);
99
100 let client = redis::Client::open(config)?;
101
102 // configuration for the pool manager itself
103 let pool = bb8::Pool::builder()
104 .max_size(1024)
105 .test_on_check_out(false)
106 .connection_timeout(Duration::from_mins(5))
107 .build_unchecked(client.clone());
108
109 // load redis configuration and create the pool
110 Ok(Arc::new(Self{
111 pool,
112 client,
113 hostname,
114 pubsub_prefix,
115 }))
116 }
117
118 /// Open a priority queue under the given key
119 pub fn priority_queue<T: Serialize + DeserializeOwned>(self: &Arc<Self>, name: String) -> PriorityQueue<T> {
120 PriorityQueue::new(name, self.clone())
121 }
122
123 /// Open a FIFO queue under the given key
124 pub fn queue<T: Serialize + DeserializeOwned>(self: &Arc<Self>, name: String, ttl: Option<Duration>) -> Queue<T> {
125 Queue::new(name, self.clone(), ttl)
126 }
127
128 /// an object that represents a set of queues with a common prefix
129 pub fn multiqueue<T: Serialize + DeserializeOwned>(self: &Arc<Self>, prefix: String) -> MultiQueue<T> {
130 MultiQueue::new(prefix, self.clone())
131 }
132
133 /// Open a hash map under the given key
134 pub fn hashmap<T: Serialize + DeserializeOwned>(self: &Arc<Self>, name: String, ttl: Option<Duration>) -> Hashmap<T> {
135 Hashmap::new(name, self.clone(), ttl)
136 }
137
138 /// Create a sink to publish messages to a named channel
139 pub fn publisher(self: &Arc<Self>, channel: String) -> Publisher {
140 Publisher::new(self.clone(), channel)
141 }
142
143 /// Write a message directly to the channel given
144 #[instrument]
145 pub async fn publish(&self, channel: &str, data: &[u8]) -> Result<usize, ErrorTypes> {
146 retry_call!(self.pool, publish, self.pubsub_prefix.clone() + channel, data)
147 }
148
149 /// Write a json message directly to the channel given
150 #[instrument(skip(value))]
151 pub async fn publish_json<T: Serialize>(&self, channel: &str, value: &T) -> Result<usize, ErrorTypes> {
152 self.publish(channel, &serde_json::to_vec(value)?).await
153 }
154
155 /// Start building a metrics exporter
156 pub fn auto_exporting_metrics<T: MetricMessage>(self: &Arc<Self>, name: String, counter_type: String) -> AutoExportingMetricsBuilder<T> {
157 AutoExportingMetricsBuilder::new(self.clone(), name, counter_type)
158 }
159
160 /// Start building a json listener that produces a stream of events
161 pub fn pubsub_json_listener<T: DeserializeOwned + Send + 'static>(self: &Arc<Self>) -> JsonListenerBuilder<T> {
162 JsonListenerBuilder::new(self.clone())
163 }
164
165 /// Build a json listener that produces a stream of events using the default configuration
166 #[instrument]
167 pub async fn subscribe_json<T: DeserializeOwned + Send + 'static>(self: &Arc<Self>, channel: String) -> mpsc::Receiver<Option<T>> {
168 self.pubsub_json_listener()
169 .subscribe(channel)
170 .listen().await
171 }
172
173 /// Start building a raw data listener that produces a stream of events
174 pub fn pubsub_listener(self: &Arc<Self>) -> ListenerBuilder {
175 ListenerBuilder::new(self.clone())
176 }
177
178 /// Build a raw data listener that produces a stream of events using the default configuration
179 #[instrument]
180 pub async fn subscribe(self: &Arc<Self>, channel: String) -> mpsc::Receiver<Option<Msg>> {
181 self.pubsub_listener()
182 .subscribe(channel)
183 .listen().await
184 }
185
186 /// Open an interface for tracking user quotas on redis
187 pub fn user_quota_tracker(self: &Arc<Self>, prefix: String) -> UserQuotaTracker {
188 UserQuotaTracker::new(self.clone(), prefix)
189 }
190
191 /// Open a set of values
192 pub fn set<T: Serialize + DeserializeOwned>(self: &Arc<Self>, name: String) -> Set<T> {
193 Set::new(name, self.clone(), None)
194 }
195
196 /// Open a set of values with an expiration policy
197 pub fn expiring_set<T: Serialize + DeserializeOwned>(self: &Arc<Self>, name: String, ttl: Option<Duration>) -> Set<T> {
198 let ttl = ttl.unwrap_or_else(|| Duration::from_secs(86400));
199 Set::new(name, self.clone(), Some(ttl))
200 }
201
202 /// Erase all data on the redis server
203 pub async fn wipe(&self) -> Result<(), ErrorTypes> {
204 let mut con = self.pool.get().await?;
205 let _: () = redis::cmd("FLUSHDB").arg("SYNC").query_async(&mut *con).await?;
206 Ok(())
207 }
208
209 /// List all keys on the redis server that satisfies the given pattern
210 #[instrument]
211 pub async fn keys(&self, pattern: &str) -> Result<Vec<String>, ErrorTypes> {
212 Ok(retry_call!(self.pool, keys, pattern)?)
213 }
214
215}
216
217/// Utility class to simplify setting up a redis object
218pub struct Builder {
219 host: String,
220 port: u16,
221 db: i64,
222 native_tls: bool,
223 pubsub_prefix: String
224}
225
226impl Default for Builder {
227 fn default() -> Self {
228 Self {
229 host: "localhost".to_string(),
230 port: 6379,
231 db: 0,
232 native_tls: true,
233 pubsub_prefix: Default::default()
234 }
235 }
236}
237
238impl Builder {
239 /// Set the host for the connection being built
240 pub fn host(mut self, host: String) -> Self {
241 self.host = host; self
242 }
243
244 /// Set the port for the connection being built
245 pub fn port(mut self, port: u16) -> Self {
246 self.port = port; self
247 }
248
249 /// Set the database index for the connection being built
250 pub fn db(mut self, db: i64) -> Self {
251 self.db = db; self
252 }
253
254 /// Set whether the connection being built should be configured
255 pub fn native_tls(mut self, native_tls: bool) -> Self {
256 self.native_tls = native_tls; self
257 }
258
259 /// Set a prefix that should be applied to all pubsub listening and publishing.
260 /// This is intended for namespacing pubsub operations for testing.
261 pub fn pubsub_prefix(mut self, pubsub_prefix: String) -> Self {
262 self.pubsub_prefix = pubsub_prefix; self
263 }
264
265 /// finalize the building process and create the redis object
266 pub fn build(self) -> Result<Arc<RedisObjects>, ErrorTypes> {
267 let config = if self.native_tls {
268 let info = redis::ConnectionAddr::TcpTls {
269 host: self.host,
270 port: self.port,
271 insecure: false,
272 tls_params: None,
273 }.into_connection_info()?;
274 let settings = info.redis_settings().clone();
275 info.set_redis_settings(settings.set_db(self.db))
276 } else {
277 let info = redis::ConnectionAddr::Tcp(self.host, self.port).into_connection_info()?;
278 let settings = info.redis_settings().clone();
279 info.set_redis_settings(settings.set_db(self.db))
280 };
281
282 RedisObjects::_open(config, self.pubsub_prefix)
283 }
284}
285
286
287/// Enumeration over all possible errors
288#[derive(Debug, thiserror::Error)]
289pub enum ErrorTypes {
290 /// Could not get a connection from the redis connection pool
291 #[error("Could not get a connection from the redis connection pool")]
292 Pool,
293 /// Redis runtime error
294 #[error("Redis runtime error: {0}")]
295 Redis(Box<redis::RedisError>),
296 /// Unexpected response from redis server
297 #[error("Unexpected response from redis server")]
298 UnknownRedisError,
299 /// Encoding or decoding issue
300 #[error("Encoding issue with message: {0}")]
301 Serde(serde_json::Error),
302}
303
304impl ErrorTypes {
305 /// Test if an error was created in serializing or deserializing data
306 pub fn is_serialize_error(&self) -> bool {
307 matches!(self, ErrorTypes::Serde(_))
308 }
309}
310
311impl From<bb8::RunError<redis::RedisError>> for ErrorTypes {
312 fn from(value: bb8::RunError<redis::RedisError>) -> Self {
313 match value {
314 bb8::RunError::User(err) => err.into(),
315 bb8::RunError::TimedOut => Self::Pool,
316 }
317 }
318}
319
320impl From<redis::RedisError> for ErrorTypes {
321 fn from(value: redis::RedisError) -> Self { Self::Redis(Box::new(value)) }
322}
323
324impl From<serde_json::Error> for ErrorTypes {
325 fn from(value: serde_json::Error) -> Self { Self::Serde(value) }
326}
327
328/// A convenience trait that lets you pass an i32 value or None for arguments
329pub trait Ii32: Into<Option<i32>> + Copy {}
330impl<T: Into<Option<i32>> + Copy> Ii32 for T {}
331
332/// A convenience trait that lets you pass an usize value or None for arguments
333pub trait Iusize: Into<Option<usize>> + Copy {}
334impl<T: Into<Option<usize>> + Copy> Iusize for T {}
335
336
337macro_rules! retry_call_timeout {
338 ($pool:expr, $method:ident, $timeout:expr, $($args:expr),+) => {
339 {
340 // track our backoff parameters
341 let mut exponent = -7.0;
342 loop {
343 // get a (fresh if needed) connection form the pool
344 let mut con = retry_call!(get_connection, $method, $pool, exponent);
345
346 // execute the method given with the argments specified
347 if $timeout == 0.0 {
348 con.set_response_timeout(Duration::from_mins(60 * 24));
349 } else {
350 con.set_response_timeout(Duration::from_secs_f64($timeout + 60.0));
351 }
352 retry_call!(handle_output, $method, con.$method($($args,)+ $timeout).await, exponent)
353 }
354 }
355 };
356
357 (method, $pool:expr, $obj:expr, $method:ident, $timeout:expr) => {
358 {
359 // track our backoff parameters
360 let mut exponent = -7.0;
361 loop {
362 // get a (fresh if needed) connection form the pool
363 let mut con = retry_call!(get_connection, $method, $pool, exponent);
364
365 // execute the method given with the argments specified
366 if $timeout == 0.0 {
367 con.set_response_timeout(Duration::from_mins(60 * 24));
368 } else {
369 con.set_response_timeout(Duration::from_secs_f64($timeout + 60.0));
370 }
371 retry_call!(handle_output, $method, $obj.$method(&mut *con).await, exponent)
372 }
373 }
374 };
375}
376
377/// A macro for retrying calls to redis when an IO error occurs
378macro_rules! retry_call {
379 (get_connection, $name:ident, $pool:expr, $exponent:ident) => {
380 match $pool.get().await {
381 Ok(connection) => connection,
382 Err(bb8::RunError::User(err)) => {
383 retry_call!(handle_error, $name, err, $exponent);
384 continue
385 },
386 Err(bb8::RunError::TimedOut) => {
387 log::warn!("Timed out getting redis connection, will retry...");
388 tokio::time::sleep(tokio::time::Duration::from_secs_f64(2f64.powf($exponent))).await;
389 $exponent = ($exponent + 1.0).min(crate::BACKOFF_MAXIMUM);
390 continue
391 }
392 }
393 };
394
395 (handle_error, $name:ident, $err:ident, $exponent:ident) => {
396 {
397 // If the error from redis is something not related to IO let the error propagate
398 if !$err.is_io_error() {
399 break Result::<_, ErrorTypes>::Err($err.into())
400 }
401
402 // For IO errors print a warning and sleep
403 log::warn!("No connection to Redis, reconnecting... [{}:{}][{}]", stringify!($name), std::line!(), $err);
404 tokio::time::sleep(tokio::time::Duration::from_secs_f64(2f64.powf($exponent))).await;
405 $exponent = ($exponent + 1.0).min(crate::BACKOFF_MAXIMUM);
406 }
407 };
408
409 (handle_output, $name:ident, $call:expr, $exponent:ident) => {
410 {
411 match $call {
412 Ok(val) => {
413 if $exponent > -7.0 {
414 log::info!("Reconnected to Redis!")
415 }
416 break Ok(val)
417 },
418 Err(err) => retry_call!(handle_error, $name, err, $exponent),
419 }
420 }
421 };
422
423 ($pool:expr, $method:ident, $($args:expr),+) => {
424 {
425 // track our backoff parameters
426 let mut exponent = -7.0;
427 loop {
428 // get a (fresh if needed) connection form the pool
429 let mut con = retry_call!(get_connection, $method, $pool, exponent);
430
431 // execute the method given with the argments specified
432 retry_call!(handle_output, $method, con.$method($($args),+).await, exponent)
433 }
434 }
435 };
436
437 (method, $pool:expr, $obj:expr, $method:ident) => {
438 {
439 // track our backoff parameters
440 let mut exponent = -7.0;
441 loop {
442 // get a (fresh if needed) connection form the pool
443 let mut con = retry_call!(get_connection, $method, $pool, exponent);
444
445 // execute the method given with the argments specified
446 retry_call!(handle_output, $method, $obj.$method(&mut *con).await, exponent)
447 }
448 }
449 };
450}
451
452pub (crate) use {retry_call, retry_call_timeout};
453
454#[cfg(test)]
455pub (crate) mod test {
456 use std::sync::Arc;
457
458 use redis::IntoConnectionInfo;
459
460 use crate::{ErrorTypes, PriorityQueue, Queue, RedisObjects};
461
462
463 pub (crate) async fn redis_connection() -> Arc<RedisObjects> {
464 RedisObjects::open(redis::ConnectionAddr::Tcp("localhost".to_string(), 6379).into_connection_info().unwrap()).unwrap()
465 }
466
467 fn init() {
468 let _ = env_logger::builder().is_test(true).try_init();
469 }
470
471 // simple function to help test that reconnect is working.
472 // Enable and run this then turn your redis server on and off.
473 // #[tokio::test]
474 // async fn reconnect() {
475 // init();
476 // let connection = redis_connection().await;
477 // let mut listener = connection.subscribe("abc123".to_string());
478 // // launch a thread that sends messages every second forever
479 // tokio::spawn(async move {
480 // loop {
481 // tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
482 // connection.publish("abc123", b"100").await.unwrap();
483 // }
484 // });
485
486 // while let Some(msg) = listener.recv().await {
487 // println!("{msg:?}");
488 // }
489 // }
490
491 #[tokio::test]
492 async fn test_sets() {
493 init();
494 let redis = redis_connection().await;
495 let s = redis.set::<String>("test-set".to_owned());
496
497 s.delete().await.unwrap();
498
499 let values = &["a", "b", "1", "2"];
500 let owned = values.iter().map(|v|v.to_string()).collect::<Vec<_>>();
501 assert_eq!(s.add_batch(&owned).await.unwrap(), 4);
502 assert_eq!(s.length().await.unwrap(), 4);
503 let members = s.members().await.unwrap();
504 assert_eq!(members.len(), 4);
505 for x in members {
506 assert!(owned.contains(&x));
507 }
508 assert!(owned.contains(&s.random().await.unwrap().unwrap()));
509 assert!(s.exist(&owned[2]).await.unwrap());
510 s.remove(&owned[2]).await.unwrap();
511 assert!(!s.exist(&owned[2]).await.unwrap());
512 let pop_val = s.pop().await.unwrap().unwrap();
513 assert!(owned.contains(&pop_val));
514 assert!(!s.exist(&pop_val).await.unwrap());
515 assert_eq!(s.length().await.unwrap(), 2);
516
517 assert!(s.limited_add(&"dog".to_owned(), 3).await.unwrap());
518 assert!(!s.limited_add(&"cat".to_owned(), 3).await.unwrap());
519 assert!(s.exist(&"dog".to_owned()).await.unwrap());
520 assert!(!s.exist(&"cat".to_owned()).await.unwrap());
521 assert_eq!(s.length().await.unwrap(), 3);
522
523 for pop_val in s.pop_all().await.unwrap() {
524 assert!(values.contains(&pop_val.as_str()) || ["cat", "dog"].contains(&pop_val.as_str()));
525 }
526 assert!(s.pop().await.unwrap().is_none());
527 assert_eq!(s.length().await.unwrap(), 0);
528 }
529
530
531 // def test_expiring_sets(redis_connection):
532 // if redis_connection:
533 // from assemblyline.remote.datatypes.set import ExpiringSet
534 // with ExpiringSet('test-expiring-set', ttl=1) as es:
535 // es.delete()
536
537 // values = ['a', 'b', 1, 2]
538 // assert es.add(*values) == 4
539 // assert es.length() == 4
540 // assert es.exist(values[2])
541 // for x in es.members():
542 // assert x in values
543 // time.sleep(1.1)
544 // assert es.length() == 0
545 // assert not es.exist(values[2])
546
547
548 // # noinspection PyShadowingNames
549 // def test_lock(redis_connection):
550 // if redis_connection:
551 // from assemblyline.remote.datatypes.lock import Lock
552
553 // def locked_execution(next_thread=None):
554 // with Lock('test', 10):
555 // if next_thread:
556 // next_thread.start()
557 // time.sleep(2)
558
559 // t2 = Thread(target=locked_execution)
560 // t1 = Thread(target=locked_execution, args=(t2,))
561 // t1.start()
562
563 // time.sleep(1)
564 // assert t1.is_alive()
565 // assert t2.is_alive()
566 // time.sleep(2)
567 // assert not t1.is_alive()
568 // assert t2.is_alive()
569 // time.sleep(2)
570 // assert not t1.is_alive()
571 // assert not t2.is_alive()
572
573 #[tokio::test]
574 async fn priority_queue() -> Result<(), ErrorTypes> {
575 let redis = redis_connection().await;
576 let pq = redis.priority_queue("test-priority-queue".to_string());
577 pq.delete().await?;
578
579 for x in 0..10 {
580 pq.push(100.0, &x.to_string()).await?;
581 }
582
583 let a_key = pq.push(101.0, &"a".to_string()).await?;
584 let z_key = pq.push(99.0, &"z".to_string()).await?;
585 assert_eq!(pq.rank(&a_key).await?.unwrap(), 0);
586 assert_eq!(pq.rank(&z_key).await?.unwrap(), pq.length().await? - 1);
587 assert!(pq.rank("onethuosentuh").await?.is_none());
588
589 assert_eq!(pq.pop(1).await?, ["a"]);
590 assert_eq!(pq.unpush(1).await?, ["z"]);
591 assert_eq!(pq.count(100.0, 100.0).await?, 10);
592 assert_eq!(pq.pop(1).await?, ["0"]);
593 assert_eq!(pq.unpush(1).await?, ["9"]);
594 assert_eq!(pq.length().await?, 8);
595 assert_eq!(pq.pop(4).await?, ["1", "2", "3", "4"]);
596 assert_eq!(pq.unpush(3).await?, ["8", "7", "6"]);
597 assert_eq!(pq.length().await?, 1);
598 // Should be [(100, 5)] at this point
599
600 // for x in 0..5 {
601 for x in 0..5 {
602 pq.push(100.0 + x as f64, &x.to_string()).await?;
603 }
604
605 assert_eq!(pq.length().await?, 6);
606 assert!(pq.dequeue_range(Some(106), None, None, None).await?.is_empty());
607 assert_eq!(pq.length().await?, 6);
608 // 3 and 4 are both options, 4 has higher score
609 assert_eq!(pq.dequeue_range(Some(103), None, None, None).await?, vec!["4"]);
610 // 2 and 3 are both options, 3 has higher score, skip it
611 assert_eq!(pq.dequeue_range(Some(102), None, Some(1), None).await?, vec!["2"]);
612 // Take some off the other end
613 assert_eq!(pq.dequeue_range(None, Some(100), None, Some(10)).await?, vec!["5", "0"]);
614 assert_eq!(pq.length().await?, 2);
615
616 let other = redis.priority_queue("second-priority-queue".to_string());
617 other.delete().await?;
618 other.push(100.0, &"a".to_string()).await?;
619 assert_eq!(PriorityQueue::all_length(&[&other, &pq]).await?, [1, 2]);
620 assert!(PriorityQueue::select(&[&other, &pq], None).await?.is_some());
621 assert!(PriorityQueue::select(&[&other, &pq], None).await?.is_some());
622 assert!(PriorityQueue::select(&[&other, &pq], None).await?.is_some());
623 assert_eq!(PriorityQueue::all_length(&[&other, &pq]).await?, [0, 0]);
624
625 pq.push(50.0, &"first".to_string()).await?;
626 pq.push(-50.0, &"second".to_string()).await?;
627
628 assert_eq!(pq.dequeue_range(Some(0), Some(100), None, None).await?, ["first"]);
629 assert_eq!(pq.dequeue_range(Some(-100), Some(0), None, None).await?, ["second"]);
630 Ok(())
631 }
632
633
634 // # noinspection PyShadowingNames,PyUnusedLocal
635 // def test_unique_priority_queue(redis_connection):
636 // from assemblyline.remote.datatypes.queues.priority import UniquePriorityQueue
637 // with UniquePriorityQueue('test-priority-queue') as pq:
638 // pq.delete()
639
640 // for x in range(10):
641 // pq.push(100, x)
642 // assert pq.length() == 10
643
644 // # Values should be unique, this should have no effect on the length
645 // for x in range(10):
646 // pq.push(100, x)
647 // assert pq.length() == 10
648
649 // pq.push(101, 'a')
650 // pq.push(99, 'z')
651
652 // assert pq.pop() == 'a'
653 // assert pq.unpush() == 'z'
654 // assert pq.count(100, 100) == 10
655 // assert pq.pop() == 0
656 // assert pq.unpush() == 9
657 // assert pq.length() == 8
658 // assert pq.pop(4) == [1, 2, 3, 4]
659 // assert pq.unpush(3) == [8, 7, 6]
660 // assert pq.length() == 1 # Should be [<100, 5>] at this point
661
662 // for x in range(5):
663 // pq.push(100 + x, x)
664
665 // assert pq.length() == 6
666 // assert pq.dequeue_range(lower_limit=106) == []
667 // assert pq.length() == 6
668 // assert pq.dequeue_range(lower_limit=103) == [4] # 3 and 4 are both options, 4 has higher score
669 // assert pq.dequeue_range(lower_limit=102, skip=1) == [2] # 2 and 3 are both options, 3 has higher score, skip it
670 // assert sorted(pq.dequeue_range(upper_limit=100, num=10)) == [0, 5] # Take some off the other end
671 // assert pq.length() == 2
672 // pq.pop(2)
673
674 // pq.push(50, 'first')
675 // pq.push(-50, 'second')
676
677 // assert pq.dequeue_range(0, 100) == ['first']
678 // assert pq.dequeue_range(-100, 0) == ['second']
679
680 #[tokio::test]
681 async fn named_queue() {
682 let redis = redis_connection().await;
683
684 let nq = redis.queue("test-named-queue".to_owned(), None);
685 nq.delete().await.unwrap();
686
687 assert!(nq.pop().await.unwrap().is_none());
688 assert!(nq.pop_batch(100).await.unwrap().is_empty());
689
690 for x in 0..5 {
691 nq.push(&x).await.unwrap();
692 }
693
694 assert_eq!(nq.content().await.unwrap(), [0, 1, 2, 3, 4]);
695
696 assert_eq!(nq.pop_batch(100).await.unwrap(), [0, 1, 2, 3, 4]);
697
698 for x in 0..5 {
699 nq.push(&x).await.unwrap();
700 }
701
702 assert_eq!(nq.length().await.unwrap(), 5);
703 nq.push_batch(&(0..5).collect::<Vec<i32>>()).await.unwrap();
704 assert_eq!(nq.length().await.unwrap(), 10);
705
706 assert_eq!(nq.peek_next().await.unwrap(), nq.pop().await.unwrap());
707 assert_eq!(nq.peek_next().await.unwrap(), Some(1));
708 let v = nq.pop().await.unwrap().unwrap();
709 assert_eq!(v, 1);
710 assert_eq!(nq.peek_next().await.unwrap().unwrap(), 2);
711 nq.unpop(&v).await.unwrap();
712 assert_eq!(nq.peek_next().await.unwrap().unwrap(), 1);
713
714 assert_eq!(Queue::select(&[&nq], None).await.unwrap().unwrap(), ("test-named-queue".to_owned(), 1));
715
716 let nq1 = redis.queue("test-named-queue-1".to_owned(), None);
717 nq1.delete().await.unwrap();
718 let nq2 = redis.queue("test-named-queue-2".to_owned(), None);
719 nq2.delete().await.unwrap();
720
721 nq1.push(&1).await.unwrap();
722 nq2.push(&2).await.unwrap();
723
724 assert_eq!(Queue::select(&[&nq1, &nq2], None).await.unwrap().unwrap(), ("test-named-queue-1".to_owned(), 1));
725 assert_eq!(Queue::select(&[&nq1, &nq2], None).await.unwrap().unwrap(), ("test-named-queue-2".to_owned(), 2));
726 }
727
728 // # noinspection PyShadowingNames
729 // def test_multi_queue(redis_connection):
730 // if redis_connection:
731 // from assemblyline.remote.datatypes.queues.multi import MultiQueue
732 // mq = MultiQueue()
733 // mq.delete('test-multi-q1')
734 // mq.delete('test-multi-q2')
735
736 // for x in range(5):
737 // mq.push('test-multi-q1', x+1)
738 // mq.push('test-multi-q2', x+6)
739
740 // assert mq.length('test-multi-q1') == 5
741 // assert mq.length('test-multi-q2') == 5
742
743 // assert mq.pop('test-multi-q1') == 1
744 // assert mq.pop('test-multi-q2') == 6
745
746 // assert mq.length('test-multi-q1') == 4
747 // assert mq.length('test-multi-q2') == 4
748
749 // mq.delete('test-multi-q1')
750 // mq.delete('test-multi-q2')
751
752 // assert mq.length('test-multi-q1') == 0
753 // assert mq.length('test-multi-q2') == 0
754
755
756 // # noinspection PyShadowingNames
757 // def test_comms_queue(redis_connection):
758 // if redis_connection:
759 // from assemblyline.remote.datatypes.queues.comms import CommsQueue
760
761 // def publish_messages(message_list):
762 // time.sleep(0.1)
763 // with CommsQueue('test-comms-queue') as cq_p:
764 // for message in message_list:
765 // cq_p.publish(message)
766
767 // msg_list = ["bob", 1, {"bob": 1}, [1, 2, 3], None, "Nice!", "stop"]
768 // t = Thread(target=publish_messages, args=(msg_list,))
769 // t.start()
770
771 // with CommsQueue('test-comms-queue') as cq:
772 // x = 0
773 // for msg in cq.listen():
774 // if msg == "stop":
775 // break
776
777 // assert msg == msg_list[x]
778
779 // x += 1
780
781 // t.join()
782 // assert not t.is_alive()
783
784
785 // # noinspection PyShadowingNames
786 // def test_user_quota_tracker(redis_connection):
787 // if redis_connection:
788 // from assemblyline.remote.datatypes.user_quota_tracker import UserQuotaTracker
789
790 // max_quota = 3
791 // timeout = 2
792 // name = get_random_id()
793 // uqt = UserQuotaTracker('test-quota', timeout=timeout)
794
795 // # First 0 to max_quota items should succeed
796 // for _ in range(max_quota):
797 // assert uqt.begin(name, max_quota) is True
798
799 // # All other items should fail until items timeout
800 // for _ in range(max_quota):
801 // assert uqt.begin(name, max_quota) is False
802
803 // # if you remove and item only one should be able to go in
804 // uqt.end(name)
805 // assert uqt.begin(name, max_quota) is True
806 // assert uqt.begin(name, max_quota) is False
807
808 // # if you wait the timeout, all items can go in
809 // time.sleep(timeout+1)
810 // for _ in range(max_quota):
811 // assert uqt.begin(name, max_quota) is True
812
813
814// def test_exact_event(redis_connection: Redis[Any]):
815// calls: list[dict[str, Any]] = []
816
817// def _track_call(data: Optional[dict[str, Any]]):
818// if data is not None:
819// calls.append(data)
820
821// watcher = EventWatcher(redis_connection)
822// try:
823// watcher.register('changes.test', _track_call)
824// watcher.start()
825// sender = EventSender('changes.', redis_connection)
826// start = time.time()
827
828// while len(calls) < 5:
829// sender.send('test', {'payload': 100})
830
831// if time.time() - start > 10:
832// pytest.fail()
833// assert len(calls) >= 5
834
835// for row in calls:
836// assert row == {'payload': 100}
837
838// finally:
839// watcher.stop()
840
841
842// def test_serialized_event(redis_connection: Redis[Any]):
843// import threading
844// started = threading.Event()
845
846// class Event(enum.IntEnum):
847// ADD = 0
848// REM = 1
849
850// @dataclass
851// class Message:
852// name: str
853// event: Event
854
855// def _serialize(message: Message):
856// return json.dumps(asdict(message))
857
858// def _deserialize(data: str) -> Message:
859// return Message(**json.loads(data))
860
861// calls: list[Message] = []
862
863// def _track_call(data: Optional[Message]):
864// if data is not None:
865// calls.append(data)
866// else:
867// started.set()
868
869// watcher = EventWatcher[Message](redis_connection, deserializer=_deserialize)
870// try:
871// watcher.register('changes.test', _track_call)
872// watcher.skip_first_refresh = False
873// watcher.start()
874// assert started.wait(timeout=5)
875// sender = EventSender[Message]('changes.', redis_connection, serializer=_serialize)
876// start = time.time()
877
878// while len(calls) < 5:
879// sender.send('test', Message(name='test', event=Event.ADD))
880
881// if time.time() - start > 10:
882// pytest.fail()
883// assert len(calls) >= 5
884
885// expected = Message(name='test', event=Event.ADD)
886// for row in calls:
887// assert row == expected
888
889// finally:
890// watcher.stop()
891
892
893// def test_pattern_event(redis_connection: Redis[Any]):
894// calls: list[dict[str, Any]] = []
895
896// def _track_call(data: Optional[dict[str, Any]]):
897// if data is not None:
898// calls.append(data)
899
900// watcher = EventWatcher(redis_connection)
901// try:
902// watcher.register('changes.*', _track_call)
903// watcher.start()
904// sender = EventSender('changes.', redis_connection)
905// start = time.time()
906
907// while len(calls) < 5:
908// sender.send(uuid.uuid4().hex, {'payload': 100})
909
910// if time.time() - start > 10:
911// pytest.fail()
912// assert len(calls) >= 5
913
914// for row in calls:
915// assert row == {'payload': 100}
916
917// finally:
918// watcher.stop()
919
920
921}