Rewrite console throttling logic; drop complex timer usage and use a very simple throttle
This also removes server process termination logic when a server is breaching the output limits. It simply continues to efficiently throttle the console output.
This commit is contained in:
Dane Everitt
83
system/locker.go
Normal file
83
system/locker.go
Normal file
@@ -0,0 +1,83 @@
|
||||
package system
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"emperror.dev/errors"
|
||||
)
|
||||
|
||||
var ErrLockerLocked = errors.Sentinel("locker: cannot acquire lock, already locked")
|
||||
|
||||
type Locker struct {
|
||||
mu sync.RWMutex
|
||||
ch chan bool
|
||||
}
|
||||
|
||||
// NewLocker returns a new Locker instance.
|
||||
func NewLocker() *Locker {
|
||||
return &Locker{
|
||||
ch: make(chan bool, 1),
|
||||
}
|
||||
}
|
||||
|
||||
// IsLocked returns the current state of the locker channel. If there is
|
||||
// currently a value in the channel, it is assumed to be locked.
|
||||
func (l *Locker) IsLocked() bool {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
return len(l.ch) == 1
|
||||
}
|
||||
|
||||
// Acquire will acquire the power lock if it is not currently locked. If it is
|
||||
// already locked, acquire will fail to acquire the lock, and will return false.
|
||||
func (l *Locker) Acquire() error {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
select {
|
||||
case l.ch <- true:
|
||||
default:
|
||||
return ErrLockerLocked
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
// TryAcquire will attempt to acquire a power-lock until the context provided
|
||||
// is canceled.
|
||||
func (l *Locker) TryAcquire(ctx context.Context) error {
|
||||
select {
|
||||
case l.ch <- true:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// Release will drain the locker channel so that we can properly re-acquire it
|
||||
// at a later time. If the channel is not currently locked this function is a
|
||||
// no-op and will immediately return.
|
||||
func (l *Locker) Release() {
|
||||
l.mu.Lock()
|
||||
select {
|
||||
case <-l.ch:
|
||||
default:
|
||||
}
|
||||
l.mu.Unlock()
|
||||
}
|
||||
|
||||
// Destroy cleans up the power locker by closing the channel.
|
||||
func (l *Locker) Destroy() {
|
||||
l.mu.Lock()
|
||||
if l.ch != nil {
|
||||
select {
|
||||
case <-l.ch:
|
||||
default:
|
||||
}
|
||||
close(l.ch)
|
||||
}
|
||||
l.mu.Unlock()
|
||||
}
|
||||
148
system/locker_test.go
Normal file
148
system/locker_test.go
Normal file
@@ -0,0 +1,148 @@
|
||||
package system
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"emperror.dev/errors"
|
||||
. "github.com/franela/goblin"
|
||||
)
|
||||
|
||||
func TestPower(t *testing.T) {
|
||||
g := Goblin(t)
|
||||
|
||||
g.Describe("Locker", func() {
|
||||
var l *Locker
|
||||
g.BeforeEach(func() {
|
||||
l = NewLocker()
|
||||
})
|
||||
|
||||
g.Describe("PowerLocker#IsLocked", func() {
|
||||
g.It("should return false when the channel is empty", func() {
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(l.IsLocked()).IsFalse()
|
||||
})
|
||||
|
||||
g.It("should return true when the channel is at capacity", func() {
|
||||
l.ch <- true
|
||||
|
||||
g.Assert(l.IsLocked()).IsTrue()
|
||||
<-l.ch
|
||||
g.Assert(l.IsLocked()).IsFalse()
|
||||
|
||||
// We don't care what the channel value is, just that there is
|
||||
// something in it.
|
||||
l.ch <- false
|
||||
g.Assert(l.IsLocked()).IsTrue()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
})
|
||||
})
|
||||
|
||||
g.Describe("PowerLocker#Acquire", func() {
|
||||
g.It("should acquire a lock when channel is empty", func() {
|
||||
err := l.Acquire()
|
||||
|
||||
g.Assert(err).IsNil()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(len(l.ch)).Equal(1)
|
||||
})
|
||||
|
||||
g.It("should return an error when the channel is full", func() {
|
||||
l.ch <- true
|
||||
|
||||
err := l.Acquire()
|
||||
|
||||
g.Assert(err).IsNotNil()
|
||||
g.Assert(errors.Is(err, ErrLockerLocked)).IsTrue()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(len(l.ch)).Equal(1)
|
||||
})
|
||||
})
|
||||
|
||||
g.Describe("PowerLocker#TryAcquire", func() {
|
||||
g.It("should acquire a lock when channel is empty", func() {
|
||||
g.Timeout(time.Second)
|
||||
|
||||
err := l.TryAcquire(context.Background())
|
||||
|
||||
g.Assert(err).IsNil()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(len(l.ch)).Equal(1)
|
||||
g.Assert(l.IsLocked()).IsTrue()
|
||||
})
|
||||
|
||||
g.It("should block until context is canceled if channel is full", func() {
|
||||
g.Timeout(time.Second)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*500)
|
||||
defer cancel()
|
||||
|
||||
l.ch <- true
|
||||
err := l.TryAcquire(ctx)
|
||||
|
||||
g.Assert(err).IsNotNil()
|
||||
g.Assert(errors.Is(err, context.DeadlineExceeded)).IsTrue()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(len(l.ch)).Equal(1)
|
||||
g.Assert(l.IsLocked()).IsTrue()
|
||||
})
|
||||
|
||||
g.It("should block until lock can be acquired", func() {
|
||||
g.Timeout(time.Second)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*200)
|
||||
defer cancel()
|
||||
|
||||
l.Acquire()
|
||||
go func() {
|
||||
time.AfterFunc(time.Millisecond * 50, func() {
|
||||
l.Release()
|
||||
})
|
||||
}()
|
||||
|
||||
err := l.TryAcquire(ctx)
|
||||
g.Assert(err).IsNil()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(len(l.ch)).Equal(1)
|
||||
g.Assert(l.IsLocked()).IsTrue()
|
||||
})
|
||||
})
|
||||
|
||||
g.Describe("PowerLocker#Release", func() {
|
||||
g.It("should release when channel is full", func() {
|
||||
l.Acquire()
|
||||
g.Assert(l.IsLocked()).IsTrue()
|
||||
l.Release()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(len(l.ch)).Equal(0)
|
||||
g.Assert(l.IsLocked()).IsFalse()
|
||||
})
|
||||
|
||||
g.It("should release when channel is empty", func() {
|
||||
g.Assert(l.IsLocked()).IsFalse()
|
||||
l.Release()
|
||||
g.Assert(cap(l.ch)).Equal(1)
|
||||
g.Assert(len(l.ch)).Equal(0)
|
||||
g.Assert(l.IsLocked()).IsFalse()
|
||||
})
|
||||
})
|
||||
|
||||
g.Describe("PowerLocker#Destroy", func() {
|
||||
g.It("should unlock and close the channel", func() {
|
||||
l.Acquire()
|
||||
g.Assert(l.IsLocked()).IsTrue()
|
||||
l.Destroy()
|
||||
g.Assert(l.IsLocked()).IsFalse()
|
||||
|
||||
defer func() {
|
||||
r := recover()
|
||||
|
||||
g.Assert(r).IsNotNil()
|
||||
g.Assert(r.(error).Error()).Equal("send on closed channel")
|
||||
}()
|
||||
|
||||
l.Acquire()
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
50
system/rate.go
Normal file
50
system/rate.go
Normal file
@@ -0,0 +1,50 @@
|
||||
package system
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Rate defines a rate limiter of n items (limit) per duration of time.
|
||||
type Rate struct {
|
||||
mu sync.Mutex
|
||||
limit uint64
|
||||
duration time.Duration
|
||||
count uint64
|
||||
last time.Time
|
||||
}
|
||||
|
||||
func NewRate(limit uint64, duration time.Duration) *Rate {
|
||||
return &Rate{
|
||||
limit: limit,
|
||||
duration: duration,
|
||||
last: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
// Try returns true if under the rate limit defined, or false if the rate limit
|
||||
// has been exceeded for the current duration.
|
||||
func (r *Rate) Try() bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
now := time.Now()
|
||||
// If it has been more than the duration, reset the timer and count.
|
||||
if now.Sub(r.last) > r.duration {
|
||||
r.count = 0
|
||||
r.last = now
|
||||
}
|
||||
if (r.count + 1) > r.limit {
|
||||
return false
|
||||
}
|
||||
// Hit this once, and return.
|
||||
r.count = r.count + 1
|
||||
return true
|
||||
}
|
||||
|
||||
// Reset resets the internal state of the rate limiter back to zero.
|
||||
func (r *Rate) Reset() {
|
||||
r.mu.Lock()
|
||||
r.count = 0
|
||||
r.last = time.Now()
|
||||
r.mu.Unlock()
|
||||
}
|
||||
67
system/rate_test.go
Normal file
67
system/rate_test.go
Normal file
@@ -0,0 +1,67 @@
|
||||
package system
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
. "github.com/franela/goblin"
|
||||
)
|
||||
|
||||
func TestRate(t *testing.T) {
|
||||
g := Goblin(t)
|
||||
|
||||
g.Describe("Rate", func() {
|
||||
g.It("properly rate limits a bucket", func() {
|
||||
r := NewRate(10, time.Millisecond*100)
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
ok := r.Try()
|
||||
if i < 10 && !ok {
|
||||
g.Failf("should not have allowed take on try %d", i)
|
||||
} else if i >= 10 && ok {
|
||||
g.Failf("should have blocked take on try %d", i)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
g.It("handles rate limiting in chunks", func() {
|
||||
var out []int
|
||||
r := NewRate(12, time.Millisecond*10)
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
if i%20 == 0 {
|
||||
// Give it time to recover.
|
||||
time.Sleep(time.Millisecond * 10)
|
||||
}
|
||||
if r.Try() {
|
||||
out = append(out, i)
|
||||
}
|
||||
}
|
||||
|
||||
g.Assert(len(out)).Equal(60)
|
||||
g.Assert(out[0]).Equal(0)
|
||||
g.Assert(out[12]).Equal(20)
|
||||
g.Assert(out[len(out)-1]).Equal(91)
|
||||
})
|
||||
|
||||
g.It("resets back to zero when called", func() {
|
||||
r := NewRate(10, time.Second)
|
||||
for i := 0; i < 100; i++ {
|
||||
if i % 10 == 0 {
|
||||
r.Reset()
|
||||
}
|
||||
g.Assert(r.Try()).IsTrue()
|
||||
}
|
||||
g.Assert(r.Try()).IsFalse("final attempt should not allow taking")
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkRate_Try(b *testing.B) {
|
||||
r := NewRate(10, time.Millisecond*100)
|
||||
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
r.Try()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user