github.com/rolandhe/saber@v0.0.4/gocc/defaultSemaphore.go (about)

     1  // Golang concurrent tools like java juc.
     2  //
     3  // Copyright 2023 The saber Authors. All rights reserved.
     4  
     5  package gocc
     6  
     7  import "time"
     8  
     9  // NewDefaultSemaphore 构建指定总量的信号量
    10  //
    11  //	limit  信号量的总量
    12  func NewDefaultSemaphore(limit uint) Semaphore {
    13  	return &semaphoreChan{
    14  		make(chan struct{}, limit),
    15  		limit,
    16  	}
    17  }
    18  
    19  // Semaphore 类似java的信号量, 但只支持单个信号量,支持超时
    20  type Semaphore interface {
    21  	// TryAcquire 尝试获取单个信号量,如果当下没有信号量,则直接返回false
    22  	TryAcquire() bool
    23  
    24  	// Acquire 获取单个信号量,如果当前没有信号量,一直阻塞等待,直到获取到位置
    25  	Acquire()
    26  
    27  	// AcquireTimeout 超时获取信号量,如果超时时间内没有信号量则返回false,否则返回true
    28  	//
    29  	//d == 0 退化成 TryAcquire
    30  	//
    31  	//d < 0 退出成 Acquire
    32  	AcquireTimeout(d time.Duration) bool
    33  
    34  	// Release 释放信号量
    35  	Release()
    36  	
    37  	// TotalTokens 信号量总数
    38  	TotalTokens() uint
    39  }
    40  
    41  // 基于channel实现信号量,这也是golang官方文档中的推荐实现
    42  type semaphoreChan struct {
    43  	tokens chan struct{}
    44  	total  uint
    45  }
    46  
    47  func (s *semaphoreChan) TryAcquire() bool {
    48  	select {
    49  	case s.tokens <- struct{}{}:
    50  		return true
    51  	default:
    52  		return false
    53  	}
    54  }
    55  
    56  func (s *semaphoreChan) Acquire() {
    57  	s.tokens <- struct{}{}
    58  }
    59  func (s *semaphoreChan) AcquireTimeout(d time.Duration) bool {
    60  	if d == 0 {
    61  		return s.TryAcquire()
    62  	}
    63  	if d < 0 {
    64  		s.Acquire()
    65  		return true
    66  	}
    67  
    68  	select {
    69  	case s.tokens <- struct{}{}:
    70  		return true
    71  	case <-time.After(d):
    72  		return false
    73  	}
    74  }
    75  
    76  func (s *semaphoreChan) Release() {
    77  	<-s.tokens
    78  }
    79  func (s *semaphoreChan) TotalTokens() uint {
    80  	return s.total
    81  }