51 lines
950 B
Go
51 lines
950 B
Go
|
// Copyright (C) 2018 Storj Labs, Inc.
|
||
|
|
// See LICENSE for copying information
|
||
|
|
|
||
|
|
package sync2
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Limiter implements concurrent goroutine limiting
|
||
|
|
type Limiter struct {
|
||
|
|
limit chan struct{}
|
||
|
|
working sync.WaitGroup
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewLimiter creates a new limiter with limit set to n
|
||
|
|
func NewLimiter(n int) *Limiter {
|
||
|
|
limiter := &Limiter{}
|
||
|
|
limiter.limit = make(chan struct{}, n)
|
||
|
|
return limiter
|
||
|
|
}
|
||
|
|
|
||
|
|
// Go tries to starts fn as a goroutine.
|
||
|
|
// When the limit is reached it will wait until it can run it
|
||
|
|
// or the context is canceled.
|
||
|
|
func (limiter *Limiter) Go(ctx context.Context, fn func()) bool {
|
||
|
|
select {
|
||
|
|
case limiter.limit <- struct{}{}:
|
||
|
|
case <-ctx.Done():
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
limiter.working.Add(1)
|
||
|
|
go func() {
|
||
|
|
defer func() {
|
||
|
|
<-limiter.limit
|
||
|
|
limiter.working.Done()
|
||
|
|
}()
|
||
|
|
|
||
|
|
fn()
|
||
|
|
}()
|
||
|
|
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
// Wait waits for all running goroutines to finish
|
||
|
|
func (limiter *Limiter) Wait() {
|
||
|
|
limiter.working.Wait()
|
||
|
|
}
|