github.com/go-board/x-go@v0.1.2-0.20220610024734-db1323f6cb15/concurrent/executor.go (about)

     1  package concurrent
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"runtime"
     7  	"time"
     8  )
     9  
    10  type TaskHandle struct {
    11  	ctx    context.Context
    12  	cancel context.CancelFunc
    13  
    14  	fn   func(ctx context.Context) error
    15  	done chan error
    16  }
    17  
    18  func (t *TaskHandle) Cancel() {
    19  	if t.cancel != nil {
    20  		t.cancel()
    21  	}
    22  }
    23  
    24  func (t *TaskHandle) Wait() error {
    25  	return <-t.done
    26  }
    27  
    28  func (t *TaskHandle) WaitTimeout(timeout time.Duration) error {
    29  	timer := time.NewTimer(timeout)
    30  	defer timer.Stop()
    31  	select {
    32  	case <-timer.C:
    33  		return errors.New("context deadline exceeded")
    34  	case err := <-t.done:
    35  		return err
    36  	}
    37  }
    38  
    39  type Pool struct {
    40  	// TODO: 使用无锁数据结构来提高性能
    41  	taskCh    chan *TaskHandle
    42  	workerNum int
    43  	done      chan struct{}
    44  }
    45  
    46  func (p *Pool) poll() {
    47  	for i := 0; i < p.workerNum; i++ {
    48  		go func() {
    49  			for {
    50  				select {
    51  				case <-p.done:
    52  					return
    53  				case task := <-p.taskCh:
    54  					select {
    55  					case <-task.ctx.Done():
    56  						task.done <- task.ctx.Err()
    57  					default:
    58  						task.done <- task.fn(task.ctx)
    59  					}
    60  				}
    61  			}
    62  		}()
    63  	}
    64  }
    65  
    66  func (p *Pool) JobQueueSize() int {
    67  	return len(p.taskCh)
    68  }
    69  
    70  func (p *Pool) Spawn(ctx context.Context, fn func(ctx context.Context) error) *TaskHandle {
    71  	ctx, cancel := context.WithCancel(ctx)
    72  	t := &TaskHandle{
    73  		ctx:    ctx,
    74  		cancel: cancel,
    75  		fn:     fn,
    76  		done:   make(chan error, 1),
    77  	}
    78  	p.taskCh <- t
    79  	return t
    80  }
    81  
    82  func (p *Pool) Shutdown() {
    83  	close(p.done)
    84  }
    85  
    86  func NewPool(workerNum int, taskNum int) *Pool {
    87  	pool := &Pool{
    88  		taskCh:    make(chan *TaskHandle, taskNum),
    89  		workerNum: workerNum,
    90  		done:      make(chan struct{}),
    91  	}
    92  	pool.poll()
    93  	return pool
    94  }
    95  
    96  var globalPool = NewPool(runtime.NumCPU(), 10000)
    97  
    98  func Spawn(ctx context.Context, fn func(ctx context.Context) error) *TaskHandle {
    99  	return globalPool.Spawn(ctx, fn)
   100  }