mirror of
https://github.com/netbirdio/netbird.git
synced 2026-04-16 07:16:38 +00:00
9ba067391f
- Remove WaitGroup, make SemaphoreGroup a pure semaphore - Make Add() return error instead of silently failing on context cancel - Remove context parameter from Done() to prevent slot leaks - Fix missing Done() call in conn.go error path
34 lines
722 B
Go
34 lines
722 B
Go
package semaphoregroup
|
|
|
|
import (
|
|
"context"
|
|
)
|
|
|
|
// SemaphoreGroup is a custom type that combines sync.WaitGroup and a semaphore.
|
|
type SemaphoreGroup struct {
|
|
semaphore chan struct{}
|
|
}
|
|
|
|
// NewSemaphoreGroup creates a new SemaphoreGroup with the specified semaphore limit.
|
|
func NewSemaphoreGroup(limit int) *SemaphoreGroup {
|
|
return &SemaphoreGroup{
|
|
semaphore: make(chan struct{}, limit),
|
|
}
|
|
}
|
|
|
|
// Add acquire a slot
|
|
func (sg *SemaphoreGroup) Add(ctx context.Context) error {
|
|
// Acquire semaphore slot
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case sg.semaphore <- struct{}{}:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// Done releases a slot. Must be called after a successful Add.
|
|
func (sg *SemaphoreGroup) Done() {
|
|
<-sg.semaphore
|
|
}
|