airobot/lib/tickets.go
2023-01-04 10:37:21 +08:00

86 lines
1.6 KiB
Go

package lib
import (
"errors"
"fmt"
"sync/atomic"
"time"
)
// Tickets 表示Goroutine票池的接口。
type Tickets interface {
// 拿走一张票。
Take(useCount, inteval int32)
// 归还一张票。
Return()
// 票池是否已被激活。
Active() bool
// 票的总数。
Total() uint32
// 剩余的票数。
Remainder() uint32
}
// myTickets 表示Goroutine票池的实现。
type myTickets struct {
total uint32 // 票的总数。
ticketCh chan struct{} // 票的容器。
active bool // 票池是否已被激活。
count int32
}
// NewTickets 会新建一个Goroutine票池。
func NewTickets(total uint32) (Tickets, error) {
gt := myTickets{}
if !gt.init(total) {
errMsg :=
fmt.Sprintf("The goroutine ticket pool can NOT be initialized! (total=%d)\n", total)
return nil, errors.New(errMsg)
}
return &gt, nil
}
func (gt *myTickets) init(total uint32) bool {
if gt.active {
return false
}
if total == 0 {
return false
}
ch := make(chan struct{}, total)
n := int(total)
for i := 0; i < n; i++ {
ch <- struct{}{}
}
gt.ticketCh = ch
gt.total = total
gt.active = true
return true
}
func (gt *myTickets) Take(useCount, inteval int32) {
count := atomic.LoadInt32(&gt.count)
if count >= int32(useCount) {
time.Sleep(time.Duration(inteval) * time.Second)
atomic.StoreInt32(&gt.count, 0)
}
<-gt.ticketCh
atomic.AddInt32(&gt.count, 1)
}
func (gt *myTickets) Return() {
gt.ticketCh <- struct{}{}
}
func (gt *myTickets) Active() bool {
return gt.active
}
func (gt *myTickets) Total() uint32 {
return gt.total
}
func (gt *myTickets) Remainder() uint32 {
return uint32(len(gt.ticketCh))
}