github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/ratelimit/counter.go (about)

     1  // Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package ratelimit
     4  
     5  import (
     6  	"github.com/TeaOSLab/EdgeNode/internal/zero"
     7  	"sync"
     8  )
     9  
    10  type Counter struct {
    11  	count     int
    12  	sem       chan zero.Zero
    13  	done      chan zero.Zero
    14  	closeOnce sync.Once
    15  }
    16  
    17  func NewCounter(count int) *Counter {
    18  	return &Counter{
    19  		count: count,
    20  		sem:   make(chan zero.Zero, count),
    21  		done:  make(chan zero.Zero),
    22  	}
    23  }
    24  
    25  func (this *Counter) Count() int {
    26  	return this.count
    27  }
    28  
    29  // Len 已占用数量
    30  func (this *Counter) Len() int {
    31  	return len(this.sem)
    32  }
    33  
    34  func (this *Counter) Ack() bool {
    35  	select {
    36  	case this.sem <- zero.New():
    37  		return true
    38  	case <-this.done:
    39  		return false
    40  	}
    41  }
    42  
    43  func (this *Counter) Release() {
    44  	select {
    45  	case <-this.sem:
    46  	default:
    47  		// 总是能Release成功
    48  	}
    49  }
    50  
    51  func (this *Counter) Close() {
    52  	this.closeOnce.Do(func() {
    53  		close(this.done)
    54  	})
    55  }