ticketgate/ticketgate.go

+package ticketgate
+
+import (
+	"os"
+	"runtime"
+	"strconv"
+	"sync"
+)
+
+type Gate struct {
+	c chan struct{}
+
+	wg sync.WaitGroup
+}
+
+var defaultConcurrency = 0
+
+func init() {
+	if os.Getenv("DEFAULT_CONCURRENCY") != "" {
+		defaultConcurrency, _ = strconv.Atoi(os.Getenv("DEFAULT_CONCURRENCY"))
+	}
+	if defaultConcurrency == 0 {
+		defaultConcurrency = runtime.NumCPU() / 2
+	}
+}
+
+func New(concurrency int) *Gate {
+	cc := concurrency
+	if concurrency < 0 {
+		cc = defaultConcurrency
+	}
+	if cc == 0 {
+		cc = 1
+	}
+	ch := make(chan struct{}, cc)
+	for i := 0; i < cc; i++ {
+		ch <- struct{}{}
+	}
+	return &Gate{
+		c: ch,
+	}
+}
+
+func (g *Gate) Take() {
+	g.wg.Add(1)
+	<-g.c
+}
+func (g *Gate) Done() {
+	g.wg.Done()
+	g.c <- struct{}{}
+}
+func (g *Gate) Wait() {
+	g.wg.Wait()
+}