1
1
use crate :: ffi:: CStr ;
2
+ #[ cfg( any( target_os = "android" , target_os = "linux" ) ) ]
3
+ use crate :: fs;
2
4
use crate :: mem:: { self , ManuallyDrop } ;
3
5
use crate :: num:: NonZero ;
4
6
#[ cfg( all( target_os = "linux" , target_env = "gnu" ) ) ]
@@ -405,6 +407,42 @@ fn truncate_cstr<const MAX_WITH_NUL: usize>(cstr: &CStr) -> [libc::c_char; MAX_W
405
407
result
406
408
}
407
409
410
+ #[ cfg( any( target_os = "android" , target_os = "linux" ) ) ]
411
+ fn count_user_threads ( ) -> Result < usize , io:: Error > {
412
+ let current_uid = unsafe { libc:: getuid ( ) } ;
413
+ let mut thread_count = 0 ;
414
+
415
+ for entry in fs:: read_dir ( "/proc" ) ? {
416
+ let entry = entry?;
417
+ let pid = entry. file_name ( ) . to_string_lossy ( ) . to_string ( ) ;
418
+
419
+ if let Ok ( _pid_num) = pid. parse :: < u32 > ( ) {
420
+ let status_path = format ! ( "/proc/{}/status" , pid) ;
421
+
422
+ if let Ok ( status) = fs:: read_to_string ( status_path) {
423
+ let mut uid: Option < libc:: uid_t > = None ;
424
+ let mut threads: Option < usize > = None ;
425
+
426
+ for line in status. lines ( ) {
427
+ if line. starts_with ( "Uid:" ) {
428
+ uid = line. split_whitespace ( ) . nth ( 1 ) . and_then ( |s| s. parse ( ) . ok ( ) ) ;
429
+ } else if line. starts_with ( "Threads:" ) {
430
+ threads = line. split_whitespace ( ) . nth ( 1 ) . and_then ( |s| s. parse ( ) . ok ( ) ) ;
431
+ }
432
+ }
433
+
434
+ if let ( Some ( uid) , Some ( t) ) = ( uid, threads) {
435
+ if uid == current_uid {
436
+ thread_count += t;
437
+ }
438
+ }
439
+ }
440
+ }
441
+ }
442
+
443
+ Ok ( thread_count)
444
+ }
445
+
408
446
pub fn available_parallelism ( ) -> io:: Result < NonZero < usize > > {
409
447
cfg_if:: cfg_if! {
410
448
if #[ cfg( any(
@@ -421,6 +459,10 @@ pub fn available_parallelism() -> io::Result<NonZero<usize>> {
421
459
#[ allow( unused_mut) ]
422
460
let mut quota = usize :: MAX ;
423
461
462
+ #[ allow( unused_assignments) ]
463
+ #[ allow( unused_mut) ]
464
+ let mut ulimit = libc:: rlim_t:: MAX ;
465
+
424
466
#[ cfg( any( target_os = "android" , target_os = "linux" ) ) ]
425
467
{
426
468
quota = cgroups:: quota( ) . max( 1 ) ;
@@ -439,14 +481,30 @@ pub fn available_parallelism() -> io::Result<NonZero<usize>> {
439
481
}
440
482
}
441
483
}
484
+
485
+ let mut r: libc:: rlimit = unsafe { mem:: zeroed( ) } ;
486
+ unsafe {
487
+ if libc:: getrlimit( libc:: RLIMIT_NPROC , & mut r) == 0 {
488
+ match r. rlim_cur {
489
+ libc:: RLIM_INFINITY => ulimit = libc:: rlim_t:: MAX ,
490
+ soft_limit => {
491
+ ulimit = match count_user_threads( ) {
492
+ Ok ( t) => if soft_limit > t as libc:: rlim_t { soft_limit - t as libc:: rlim_t} else { 1 } ,
493
+ _ => 1
494
+ }
495
+ }
496
+ }
497
+ }
498
+ }
442
499
}
500
+
443
501
match unsafe { libc:: sysconf( libc:: _SC_NPROCESSORS_ONLN) } {
444
502
-1 => Err ( io:: Error :: last_os_error( ) ) ,
445
503
0 => Err ( io:: Error :: UNKNOWN_THREAD_COUNT ) ,
446
504
cpus => {
447
505
let count = cpus as usize ;
448
506
// Cover the unusual situation where we were able to get the quota but not the affinity mask
449
- let count = count. min( quota) ;
507
+ let count = count. min( quota. min ( ulimit . try_into ( ) . unwrap_or ( usize :: MAX ) ) ) ;
450
508
Ok ( unsafe { NonZero :: new_unchecked( count) } )
451
509
}
452
510
}
0 commit comments