-
-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add MultiPool and MultiPoolWithFunc (#305)
- Loading branch information
Showing
7 changed files
with
524 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,205 @@ | ||
| // MIT License | ||
|
|
||
| // Copyright (c) 2023 Andy Pan | ||
|
|
||
| // Permission is hereby granted, free of charge, to any person obtaining a copy | ||
| // of this software and associated documentation files (the "Software"), to deal | ||
| // in the Software without restriction, including without limitation the rights | ||
| // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | ||
| // copies of the Software, and to permit persons to whom the Software is | ||
| // furnished to do so, subject to the following conditions: | ||
| // | ||
| // The above copyright notice and this permission notice shall be included in all | ||
| // copies or substantial portions of the Software. | ||
| // | ||
| // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | ||
| // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | ||
| // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | ||
| // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | ||
| // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | ||
| // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | ||
| // SOFTWARE. | ||
|
|
||
| package ants | ||
|
|
||
| import ( | ||
| "errors" | ||
| "fmt" | ||
| "strings" | ||
| "sync/atomic" | ||
| "time" | ||
| ) | ||
|
|
||
| // LoadBalancingStrategy represents the type of load-balancing algorithm. | ||
| type LoadBalancingStrategy int | ||
|
|
||
| const ( | ||
| // RoundRobin distributes task to a list of pools in rotation. | ||
| RoundRobin LoadBalancingStrategy = 1 << (iota + 1) | ||
|
|
||
| // LeastTasks always selects the pool with the least number of pending tasks. | ||
| LeastTasks | ||
| ) | ||
|
|
||
| // MultiPool consists of multiple pools, from which you will benefit the | ||
| // performance improvement on basis of the fine-grained locking that reduces | ||
| // the lock contention. | ||
| // MultiPool is a good fit for the scenario that you have a large number of | ||
| // tasks to submit, and you don't want the single pool to be the bottleneck. | ||
| type MultiPool struct { | ||
| pools []*Pool | ||
| index uint32 | ||
| state int32 | ||
| lbs LoadBalancingStrategy | ||
| } | ||
|
|
||
| // NewMultiPool instantiates a MultiPool with a size of the pool list and a size | ||
| // per pool, and the load-balancing strategy. | ||
| func NewMultiPool(size, sizePerPool int, lbs LoadBalancingStrategy, options ...Option) (*MultiPool, error) { | ||
| pools := make([]*Pool, size) | ||
| for i := 0; i < size; i++ { | ||
| pool, err := NewPool(sizePerPool, options...) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| pools[i] = pool | ||
| } | ||
| if lbs != RoundRobin && lbs != LeastTasks { | ||
| return nil, ErrInvalidLoadBalancingStrategy | ||
| } | ||
| return &MultiPool{pools: pools, lbs: lbs}, nil | ||
| } | ||
|
|
||
| func (mp *MultiPool) next() (idx int) { | ||
| switch mp.lbs { | ||
| case RoundRobin: | ||
| if idx = int((atomic.AddUint32(&mp.index, 1) - 1) % uint32(len(mp.pools))); idx == -1 { | ||
| idx = 0 | ||
| } | ||
| return | ||
| case LeastTasks: | ||
| leastTasks := 1<<31 - 1 | ||
| for i, pool := range mp.pools { | ||
| if n := pool.Running(); n < leastTasks { | ||
| leastTasks = n | ||
| idx = i | ||
| } | ||
| } | ||
| return | ||
| } | ||
| return -1 | ||
| } | ||
|
|
||
| // Submit submits a task to a pool selected by the load-balancing strategy. | ||
| func (mp *MultiPool) Submit(task func()) error { | ||
| if mp.IsClosed() { | ||
| return ErrPoolClosed | ||
| } | ||
| return mp.pools[mp.next()].Submit(task) | ||
| } | ||
|
|
||
| // Running returns the number of the currently running workers across all pools. | ||
| func (mp *MultiPool) Running() (n int) { | ||
| for _, pool := range mp.pools { | ||
| n += pool.Running() | ||
| } | ||
| return | ||
| } | ||
|
|
||
| // RunningByIndex returns the number of the currently running workers in the specific pool. | ||
| func (mp *MultiPool) RunningByIndex(idx int) (int, error) { | ||
| if idx < 0 || idx >= len(mp.pools) { | ||
| return -1, ErrInvalidPoolIndex | ||
| } | ||
| return mp.pools[idx].Running(), nil | ||
| } | ||
|
|
||
| // Free returns the number of available workers across all pools. | ||
| func (mp *MultiPool) Free() (n int) { | ||
| for _, pool := range mp.pools { | ||
| n += pool.Free() | ||
| } | ||
| return | ||
| } | ||
|
|
||
| // FreeByIndex returns the number of available workers in the specific pool. | ||
| func (mp *MultiPool) FreeByIndex(idx int) (int, error) { | ||
| if idx < 0 || idx >= len(mp.pools) { | ||
| return -1, ErrInvalidPoolIndex | ||
| } | ||
| return mp.pools[idx].Free(), nil | ||
| } | ||
|
|
||
| // Waiting returns the number of the currently waiting tasks across all pools. | ||
| func (mp *MultiPool) Waiting() (n int) { | ||
| for _, pool := range mp.pools { | ||
| n += pool.Waiting() | ||
| } | ||
| return | ||
| } | ||
|
|
||
| // WaitingByIndex returns the number of the currently waiting tasks in the specific pool. | ||
| func (mp *MultiPool) WaitingByIndex(idx int) (int, error) { | ||
| if idx < 0 || idx >= len(mp.pools) { | ||
| return -1, ErrInvalidPoolIndex | ||
| } | ||
| return mp.pools[idx].Waiting(), nil | ||
| } | ||
|
|
||
| // Cap returns the capacity of this multi-pool. | ||
| func (mp *MultiPool) Cap() (n int) { | ||
| for _, pool := range mp.pools { | ||
| n += pool.Cap() | ||
| } | ||
| return | ||
| } | ||
|
|
||
| // Tune resizes each pool in multi-pool. | ||
| // | ||
| // Note that this method doesn't resize the overall | ||
| // capacity of multi-pool. | ||
| func (mp *MultiPool) Tune(size int) { | ||
| for _, pool := range mp.pools { | ||
| pool.Tune(size) | ||
| } | ||
| } | ||
|
|
||
| // IsClosed indicates whether the multi-pool is closed. | ||
| func (mp *MultiPool) IsClosed() bool { | ||
| return atomic.LoadInt32(&mp.state) == CLOSED | ||
| } | ||
|
|
||
| // ReleaseTimeout closes the multi-pool with a timeout, | ||
| // it waits all pools to be closed before timing out. | ||
| func (mp *MultiPool) ReleaseTimeout(timeout time.Duration) error { | ||
| if !atomic.CompareAndSwapInt32(&mp.state, OPENED, CLOSED) { | ||
| return ErrPoolClosed | ||
| } | ||
|
|
||
| var errStr strings.Builder | ||
| for i, pool := range mp.pools { | ||
| if err := pool.ReleaseTimeout(timeout); err != nil { | ||
| errStr.WriteString(fmt.Sprintf("pool %d: %v\n", i, err)) | ||
| if i < len(mp.pools)-1 { | ||
| errStr.WriteString(" | ") | ||
| } | ||
| return err | ||
| } | ||
| } | ||
|
|
||
| if errStr.Len() == 0 { | ||
| return nil | ||
| } | ||
|
|
||
| return errors.New(errStr.String()) | ||
| } | ||
|
|
||
| // Reboot reboots a released multi-pool. | ||
| func (mp *MultiPool) Reboot() { | ||
| if atomic.CompareAndSwapInt32(&mp.state, CLOSED, OPENED) { | ||
| atomic.StoreUint32(&mp.index, 0) | ||
| for _, pool := range mp.pools { | ||
| pool.Reboot() | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.