1
1
use diesel:: { ExpressionMethods , QueryDsl } ;
2
2
use diesel_async:: RunQueryDsl ;
3
+ use kitsune_config:: instance:: StatisticsMode ;
3
4
use kitsune_db:: {
4
5
schema:: { accounts, posts, users} ,
5
6
with_connection, PgPool ,
@@ -8,7 +9,7 @@ use kitsune_derive::kitsune_service;
8
9
use kitsune_error:: { Error , Result } ;
9
10
use rand:: seq:: IteratorRandom ;
10
11
use smol_str:: SmolStr ;
11
- use std:: ops:: RangeInclusive ;
12
+ use std:: { future :: Future , ops:: RangeInclusive } ;
12
13
13
14
const STATISTICS_RANGE : RangeInclusive < u64 > = 24 ..=1312_1312 ;
14
15
@@ -25,8 +26,20 @@ pub struct InstanceService {
25
26
#[ builder( setter( into) ) ]
26
27
description : SmolStr ,
27
28
character_limit : usize ,
28
- randomize_statistics : bool ,
29
29
registrations_open : bool ,
30
+ statistics_mode : StatisticsMode ,
31
+ }
32
+
33
+ #[ inline]
34
+ async fn with_statistics_mode < F , E > ( mode : StatisticsMode , fut : F ) -> Result < u64 , E >
35
+ where
36
+ F : Future < Output = Result < u64 , E > > ,
37
+ {
38
+ match mode {
39
+ StatisticsMode :: Random => Ok ( random_statistic ( ) ) ,
40
+ StatisticsMode :: Regular => fut. await ,
41
+ StatisticsMode :: Zero => Ok ( 0 ) ,
42
+ }
30
43
}
31
44
32
45
impl InstanceService {
@@ -46,37 +59,35 @@ impl InstanceService {
46
59
}
47
60
48
61
pub async fn known_instances ( & self ) -> Result < u64 > {
49
- if self . randomize_statistics {
50
- return Ok ( random_statistic ( ) ) ;
51
- }
52
-
53
- with_connection ! ( self . db_pool, |db_conn| {
54
- accounts:: table
55
- . filter( accounts:: local. eq( false ) )
56
- . select( accounts:: domain)
57
- . distinct( )
58
- . count( )
59
- . get_result:: <i64 >( db_conn)
60
- . await
61
- . map( |count| count as u64 )
62
+ with_statistics_mode ( self . statistics_mode , async {
63
+ with_connection ! ( self . db_pool, |db_conn| {
64
+ accounts:: table
65
+ . filter( accounts:: local. eq( false ) )
66
+ . select( accounts:: domain)
67
+ . distinct( )
68
+ . count( )
69
+ . get_result:: <i64 >( db_conn)
70
+ . await
71
+ . map( |count| count as u64 )
72
+ } )
73
+ . map_err ( Error :: from)
62
74
} )
63
- . map_err ( Error :: from )
75
+ . await
64
76
}
65
77
66
78
pub async fn local_post_count ( & self ) -> Result < u64 > {
67
- if self . randomize_statistics {
68
- return Ok ( random_statistic ( ) ) ;
69
- }
70
-
71
- with_connection ! ( self . db_pool, |db_conn| {
72
- posts:: table
73
- . filter( posts:: is_local. eq( true ) )
74
- . count( )
75
- . get_result:: <i64 >( db_conn)
76
- . await
77
- . map( |count| count as u64 )
79
+ with_statistics_mode ( self . statistics_mode , async {
80
+ with_connection ! ( self . db_pool, |db_conn| {
81
+ posts:: table
82
+ . filter( posts:: is_local. eq( true ) )
83
+ . count( )
84
+ . get_result:: <i64 >( db_conn)
85
+ . await
86
+ . map( |count| count as u64 )
87
+ } )
88
+ . map_err ( Error :: from)
78
89
} )
79
- . map_err ( Error :: from )
90
+ . await
80
91
}
81
92
82
93
#[ must_use]
@@ -85,17 +96,16 @@ impl InstanceService {
85
96
}
86
97
87
98
pub async fn user_count ( & self ) -> Result < u64 > {
88
- if self . randomize_statistics {
89
- return Ok ( random_statistic ( ) ) ;
90
- }
91
-
92
- with_connection ! ( self . db_pool, |db_conn| {
93
- users:: table
94
- . count( )
95
- . get_result:: <i64 >( db_conn)
96
- . await
97
- . map( |count| count as u64 )
99
+ with_statistics_mode ( self . statistics_mode , async {
100
+ with_connection ! ( self . db_pool, |db_conn| {
101
+ users:: table
102
+ . count( )
103
+ . get_result:: <i64 >( db_conn)
104
+ . await
105
+ . map( |count| count as u64 )
106
+ } )
107
+ . map_err ( Error :: from)
98
108
} )
99
- . map_err ( Error :: from )
109
+ . await
100
110
}
101
111
}
0 commit comments