Andrew Thornton
6 years ago
11 changed files with 540 additions and 360 deletions
@ -1,82 +0,0 @@
|
||||
// Copyright 2019 The Gitea Authors. All rights reserved.
|
||||
// Use of this source code is governed by a MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package queue |
||||
|
||||
import ( |
||||
"context" |
||||
"time" |
||||
|
||||
"code.gitea.io/gitea/modules/log" |
||||
) |
||||
|
||||
// BatchedChannelQueueType is the type for batched channel queue
|
||||
const BatchedChannelQueueType Type = "batched-channel" |
||||
|
||||
// BatchedChannelQueueConfiguration is the configuration for a BatchedChannelQueue
|
||||
type BatchedChannelQueueConfiguration struct { |
||||
QueueLength int |
||||
BatchLength int |
||||
Workers int |
||||
} |
||||
|
||||
// BatchedChannelQueue implements
|
||||
type BatchedChannelQueue struct { |
||||
*ChannelQueue |
||||
batchLength int |
||||
} |
||||
|
||||
// NewBatchedChannelQueue create a memory channel queue
|
||||
func NewBatchedChannelQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, error) { |
||||
configInterface, err := toConfig(BatchedChannelQueueConfiguration{}, cfg) |
||||
if err != nil { |
||||
return nil, err |
||||
} |
||||
config := configInterface.(BatchedChannelQueueConfiguration) |
||||
return &BatchedChannelQueue{ |
||||
&ChannelQueue{ |
||||
queue: make(chan Data, config.QueueLength), |
||||
handle: handle, |
||||
exemplar: exemplar, |
||||
workers: config.Workers, |
||||
}, |
||||
config.BatchLength, |
||||
}, nil |
||||
} |
||||
|
||||
// Run starts to run the queue
|
||||
func (c *BatchedChannelQueue) Run(atShutdown, atTerminate func(context.Context, func())) { |
||||
atShutdown(context.Background(), func() { |
||||
log.Warn("BatchedChannelQueue is not shutdownable!") |
||||
}) |
||||
atTerminate(context.Background(), func() { |
||||
log.Warn("BatchedChannelQueue is not terminatable!") |
||||
}) |
||||
for i := 0; i < c.workers; i++ { |
||||
go func() { |
||||
delay := time.Millisecond * 300 |
||||
var datas = make([]Data, 0, c.batchLength) |
||||
for { |
||||
select { |
||||
case data := <-c.queue: |
||||
datas = append(datas, data) |
||||
if len(datas) >= c.batchLength { |
||||
c.handle(datas...) |
||||
datas = make([]Data, 0, c.batchLength) |
||||
} |
||||
case <-time.After(delay): |
||||
delay = time.Millisecond * 100 |
||||
if len(datas) > 0 { |
||||
c.handle(datas...) |
||||
datas = make([]Data, 0, c.batchLength) |
||||
} |
||||
} |
||||
} |
||||
}() |
||||
} |
||||
} |
||||
|
||||
func init() { |
||||
queuesMap[BatchedChannelQueueType] = NewBatchedChannelQueue |
||||
} |
||||
@ -1,46 +0,0 @@
|
||||
// Copyright 2019 The Gitea Authors. All rights reserved.
|
||||
// Use of this source code is governed by a MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package queue |
||||
|
||||
import "testing" |
||||
|
||||
import "github.com/stretchr/testify/assert" |
||||
|
||||
import "context" |
||||
|
||||
func TestBatchedChannelQueue(t *testing.T) { |
||||
handleChan := make(chan *testData) |
||||
handle := func(data ...Data) { |
||||
assert.True(t, len(data) == 2) |
||||
for _, datum := range data { |
||||
testDatum := datum.(*testData) |
||||
handleChan <- testDatum |
||||
} |
||||
} |
||||
|
||||
nilFn := func(_ context.Context, _ func()) {} |
||||
|
||||
queue, err := NewBatchedChannelQueue(handle, BatchedChannelQueueConfiguration{QueueLength: 20, BatchLength: 2, Workers: 1}, &testData{}) |
||||
assert.NoError(t, err) |
||||
|
||||
go queue.Run(nilFn, nilFn) |
||||
|
||||
test1 := testData{"A", 1} |
||||
test2 := testData{"B", 2} |
||||
|
||||
queue.Push(&test1) |
||||
go queue.Push(&test2) |
||||
|
||||
result1 := <-handleChan |
||||
assert.Equal(t, test1.TestString, result1.TestString) |
||||
assert.Equal(t, test1.TestInt, result1.TestInt) |
||||
|
||||
result2 := <-handleChan |
||||
assert.Equal(t, test2.TestString, result2.TestString) |
||||
assert.Equal(t, test2.TestInt, result2.TestInt) |
||||
|
||||
err = queue.Push(test1) |
||||
assert.Error(t, err) |
||||
} |
||||
@ -0,0 +1,239 @@
|
||||
// Copyright 2019 The Gitea Authors. All rights reserved.
|
||||
// Use of this source code is governed by a MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package queue |
||||
|
||||
import ( |
||||
"context" |
||||
"sync" |
||||
"time" |
||||
|
||||
"code.gitea.io/gitea/modules/log" |
||||
) |
||||
|
||||
// WorkerPool takes
|
||||
type WorkerPool struct { |
||||
lock sync.Mutex |
||||
baseCtx context.Context |
||||
cancel context.CancelFunc |
||||
cond *sync.Cond |
||||
numberOfWorkers int |
||||
batchLength int |
||||
handle HandlerFunc |
||||
dataChan chan Data |
||||
blockTimeout time.Duration |
||||
boostTimeout time.Duration |
||||
boostWorkers int |
||||
} |
||||
|
||||
// Push pushes the data to the internal channel
|
||||
func (p *WorkerPool) Push(data Data) { |
||||
p.lock.Lock() |
||||
if p.blockTimeout > 0 && p.boostTimeout > 0 { |
||||
p.lock.Unlock() |
||||
p.pushBoost(data) |
||||
} else { |
||||
p.lock.Unlock() |
||||
p.dataChan <- data |
||||
} |
||||
} |
||||
|
||||
func (p *WorkerPool) pushBoost(data Data) { |
||||
select { |
||||
case p.dataChan <- data: |
||||
default: |
||||
p.lock.Lock() |
||||
if p.blockTimeout <= 0 { |
||||
p.lock.Unlock() |
||||
p.dataChan <- data |
||||
return |
||||
} |
||||
ourTimeout := p.blockTimeout |
||||
timer := time.NewTimer(p.blockTimeout) |
||||
p.lock.Unlock() |
||||
select { |
||||
case p.dataChan <- data: |
||||
if timer.Stop() { |
||||
select { |
||||
case <-timer.C: |
||||
default: |
||||
} |
||||
} |
||||
case <-timer.C: |
||||
p.lock.Lock() |
||||
if p.blockTimeout > ourTimeout { |
||||
p.lock.Unlock() |
||||
p.dataChan <- data |
||||
return |
||||
} |
||||
p.blockTimeout *= 2 |
||||
log.Warn("Worker Channel blocked for %v - adding %d temporary workers for %s, block timeout now %v", ourTimeout, p.boostWorkers, p.boostTimeout, p.blockTimeout) |
||||
ctx, cancel := context.WithCancel(p.baseCtx) |
||||
go func() { |
||||
<-time.After(p.boostTimeout) |
||||
cancel() |
||||
p.lock.Lock() |
||||
p.blockTimeout /= 2 |
||||
p.lock.Unlock() |
||||
}() |
||||
p.addWorkers(ctx, p.boostWorkers) |
||||
p.lock.Unlock() |
||||
p.dataChan <- data |
||||
} |
||||
} |
||||
} |
||||
|
||||
// NumberOfWorkers returns the number of current workers in the pool
|
||||
func (p *WorkerPool) NumberOfWorkers() int { |
||||
p.lock.Lock() |
||||
defer p.lock.Unlock() |
||||
return p.numberOfWorkers |
||||
} |
||||
|
||||
// AddWorkers adds workers to the pool
|
||||
func (p *WorkerPool) AddWorkers(number int, timeout time.Duration) context.CancelFunc { |
||||
var ctx context.Context |
||||
var cancel context.CancelFunc |
||||
if timeout > 0 { |
||||
ctx, cancel = context.WithTimeout(p.baseCtx, timeout) |
||||
} else { |
||||
ctx, cancel = context.WithCancel(p.baseCtx) |
||||
} |
||||
|
||||
p.addWorkers(ctx, number) |
||||
return cancel |
||||
} |
||||
|
||||
// addWorkers adds workers to the pool
|
||||
func (p *WorkerPool) addWorkers(ctx context.Context, number int) { |
||||
for i := 0; i < number; i++ { |
||||
p.lock.Lock() |
||||
if p.cond == nil { |
||||
p.cond = sync.NewCond(&p.lock) |
||||
} |
||||
p.numberOfWorkers++ |
||||
p.lock.Unlock() |
||||
go func() { |
||||
p.doWork(ctx) |
||||
|
||||
p.lock.Lock() |
||||
p.numberOfWorkers-- |
||||
if p.numberOfWorkers <= 0 { |
||||
// numberOfWorkers can't go negative but...
|
||||
p.numberOfWorkers = 0 |
||||
p.cond.Broadcast() |
||||
} |
||||
p.lock.Unlock() |
||||
}() |
||||
} |
||||
} |
||||
|
||||
// Wait for WorkerPool to finish
|
||||
func (p *WorkerPool) Wait() { |
||||
p.lock.Lock() |
||||
defer p.lock.Unlock() |
||||
if p.cond == nil { |
||||
p.cond = sync.NewCond(&p.lock) |
||||
} |
||||
if p.numberOfWorkers <= 0 { |
||||
return |
||||
} |
||||
p.cond.Wait() |
||||
} |
||||
|
||||
// CleanUp will drain the remaining contents of the channel
|
||||
// This should be called after AddWorkers context is closed
|
||||
func (p *WorkerPool) CleanUp(ctx context.Context) { |
||||
log.Trace("CleanUp") |
||||
close(p.dataChan) |
||||
for data := range p.dataChan { |
||||
p.handle(data) |
||||
select { |
||||
case <-ctx.Done(): |
||||
log.Warn("Cleanup context closed before finishing clean-up") |
||||
return |
||||
default: |
||||
} |
||||
} |
||||
log.Trace("CleanUp done") |
||||
} |
||||
|
||||
func (p *WorkerPool) doWork(ctx context.Context) { |
||||
delay := time.Millisecond * 300 |
||||
var data = make([]Data, 0, p.batchLength) |
||||
for { |
||||
select { |
||||
case <-ctx.Done(): |
||||
if len(data) > 0 { |
||||
log.Trace("Handling: %d data, %v", len(data), data) |
||||
p.handle(data...) |
||||
} |
||||
log.Trace("Worker shutting down") |
||||
return |
||||
case datum, ok := <-p.dataChan: |
||||
if !ok { |
||||
// the dataChan has been closed - we should finish up:
|
||||
if len(data) > 0 { |
||||
log.Trace("Handling: %d data, %v", len(data), data) |
||||
p.handle(data...) |
||||
} |
||||
log.Trace("Worker shutting down") |
||||
return |
||||
} |
||||
data = append(data, datum) |
||||
if len(data) >= p.batchLength { |
||||
log.Trace("Handling: %d data, %v", len(data), data) |
||||
p.handle(data...) |
||||
data = make([]Data, 0, p.batchLength) |
||||
} |
||||
default: |
||||
timer := time.NewTimer(delay) |
||||
select { |
||||
case <-ctx.Done(): |
||||
if timer.Stop() { |
||||
select { |
||||
case <-timer.C: |
||||
default: |
||||
} |
||||
} |
||||
if len(data) > 0 { |
||||
log.Trace("Handling: %d data, %v", len(data), data) |
||||
p.handle(data...) |
||||
} |
||||
log.Trace("Worker shutting down") |
||||
return |
||||
case datum, ok := <-p.dataChan: |
||||
if timer.Stop() { |
||||
select { |
||||
case <-timer.C: |
||||
default: |
||||
} |
||||
} |
||||
if !ok { |
||||
// the dataChan has been closed - we should finish up:
|
||||
if len(data) > 0 { |
||||
log.Trace("Handling: %d data, %v", len(data), data) |
||||
p.handle(data...) |
||||
} |
||||
log.Trace("Worker shutting down") |
||||
return |
||||
} |
||||
data = append(data, datum) |
||||
if len(data) >= p.batchLength { |
||||
log.Trace("Handling: %d data, %v", len(data), data) |
||||
p.handle(data...) |
||||
data = make([]Data, 0, p.batchLength) |
||||
} |
||||
case <-timer.C: |
||||
delay = time.Millisecond * 100 |
||||
if len(data) > 0 { |
||||
log.Trace("Handling: %d data, %v", len(data), data) |
||||
p.handle(data...) |
||||
data = make([]Data, 0, p.batchLength) |
||||
} |
||||
|
||||
} |
||||
} |
||||
} |
||||
} |
||||
Loading…
Reference in new issue