github.com/dfklegend/cell2/utils@v0.0.0-20240402033734-a0a9f3d9335d/waterfall/waterfall_sche.go (about)

     1  package waterfall
     2  
     3  import (
     4  	scheP "github.com/dfklegend/cell2/utils/sche"
     5  )
     6  
     7  // 实现一个类似nodejs waterfall
     8  // 每个回调都在调用routine中调用
     9  
    10  type Chain struct {
    11  	// 任务列表
    12  	tasks []Task
    13  	// 最后的任务
    14  	final  FinalCallback
    15  	cursor int
    16  	// 执行
    17  	callbackFunc Callback
    18  }
    19  
    20  func (c *Chain) invokeTask(index int, args ...interface{}) {
    21  	c.tasks[index](c.callbackFunc, args...)
    22  }
    23  
    24  func (c *Chain) tryExec(index int, args ...interface{}) {
    25  	if index < len(c.tasks) {
    26  		c.invokeTask(index, args...)
    27  	} else {
    28  		c.invokeFinal(false, args...)
    29  	}
    30  }
    31  
    32  func (c *Chain) next(args ...interface{}) {
    33  	c.cursor++
    34  	c.tryExec(c.cursor, args...)
    35  }
    36  
    37  func (c *Chain) invokeFinal(err bool, args ...interface{}) {
    38  	c.final(err, args...)
    39  }
    40  
    41  func (c *Chain) invokeCallback(err bool, args ...interface{}) {
    42  	if err {
    43  		c.invokeFinal(true, args...)
    44  		return
    45  	}
    46  
    47  	// next
    48  	c.next(args...)
    49  }
    50  
    51  // 	一般远程任务，传入的回调函数，可能在其他的执行环境内触发
    52  // 	保证Task和Final都在sche内驱动执行
    53  func Sche(sche *scheP.Sche, tasks []Task, final FinalCallback) *Chain {
    54  	thisChain := &Chain{
    55  		tasks: tasks,
    56  		final: final,
    57  	}
    58  
    59  	thisChain.callbackFunc = func(err bool, args ...interface{}) {
    60  		// 去到环境中执行
    61  		sche.Post(func() {
    62  			thisChain.invokeCallback(err, args...)
    63  		})
    64  	}
    65  
    66  	sche.Post(func() {
    67  		thisChain.tryExec(0)
    68  	})
    69  	return thisChain
    70  }
    71  
    72  //	Builder of sche
    73  //	NewBuilder().Next().Next().Final().Do()
    74  type Builder struct {
    75  	sche  *scheP.Sche
    76  	tasks []Task
    77  	final FinalCallback
    78  }
    79  
    80  func NewBuilder(s *scheP.Sche) *Builder {
    81  	return &Builder{
    82  		sche:  s,
    83  		tasks: make([]Task, 0),
    84  	}
    85  }
    86  
    87  func (b *Builder) Next(task Task) *Builder {
    88  	b.tasks = append(b.tasks, task)
    89  	return b
    90  }
    91  
    92  func (b *Builder) Final(final FinalCallback) *Builder {
    93  	b.final = final
    94  	return b
    95  }
    96  
    97  func (b *Builder) Do() {
    98  	Sche(b.sche, b.tasks, b.final)
    99  }