@@ -7,7 +7,7 @@ use crate::sys::weak::dlsym;
7
7
use crate :: sys:: weak:: weak;
8
8
use crate :: sys:: { os, stack_overflow} ;
9
9
use crate :: time:: { Duration , Instant } ;
10
- use crate :: { cmp, io, ptr} ;
10
+ use crate :: { cmp, fs , io, ptr} ;
11
11
#[ cfg( not( any(
12
12
target_os = "l4re" ,
13
13
target_os = "vxworks" ,
@@ -405,6 +405,41 @@ fn truncate_cstr<const MAX_WITH_NUL: usize>(cstr: &CStr) -> [libc::c_char; MAX_W
405
405
result
406
406
}
407
407
408
+ fn count_user_threads ( ) -> Result < usize , io:: Error > {
409
+ let current_uid = unsafe { libc:: getuid ( ) } ;
410
+ let mut thread_count = 0 ;
411
+
412
+ for entry in fs:: read_dir ( "/proc" ) ? {
413
+ let entry = entry?;
414
+ let pid = entry. file_name ( ) . to_string_lossy ( ) . to_string ( ) ;
415
+
416
+ if let Ok ( _pid_num) = pid. parse :: < u32 > ( ) {
417
+ let status_path = format ! ( "/proc/{}/status" , pid) ;
418
+
419
+ if let Ok ( status) = fs:: read_to_string ( status_path) {
420
+ let mut uid: Option < libc:: uid_t > = None ;
421
+ let mut threads: Option < usize > = None ;
422
+
423
+ for line in status. lines ( ) {
424
+ if line. starts_with ( "Uid:" ) {
425
+ uid = line. split_whitespace ( ) . nth ( 1 ) . and_then ( |s| s. parse ( ) . ok ( ) ) ;
426
+ } else if line. starts_with ( "Threads:" ) {
427
+ threads = line. split_whitespace ( ) . nth ( 1 ) . and_then ( |s| s. parse ( ) . ok ( ) ) ;
428
+ }
429
+ }
430
+
431
+ if let ( Some ( uid) , Some ( t) ) = ( uid, threads) {
432
+ if uid == current_uid {
433
+ thread_count += t;
434
+ }
435
+ }
436
+ }
437
+ }
438
+ }
439
+
440
+ Ok ( thread_count)
441
+ }
442
+
408
443
pub fn available_parallelism ( ) -> io:: Result < NonZero < usize > > {
409
444
cfg_if:: cfg_if! {
410
445
if #[ cfg( any(
@@ -420,6 +455,7 @@ pub fn available_parallelism() -> io::Result<NonZero<usize>> {
420
455
#[ allow( unused_assignments) ]
421
456
#[ allow( unused_mut) ]
422
457
let mut quota = usize :: MAX ;
458
+ let mut ulimit = libc:: rlim_t:: MAX ;
423
459
424
460
#[ cfg( any( target_os = "android" , target_os = "linux" ) ) ]
425
461
{
@@ -439,14 +475,30 @@ pub fn available_parallelism() -> io::Result<NonZero<usize>> {
439
475
}
440
476
}
441
477
}
478
+
479
+ let mut r: libc:: rlimit = unsafe { mem:: zeroed( ) } ;
480
+ unsafe {
481
+ if libc:: getrlimit( libc:: RLIMIT_NPROC , & mut r) == 0 {
482
+ match r. rlim_cur {
483
+ libc:: RLIM_INFINITY => ulimit = libc:: rlim_t:: MAX ,
484
+ soft_limit => {
485
+ ulimit = match count_user_threads( ) {
486
+ Ok ( t) => soft_limit - t as libc:: rlim_t,
487
+ _ => 1
488
+ }
489
+ }
490
+ }
491
+ }
492
+ }
442
493
}
494
+
443
495
match unsafe { libc:: sysconf( libc:: _SC_NPROCESSORS_ONLN) } {
444
496
-1 => Err ( io:: Error :: last_os_error( ) ) ,
445
497
0 => Err ( io:: Error :: UNKNOWN_THREAD_COUNT ) ,
446
498
cpus => {
447
499
let count = cpus as usize ;
448
500
// Cover the unusual situation where we were able to get the quota but not the affinity mask
449
- let count = count. min( quota) ;
501
+ let count = count. min( quota. min ( ulimit . try_into ( ) . unwrap_or ( usize :: MAX ) ) ) ;
450
502
Ok ( unsafe { NonZero :: new_unchecked( count) } )
451
503
}
452
504
}
0 commit comments