github.com/unicornultrafoundation/go-u2u@v1.0.0-rc1.0.20240205080301-e74a83d3fadc/utils/dbutil/threads/pool.go (about)

     1  package threads
     2  
     3  import (
     4  	"math"
     5  	"runtime/debug"
     6  	"sync/atomic"
     7  )
     8  
     9  const GoroutinesPerThread = 0.8
    10  
    11  // threadPool counts threads in use
    12  type ThreadPool struct {
    13  	cap  int32
    14  	left int32
    15  }
    16  
    17  var GlobalPool ThreadPool
    18  
    19  // init ThreadPool only on demand to give time to other packages
    20  // call debug.SetMaxThreads() if they need.
    21  func (p *ThreadPool) init() {
    22  	initialized := !atomic.CompareAndSwapInt32(&p.cap, 0, math.MaxInt32)
    23  	if initialized {
    24  		return
    25  	}
    26  	cap := int32(getMaxThreads() * GoroutinesPerThread)
    27  	atomic.StoreInt32(&p.left, cap)
    28  	atomic.StoreInt32(&p.cap, cap)
    29  }
    30  
    31  // Capacity of pool.
    32  // Note: first call may return greater value than nexts. Don't cache it.
    33  func (p *ThreadPool) Cap() int {
    34  	p.init()
    35  	cap := atomic.LoadInt32(&p.cap)
    36  	return int(cap)
    37  }
    38  
    39  func (p *ThreadPool) Lock(want int) (got int, release func(count int)) {
    40  	p.init()
    41  
    42  	if want < 0 {
    43  		want = 0
    44  	}
    45  	if want > math.MaxInt32 {
    46  		want = math.MaxInt32
    47  	}
    48  
    49  	left := atomic.AddInt32(&p.left, -int32(want))
    50  	got = want
    51  	if left < 0 {
    52  		if left < -int32(got) {
    53  			left = -int32(got)
    54  		}
    55  		atomic.AddInt32(&p.left, -left)
    56  		got += int(left)
    57  	}
    58  
    59  	release = func(count int) {
    60  		if 0 > count || count > got {
    61  			count = got
    62  		}
    63  
    64  		got -= count
    65  		atomic.AddInt32(&p.left, int32(count))
    66  	}
    67  
    68  	return
    69  }
    70  
    71  func getMaxThreads() float64 {
    72  	was := debug.SetMaxThreads(10000)
    73  	debug.SetMaxThreads(was)
    74  	return float64(was)
    75  }