storj/internal/sync2/limiter_test.go
2019-04-17 13:09:44 +03:00

74 lines
1.3 KiB
Go

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information
package sync2_test
import (
"context"
"sync/atomic"
"testing"
"time"
"storj.io/storj/internal/sync2"
)
func TestLimiterLimiting(t *testing.T) {
t.Parallel()
const N, Limit = 1000, 10
ctx := context.Background()
limiter := sync2.NewLimiter(Limit)
counter := int32(0)
for i := 0; i < N; i++ {
limiter.Go(ctx, func() {
if atomic.AddInt32(&counter, 1) > Limit {
panic("limit exceeded")
}
time.Sleep(time.Millisecond)
atomic.AddInt32(&counter, -1)
})
}
limiter.Wait()
}
func TestLimiterCanceling(t *testing.T) {
t.Parallel()
const N, Limit = 1000, 10
limiter := sync2.NewLimiter(Limit)
ctx, cancel := context.WithCancel(context.Background())
counter := int32(0)
waitForCancel := make(chan struct{}, N)
block := make(chan struct{})
allreturned := make(chan struct{})
go func() {
for i := 0; i < N; i++ {
limiter.Go(ctx, func() {
if atomic.AddInt32(&counter, 1) > Limit {
panic("limit exceeded")
}
waitForCancel <- struct{}{}
<-block
})
}
close(allreturned)
}()
for i := 0; i < Limit; i++ {
<-waitForCancel
}
cancel()
<-allreturned
close(block)
limiter.Wait()
if counter > Limit {
t.Fatal("too many times run")
}
}