@@ -13,12 +13,13 @@ use crate::error::Error;
13
13
use crate :: http:: user:: timing:: sample_time;
14
14
use crate :: metrics:: {
15
15
CONCURRENT_CONNECTIONS_COUNT , CONNECTION_ALIVE_DURATION , CONNECTION_CREATE_TIME ,
16
+ TOTAL_RESPONSE_SIZE_HIST ,
16
17
} ;
17
18
use crate :: namespace:: meta_store:: MetaStore ;
18
19
use crate :: namespace:: NamespaceName ;
19
20
use crate :: query:: { Params , Query } ;
20
21
use crate :: query_analysis:: Statement ;
21
- use crate :: query_result_builder:: { IgnoreResult , QueryResultBuilder } ;
22
+ use crate :: query_result_builder:: { IgnoreResult , QueryResultBuilder , TOTAL_RESPONSE_SIZE } ;
22
23
use crate :: replication:: FrameNo ;
23
24
use crate :: Result ;
24
25
@@ -205,6 +206,7 @@ pub trait MakeConnection: Send + Sync + 'static {
205
206
timeout : Option < Duration > ,
206
207
max_total_response_size : u64 ,
207
208
max_concurrent_requests : u64 ,
209
+ disable_intelligent_throttling : bool ,
208
210
) -> MakeThrottledConnection < Self >
209
211
where
210
212
Self : Sized ,
@@ -215,6 +217,7 @@ pub trait MakeConnection: Send + Sync + 'static {
215
217
timeout,
216
218
max_total_response_size,
217
219
max_concurrent_requests,
220
+ disable_intelligent_throttling,
218
221
)
219
222
}
220
223
@@ -280,6 +283,7 @@ pub struct MakeThrottledConnection<F> {
280
283
max_total_response_size : u64 ,
281
284
waiters : AtomicUsize ,
282
285
max_concurrent_requests : u64 ,
286
+ disable_intelligent_throttling : bool ,
283
287
}
284
288
285
289
impl < F > MakeThrottledConnection < F > {
@@ -289,6 +293,7 @@ impl<F> MakeThrottledConnection<F> {
289
293
timeout : Option < Duration > ,
290
294
max_total_response_size : u64 ,
291
295
max_concurrent_requests : u64 ,
296
+ disable_intelligent_throttling : bool ,
292
297
) -> Self {
293
298
Self {
294
299
semaphore,
@@ -297,12 +302,16 @@ impl<F> MakeThrottledConnection<F> {
297
302
max_total_response_size,
298
303
waiters : AtomicUsize :: new ( 0 ) ,
299
304
max_concurrent_requests,
305
+ disable_intelligent_throttling,
300
306
}
301
307
}
302
308
303
309
// How many units should be acquired from the semaphore,
304
310
// depending on current memory pressure.
305
311
fn units_to_take ( & self ) -> u32 {
312
+ if self . disable_intelligent_throttling {
313
+ return 1 ;
314
+ }
306
315
let total_response_size = crate :: query_result_builder:: TOTAL_RESPONSE_SIZE
307
316
. load ( std:: sync:: atomic:: Ordering :: Relaxed ) as u64 ;
308
317
if total_response_size * 2 > self . max_total_response_size {
@@ -352,6 +361,8 @@ impl<F: MakeConnection> MakeConnection for MakeThrottledConnection<F> {
352
361
"Available semaphore units: {}" ,
353
362
self . semaphore. available_permits( )
354
363
) ;
364
+ TOTAL_RESPONSE_SIZE_HIST
365
+ . record ( TOTAL_RESPONSE_SIZE . load ( std:: sync:: atomic:: Ordering :: Relaxed ) as f64 ) ;
355
366
let units = self . units_to_take ( ) ;
356
367
let waiters_guard = WaitersGuard :: new ( & self . waiters ) ;
357
368
if ( waiters_guard. waiters . load ( Ordering :: Relaxed ) as u64 ) >= self . max_concurrent_requests {
@@ -519,6 +530,7 @@ pub mod test {
519
530
Some ( Duration :: from_millis ( 100 ) ) ,
520
531
u64:: MAX ,
521
532
u64:: MAX ,
533
+ false ,
522
534
) ;
523
535
524
536
let mut conns = Vec :: with_capacity ( 10 ) ;
0 commit comments