67 lines
1.0 KiB
Go
67 lines
1.0 KiB
Go
package tcp
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
)
|
|
|
|
// TaskPool runs inbound business (MySQL included). EventLoop must only Submit.
|
|
type TaskPool interface {
|
|
Submit(ctx context.Context, task func()) error
|
|
Close()
|
|
}
|
|
|
|
type pool struct {
|
|
tasks chan func()
|
|
workers int
|
|
wg sync.WaitGroup
|
|
once sync.Once
|
|
}
|
|
|
|
func NewTaskPool(workers, queue int) TaskPool {
|
|
if workers < 1 {
|
|
workers = 1
|
|
}
|
|
if queue < 1 {
|
|
queue = 1
|
|
}
|
|
p := &pool{tasks: make(chan func(), queue), workers: workers}
|
|
for i := 0; i < workers; i++ {
|
|
p.wg.Add(1)
|
|
go p.loop()
|
|
}
|
|
return p
|
|
}
|
|
|
|
func (p *pool) loop() {
|
|
defer p.wg.Done()
|
|
for task := range p.tasks {
|
|
if task != nil {
|
|
task()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (p *pool) Submit(ctx context.Context, task func()) error {
|
|
if task == nil {
|
|
return errs.ErrInvalidArgument
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case p.tasks <- task:
|
|
return nil
|
|
default:
|
|
return errs.ErrBusy
|
|
}
|
|
}
|
|
|
|
func (p *pool) Close() {
|
|
p.once.Do(func() {
|
|
close(p.tasks)
|
|
p.wg.Wait()
|
|
})
|
|
}
|