github.com/clubpay/ronykit/kit@v0.14.4-0.20240515065620-d0dace45cbc7/utils/batch/batcher.go (about)

     1  package batch
     2  
     3  import (
     4  	"sync/atomic"
     5  	"time"
     6  	_ "unsafe"
     7  
     8  	"github.com/clubpay/ronykit/kit/utils"
     9  )
    10  
    11  /*
    12     Creation Time: 2022 - Jul - 22
    13     Created by:  (ehsan)
    14     Maintainers:
    15        1.  Ehsan N. Moosa (E2)
    16     Auditor: Ehsan N. Moosa (E2)
    17  */
    18  
    19  type NA = struct{}
    20  
    21  type Func[IN, OUT any] func(targetID string, entries []Entry[IN, OUT])
    22  
    23  type MultiBatcher[IN, OUT any] struct {
    24  	cfg         config
    25  	batcherFunc Func[IN, OUT]
    26  	poolMtx     utils.SpinLock
    27  	pool        map[string]*Batcher[IN, OUT]
    28  }
    29  
    30  // NewMulti creates a pool of Batcher functions.
    31  // By calling Enter or EnterAndWait you add the item into the Batcher which
    32  // is identified by 'tagID'.
    33  func NewMulti[IN, OUT any](f Func[IN, OUT], opt ...Option) *MultiBatcher[IN, OUT] {
    34  	cfg := defaultConfig
    35  	for _, o := range opt {
    36  		o(&cfg)
    37  	}
    38  
    39  	fp := &MultiBatcher[IN, OUT]{
    40  		cfg:         cfg,
    41  		batcherFunc: f,
    42  		pool:        make(map[string]*Batcher[IN, OUT], 16),
    43  	}
    44  
    45  	return fp
    46  }
    47  
    48  func (fp *MultiBatcher[IN, OUT]) getBatcher(tagID string) *Batcher[IN, OUT] {
    49  	fp.poolMtx.Lock()
    50  	f := fp.pool[tagID]
    51  	if f == nil {
    52  		f = newBatcher[IN, OUT](fp.batcherFunc, tagID, fp.cfg)
    53  		fp.pool[tagID] = f
    54  	}
    55  	fp.poolMtx.Unlock()
    56  
    57  	return f
    58  }
    59  
    60  func (fp *MultiBatcher[IN, OUT]) Enter(targetID string, entry Entry[IN, OUT]) {
    61  	fp.getBatcher(targetID).Enter(entry)
    62  }
    63  
    64  func (fp *MultiBatcher[IN, OUT]) EnterAndWait(targetID string, entry Entry[IN, OUT]) {
    65  	fp.getBatcher(targetID).EnterAndWait(entry)
    66  }
    67  
    68  type Batcher[IN, OUT any] struct {
    69  	utils.SpinLock
    70  
    71  	readyWorkers int32
    72  	batchSize    int32
    73  	minWaitTime  time.Duration
    74  	flusherFunc  Func[IN, OUT]
    75  	entryChan    chan Entry[IN, OUT]
    76  	tagID        string
    77  }
    78  
    79  func NewBatcher[IN, OUT any](f Func[IN, OUT], tagID string, opt ...Option) *Batcher[IN, OUT] {
    80  	cfg := defaultConfig
    81  	for _, o := range opt {
    82  		o(&cfg)
    83  	}
    84  
    85  	return newBatcher[IN, OUT](f, tagID, cfg)
    86  }
    87  
    88  func newBatcher[IN, OUT any](f Func[IN, OUT], tagID string, cfg config) *Batcher[IN, OUT] {
    89  	return &Batcher[IN, OUT]{
    90  		readyWorkers: cfg.maxWorkers,
    91  		batchSize:    cfg.batchSize,
    92  		minWaitTime:  cfg.minWaitTime,
    93  		flusherFunc:  f,
    94  		entryChan:    make(chan Entry[IN, OUT], cfg.batchSize),
    95  		tagID:        tagID,
    96  	}
    97  }
    98  
    99  func (f *Batcher[IN, OUT]) startWorker() {
   100  	f.Lock()
   101  	if atomic.AddInt32(&f.readyWorkers, -1) < 0 {
   102  		atomic.AddInt32(&f.readyWorkers, 1)
   103  		f.Unlock()
   104  
   105  		return
   106  	}
   107  	f.Unlock()
   108  
   109  	w := &worker[IN, OUT]{
   110  		f:  f,
   111  		bs: int(f.batchSize),
   112  	}
   113  
   114  	go w.run()
   115  }
   116  
   117  func (f *Batcher[IN, OUT]) Enter(entry Entry[IN, OUT]) {
   118  	f.entryChan <- entry
   119  	f.startWorker()
   120  }
   121  
   122  func (f *Batcher[IN, OUT]) EnterAndWait(entry Entry[IN, OUT]) {
   123  	f.Enter(entry)
   124  	entry.wait()
   125  }
   126  
   127  type worker[IN, OUT any] struct {
   128  	f  *Batcher[IN, OUT]
   129  	bs int
   130  }
   131  
   132  func (w *worker[IN, OUT]) run() {
   133  	var (
   134  		el        = make([]Entry[IN, OUT], 0, w.bs)
   135  		startTime = utils.NanoTime()
   136  	)
   137  	for {
   138  		for {
   139  			select {
   140  			case e := <-w.f.entryChan:
   141  				el = append(el, e)
   142  				if len(el) < w.bs {
   143  					continue
   144  				}
   145  			default:
   146  			}
   147  
   148  			break
   149  		}
   150  
   151  		if w.f.minWaitTime > 0 && len(el) < w.bs {
   152  			delta := w.f.minWaitTime - time.Duration(utils.NanoTime()-startTime)
   153  			if delta > 0 {
   154  				time.Sleep(delta)
   155  
   156  				continue
   157  			}
   158  		}
   159  		w.f.Lock()
   160  		if len(el) == 0 {
   161  			// clean up and shutdown the worker
   162  			atomic.AddInt32(&w.f.readyWorkers, 1)
   163  			w.f.Unlock()
   164  
   165  			break
   166  		}
   167  		w.f.Unlock()
   168  		w.f.flusherFunc(w.f.tagID, el)
   169  		for idx := range el {
   170  			el[idx].done()
   171  		}
   172  		el = el[:0]
   173  	}
   174  }