github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/fs/limiter.go (about)

     1  // Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package fsutils
     4  
     5  import (
     6  	"runtime"
     7  	"time"
     8  )
     9  
    10  var maxThreads = runtime.NumCPU()
    11  var WriterLimiter = NewLimiter(max(maxThreads, 8))
    12  var ReaderLimiter = NewLimiter(max(maxThreads, 8))
    13  
    14  type Limiter struct {
    15  	threads      chan struct{}
    16  	count        int
    17  	countDefault int
    18  	timers       chan *time.Timer
    19  }
    20  
    21  func NewLimiter(threads int) *Limiter {
    22  	if threads < 4 {
    23  		threads = 4
    24  	}
    25  	if threads > 64 {
    26  		threads = 64
    27  	}
    28  
    29  	var threadsChan = make(chan struct{}, threads)
    30  	for i := 0; i < threads; i++ {
    31  		threadsChan <- struct{}{}
    32  	}
    33  
    34  	return &Limiter{
    35  		countDefault: threads,
    36  		count:        threads,
    37  		threads:      threadsChan,
    38  		timers:       make(chan *time.Timer, 4096),
    39  	}
    40  }
    41  
    42  func (this *Limiter) SetThreads(newThreads int) {
    43  	if newThreads <= 0 {
    44  		newThreads = this.countDefault
    45  	}
    46  
    47  	if newThreads != this.count {
    48  		var threadsChan = make(chan struct{}, newThreads)
    49  		for i := 0; i < newThreads; i++ {
    50  			threadsChan <- struct{}{}
    51  		}
    52  
    53  		this.threads = threadsChan
    54  		this.count = newThreads
    55  	}
    56  }
    57  
    58  func (this *Limiter) Ack() {
    59  	<-this.threads
    60  }
    61  
    62  func (this *Limiter) TryAck() bool {
    63  	const timeoutDuration = 500 * time.Millisecond
    64  
    65  	var timeout *time.Timer
    66  	select {
    67  	case timeout = <-this.timers:
    68  		timeout.Reset(timeoutDuration)
    69  	default:
    70  		timeout = time.NewTimer(timeoutDuration)
    71  	}
    72  
    73  	defer func() {
    74  		timeout.Stop()
    75  
    76  		select {
    77  		case this.timers <- timeout:
    78  		default:
    79  		}
    80  	}()
    81  
    82  	select {
    83  	case <-this.threads:
    84  		return true
    85  	case <-timeout.C:
    86  		return false
    87  	}
    88  }
    89  
    90  func (this *Limiter) Release() {
    91  	select {
    92  	case this.threads <- struct{}{}:
    93  	default:
    94  		// 由于容量可能有变化，这里忽略多余的thread
    95  	}
    96  }
    97  
    98  func (this *Limiter) FreeThreads() int {
    99  	return len(this.threads)
   100  }