github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/worker/active.go (about)

     1  package worker
     2  
     3  import (
     4  	"github.com/15mga/kiwi/util"
     5  	"time"
     6  
     7  	"github.com/15mga/kiwi/ds"
     8  )
     9  
    10  const (
    11  	cmdActiveCheck   = "check"
    12  	cmdActiveDispose = "dispose"
    13  	cmdActivePush    = "push"
    14  )
    15  
    16  type (
    17  	ActiveOption func(o *activeOption)
    18  	activeOption struct {
    19  		tickSecs int64
    20  		cap      int
    21  	}
    22  )
    23  
    24  // ActiveTickSecs 设置活跃的检测间隔时间
    25  func ActiveTickSecs(seconds int64) ActiveOption {
    26  	return func(opt *activeOption) {
    27  		opt.tickSecs = seconds
    28  	}
    29  }
    30  
    31  // ActiveCap 设置活跃初始容量
    32  func ActiveCap(cap int) ActiveOption {
    33  	return func(opt *activeOption) {
    34  		opt.cap = cap
    35  	}
    36  }
    37  
    38  type activeData struct {
    39  	ts int64
    40  	id string
    41  }
    42  
    43  var (
    44  	_Active *active
    45  )
    46  
    47  func Active() *active {
    48  	return _Active
    49  }
    50  
    51  func InitActive(opts ...ActiveOption) {
    52  	if _Active != nil {
    53  		return
    54  	}
    55  	opt := &activeOption{
    56  		tickSecs: 32,
    57  		cap:      4096,
    58  	}
    59  	for _, o := range opts {
    60  		o(opt)
    61  	}
    62  	_Active = &active{
    63  		option:  opt,
    64  		closeCh: make(chan struct{}, 1),
    65  		activeWorkers: ds.NewKSet[string, *activeWorker](opt.cap, func(a *activeWorker) string {
    66  			return a.id
    67  		}),
    68  		activeTimeStamp: ds.NewKSet[string, *activeData](opt.cap, func(data *activeData) string {
    69  			return data.id
    70  		}),
    71  		activeStopSeconds: opt.tickSecs << 1,
    72  	}
    73  	_Active.worker = NewJobWorker(_Active.process)
    74  	go func() {
    75  		ticker := time.NewTicker(time.Duration(_Active.option.tickSecs) * time.Second)
    76  		defer func() {
    77  			ticker.Stop()
    78  			close(_Active.closeCh)
    79  		}()
    80  		for {
    81  			select {
    82  			case <-ticker.C:
    83  				_Active.worker.Push(cmdActiveCheck)
    84  			case <-_Active.closeCh:
    85  				return
    86  			}
    87  		}
    88  	}()
    89  	_Active.worker.Start()
    90  }
    91  
    92  type active struct {
    93  	option            *activeOption
    94  	worker            *JobWorker
    95  	closeCh           chan struct{}
    96  	activeWorkers     *ds.KSet[string, *activeWorker]
    97  	activeStopSeconds int64
    98  	activeTimeStamp   *ds.KSet[string, *activeData]
    99  }
   100  
   101  func (a *active) Dispose() {
   102  	a.worker.Push(cmdActiveDispose)
   103  	a.worker.Dispose()
   104  }
   105  
   106  func (a *active) Push(id string, fn util.FnAnySlc, params ...any) {
   107  	a.worker.Push(cmdActivePush, id, fn, params)
   108  }
   109  
   110  func (a *active) process(job *Job) {
   111  	switch job.Name {
   112  	case cmdActiveCheck:
   113  		now := time.Now().Unix()
   114  		//移除不活跃的协程
   115  		a.activeTimeStamp.TestDel(func(id string, item *activeData) (del bool, brk bool) {
   116  			if now-item.ts < a.activeStopSeconds {
   117  				return
   118  			}
   119  			obj, _ := a.activeWorkers.Del(id)
   120  			obj.Dispose()
   121  			del = true
   122  			return
   123  		})
   124  	case cmdActiveDispose:
   125  		a.activeWorkers.Iter(func(item *activeWorker) {
   126  			item.Dispose()
   127  		})
   128  	case cmdActivePush:
   129  		id, fn, params := util.SplitSlc3[string, util.FnAnySlc, []any](job.Data)
   130  		worker, ok := a.activeWorkers.Get(id)
   131  		now := time.Now().Unix()
   132  		if ok {
   133  			d, _ := a.activeTimeStamp.Get(id)
   134  			d.ts = now
   135  		} else {
   136  			worker = newActiveWorker(id)
   137  			worker.Start()
   138  			a.activeWorkers.Set(worker)
   139  			_ = a.activeTimeStamp.Add(&activeData{
   140  				ts: now,
   141  				id: id,
   142  			})
   143  		}
   144  		worker.Push(fn, params...)
   145  	}
   146  }
   147  
   148  func newActiveWorker(id string) *activeWorker {
   149  	a := &activeWorker{
   150  		FnWorker: NewFnWorker(),
   151  		id:       id,
   152  	}
   153  	return a
   154  }
   155  
   156  type activeWorker struct {
   157  	*FnWorker
   158  	id string
   159  }