gitee.com/sy_183/go-common@v1.0.5-0.20231205030221-958cfe129b47/lifecycle/task/task-executor.go (about)

     1  package task
     2  
     3  import (
     4  	"gitee.com/sy_183/go-common/chans"
     5  	"gitee.com/sy_183/go-common/errors"
     6  	"gitee.com/sy_183/go-common/lifecycle"
     7  	"sync"
     8  	"sync/atomic"
     9  )
    10  
    11  var TaskExecutorClosedError = errors.New("任务执行器已关闭")
    12  
    13  type TaskExecutor struct {
    14  	lifecycle.Lifecycle
    15  	maxTask int
    16  	ch      atomic.Pointer[chan Task]
    17  }
    18  
    19  func NewTaskExecutor(maxTask int) *TaskExecutor {
    20  	executor := &TaskExecutor{maxTask: maxTask}
    21  	executor.Lifecycle = lifecycle.NewWithInterruptedRun(executor.start, executor.run)
    22  	return executor
    23  }
    24  
    25  func (e *TaskExecutor) loadChannel() chan Task {
    26  	if ch := e.ch.Load(); ch != nil {
    27  		return *ch
    28  	}
    29  	return nil
    30  }
    31  
    32  func (e *TaskExecutor) start(lifecycle.Lifecycle, chan struct{}) error {
    33  	ch := make(chan Task, e.maxTask)
    34  	e.ch.Store(&ch)
    35  	return nil
    36  }
    37  
    38  func (e *TaskExecutor) run(_ lifecycle.Lifecycle, interrupter chan struct{}) error {
    39  	ch := e.loadChannel()
    40  	defer func() {
    41  		e.ch.Store(nil)
    42  		close(ch)
    43  		for task := range ch {
    44  			task.Do(nil)
    45  		}
    46  	}()
    47  	for {
    48  		select {
    49  		case task := <-ch:
    50  			if task.Do(interrupter) {
    51  				return nil
    52  			}
    53  		case <-interrupter:
    54  			return nil
    55  		}
    56  	}
    57  }
    58  
    59  func (e *TaskExecutor) StartFunc() lifecycle.InterruptedStartFunc {
    60  	return e.start
    61  }
    62  
    63  func (e *TaskExecutor) RunFunc() lifecycle.InterruptedRunFunc {
    64  	return e.run
    65  }
    66  
    67  func (e *TaskExecutor) notPanicPush(ch chan Task, task Task) (err error) {
    68  	defer func() {
    69  		if _, is := recover().(error); is {
    70  			err = TaskExecutorClosedError
    71  		}
    72  	}()
    73  	ch <- task
    74  	return nil
    75  }
    76  
    77  func (e *TaskExecutor) notPanicTryPush(ch chan Task, task Task) (ok bool, err error) {
    78  	defer func() {
    79  		if _, is := recover().(error); is {
    80  			err = TaskExecutorClosedError
    81  		}
    82  	}()
    83  	return chans.TryPush(ch, task), nil
    84  }
    85  
    86  func (e *TaskExecutor) Async(task Task) error {
    87  	ch := e.loadChannel()
    88  	if ch == nil {
    89  		return TaskExecutorClosedError
    90  	}
    91  	return e.notPanicPush(ch, task)
    92  }
    93  
    94  func (e *TaskExecutor) Sync(task Task) error {
    95  	ch := e.loadChannel()
    96  	if ch == nil {
    97  		return TaskExecutorClosedError
    98  	}
    99  	waiter := sync.WaitGroup{}
   100  	waiter.Add(1)
   101  	if err := e.notPanicPush(ch, Interrupted(func(interrupter chan struct{}) (interrupted bool) {
   102  		defer waiter.Done()
   103  		return task.Do(interrupter)
   104  	})); err != nil {
   105  		return err
   106  	}
   107  	waiter.Wait()
   108  	return nil
   109  }
   110  
   111  func (e *TaskExecutor) Try(task Task) (ok bool, err error) {
   112  	ch := e.loadChannel()
   113  	if ch == nil {
   114  		return false, TaskExecutorClosedError
   115  	}
   116  	return e.notPanicTryPush(ch, task)
   117  }
   118  
   119  func (e *TaskExecutor) Wait() error {
   120  	return e.Sync(Nop())
   121  }