@@ -44,52 +44,23 @@ package pool
4444
4545import (
4646 "context"
47- "sync"
48-
49- "github.com/DoNewsCode/core/ctxmeta"
50-
51- "github.com/oklog/run"
5247)
5348
54- // NewPool returned func(contract.Dispatcher) *Pool
55- func NewPool (options ... ProviderOptionFunc ) * Pool {
49+ // NewPool returns *Pool
50+ func NewPool (manager * Manager , cap int ) * Pool {
5651 pool := Pool {
57- ch : make (chan job ),
58- concurrency : 10 ,
59- }
60- for _ , f := range options {
61- f (& pool )
52+ manager : manager ,
53+ concurrency : make (chan struct {}, cap ),
6254 }
6355 return & pool
6456}
6557
66- type job struct {
67- fn func ()
68- }
69-
7058// Pool is an async worker pool. It can be used to dispatch the async jobs from
7159// web servers. See the package documentation about its advantage over creating a
7260// goroutine directly.
7361type Pool struct {
74- ch chan job
75- concurrency int
76- counter * Counter
77- }
78-
79- // ProvideRunGroup implements core.RunProvider
80- func (p * Pool ) ProvideRunGroup (group * run.Group ) {
81- ctx , cancel := context .WithCancel (context .Background ())
82-
83- group .Add (func () error {
84- return p .Run (ctx )
85- }, func (err error ) {
86- cancel ()
87- })
88- }
89-
90- // Module implements di.Modular
91- func (p * Pool ) Module () interface {} {
92- return p
62+ manager * Manager
63+ concurrency chan struct {}
9364}
9465
9566// Go dispatchers a job to the async worker pool. requestContext is the context
@@ -98,37 +69,18 @@ func (p *Pool) Module() interface{} {
9869// nothing to do with the request. If the pool has reached max concurrency, the job will
9970// be executed in the current goroutine. In other word, the job will be executed synchronously.
10071func (p * Pool ) Go (requestContext context.Context , function func (asyncContext context.Context )) {
101- j := job {
102- fn : func () {
103- function (ctxmeta .WithoutCancel (requestContext ))
104- },
105- }
106- select {
107- case p .ch <- j :
108- default :
109- p .counter .IncSyncJob ()
110- j .fn ()
111- }
72+ p .concurrency <- struct {}{}
73+ p .manager .Go (requestContext , func (ctx context.Context ) {
74+ defer func () {
75+ <- p .concurrency
76+ }()
77+ function (ctx )
78+ })
11279}
11380
114- // Run starts the async worker pool and block until it finishes.
115- func (p * Pool ) Run (ctx context.Context ) error {
116- var wg sync.WaitGroup
117- for i := 0 ; i < p .concurrency ; i ++ {
118- wg .Add (1 )
119- go func () {
120- defer wg .Done ()
121- for {
122- select {
123- case j := <- p .ch :
124- p .counter .IncAsyncJob ()
125- j .fn ()
126- case <- ctx .Done ():
127- return
128- }
129- }
130- }()
81+ // Wait waits for all the async jobs to finish.
82+ func (p * Pool ) Wait () {
83+ for i := 0 ; i < cap (p .concurrency ); i ++ {
84+ p .concurrency <- struct {}{}
13185 }
132- wg .Wait ()
133- return nil
13486}
0 commit comments