77 lines
1.1 KiB
Go
77 lines
1.1 KiB
Go
|
package taskPools
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"log"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Pools struct {
|
||
|
|
ch chan struct{}
|
||
|
|
wg *sync.WaitGroup
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewPools(n int) *Pools {
|
||
|
|
if n <= 0 {
|
||
|
|
panic("n must >= 1")
|
||
|
|
}
|
||
|
|
c := make(chan struct{}, n)
|
||
|
|
for i := 0; i < n; i++ {
|
||
|
|
c <- struct{}{}
|
||
|
|
}
|
||
|
|
return &Pools{
|
||
|
|
ch: c,
|
||
|
|
wg: &sync.WaitGroup{},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *Pools) ExecuteWithTimeOut(timeout time.Duration, fn func(), args ...string) {
|
||
|
|
if timeout <= 0 {
|
||
|
|
p.Execute(fn)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
p.wg.Add(1)
|
||
|
|
q := <-p.ch
|
||
|
|
go func() {
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||
|
|
done := make(chan struct{}, 1)
|
||
|
|
defer func() {
|
||
|
|
cancel()
|
||
|
|
p.wg.Done()
|
||
|
|
p.ch <- q
|
||
|
|
}()
|
||
|
|
go func() {
|
||
|
|
fn()
|
||
|
|
done <- struct{}{}
|
||
|
|
}()
|
||
|
|
select {
|
||
|
|
case <-ctx.Done():
|
||
|
|
if len(args) > 0 && args[0] != "" {
|
||
|
|
log.Printf("执行%s超时", args[0])
|
||
|
|
}
|
||
|
|
case <-done:
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *Pools) Execute(fn func()) {
|
||
|
|
if cap(p.ch) == 1 {
|
||
|
|
fn()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
p.wg.Add(1)
|
||
|
|
q := <-p.ch
|
||
|
|
go func() {
|
||
|
|
defer func() {
|
||
|
|
p.wg.Done()
|
||
|
|
p.ch <- q
|
||
|
|
}()
|
||
|
|
fn()
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *Pools) Wait() {
|
||
|
|
p.wg.Wait()
|
||
|
|
}
|