github.com/ronaksoft/rony@v0.16.26-0.20230807065236-1743dbfe6959/tools/flusher.go (about)

     1  package tools
     2  
     3  import (
     4  	"sync/atomic"
     5  	"time"
     6  )
     7  
     8  /*
     9     Creation Time: 2020 - Dec - 31
    10     Created by:  (ehsan)
    11     Maintainers:
    12        1.  Ehsan N. Moosa (E2)
    13     Auditor: Ehsan N. Moosa (E2)
    14     Copyright Ronak Software Group 2020
    15  */
    16  
    17  type FlushEntry interface {
    18  	wait()
    19  	done()
    20  	Value() interface{}
    21  }
    22  
    23  type entry struct {
    24  	v  interface{}
    25  	ch chan struct{}
    26  	cb func()
    27  }
    28  
    29  func NewEntry(v interface{}) FlushEntry {
    30  	return &entry{
    31  		v:  v,
    32  		ch: make(chan struct{}, 1),
    33  	}
    34  }
    35  
    36  func NewEntryWithCallback(v interface{}, cb func()) FlushEntry {
    37  	return &entry{
    38  		v:  v,
    39  		ch: make(chan struct{}, 1),
    40  		cb: cb,
    41  	}
    42  }
    43  
    44  func (e *entry) wait() {
    45  	<-e.ch
    46  }
    47  
    48  func (e *entry) done() {
    49  	if e.cb != nil {
    50  		e.cb()
    51  	}
    52  	e.ch <- struct{}{}
    53  }
    54  
    55  func (e *entry) Value() interface{} {
    56  	return e.v
    57  }
    58  
    59  type FlusherFunc func(targetID string, entries []FlushEntry)
    60  
    61  type FlusherPool struct {
    62  	maxWorkers  int32
    63  	batchSize   int32
    64  	minWaitTime time.Duration
    65  	flusherFunc FlusherFunc
    66  	poolMtx     SpinLock
    67  	pool        map[string]*flusher
    68  }
    69  
    70  // NewFlusherPool creates a pool of flusher funcs. By calling Enter or EnterAndWait you add
    71  // the item into the flusher which identified by 'targetID'.
    72  func NewFlusherPool(maxWorkers, batchSize int32, f FlusherFunc) *FlusherPool {
    73  	return NewFlusherPoolWithWaitTime(maxWorkers, batchSize, 0, f)
    74  }
    75  
    76  func NewFlusherPoolWithWaitTime(maxWorkers, batchSize int32, minWaitTime time.Duration, f FlusherFunc) *FlusherPool {
    77  	fp := &FlusherPool{
    78  		maxWorkers:  maxWorkers,
    79  		batchSize:   batchSize,
    80  		minWaitTime: minWaitTime,
    81  		flusherFunc: f,
    82  		pool:        make(map[string]*flusher, 16),
    83  	}
    84  
    85  	return fp
    86  }
    87  
    88  func (fp *FlusherPool) getFlusher(targetID string) *flusher {
    89  	fp.poolMtx.Lock()
    90  	f := fp.pool[targetID]
    91  	if f == nil {
    92  		f = &flusher{
    93  			readyWorkers: fp.maxWorkers,
    94  			batchSize:    fp.batchSize,
    95  			minWaitTime:  fp.minWaitTime,
    96  			flusherFunc:  fp.flusherFunc,
    97  			entryChan:    make(chan FlushEntry, fp.batchSize),
    98  			targetID:     targetID,
    99  		}
   100  		fp.pool[targetID] = f
   101  	}
   102  	fp.poolMtx.Unlock()
   103  
   104  	return f
   105  }
   106  
   107  func (fp *FlusherPool) Enter(targetID string, entry FlushEntry) {
   108  	fp.getFlusher(targetID).enter(entry)
   109  }
   110  
   111  func (fp *FlusherPool) EnterAndWait(targetID string, entry FlushEntry) {
   112  	fp.getFlusher(targetID).enterAndWait(entry)
   113  }
   114  
   115  type flusher struct {
   116  	SpinLock
   117  	readyWorkers int32
   118  	batchSize    int32
   119  	minWaitTime  time.Duration
   120  	flusherFunc  FlusherFunc
   121  	entryChan    chan FlushEntry
   122  	targetID     string
   123  }
   124  
   125  func (f *flusher) startWorker() {
   126  	f.Lock()
   127  	if atomic.AddInt32(&f.readyWorkers, -1) < 0 {
   128  		atomic.AddInt32(&f.readyWorkers, 1)
   129  		f.Unlock()
   130  
   131  		return
   132  	}
   133  	f.Unlock()
   134  
   135  	w := &worker{
   136  		f:  f,
   137  		bs: int(f.batchSize),
   138  	}
   139  	go w.run()
   140  }
   141  
   142  func (f *flusher) enter(entry FlushEntry) {
   143  	f.entryChan <- entry
   144  	f.startWorker()
   145  }
   146  
   147  func (f *flusher) enterAndWait(entry FlushEntry) {
   148  	f.entryChan <- entry
   149  	f.startWorker()
   150  	entry.wait()
   151  }
   152  
   153  type worker struct {
   154  	f  *flusher
   155  	bs int
   156  }
   157  
   158  func (w *worker) run() {
   159  	var (
   160  		el        = make([]FlushEntry, 0, w.bs)
   161  		startTime = NanoTime()
   162  	)
   163  	for {
   164  		for {
   165  			select {
   166  			case e := <-w.f.entryChan:
   167  				el = append(el, e)
   168  				if len(el) < w.bs {
   169  					continue
   170  				}
   171  			default:
   172  			}
   173  
   174  			break
   175  		}
   176  
   177  		if w.f.minWaitTime > 0 {
   178  			delta := w.f.minWaitTime - time.Duration(NanoTime()-startTime)
   179  			if delta > 0 {
   180  				time.Sleep(delta)
   181  
   182  				continue
   183  			}
   184  		}
   185  		w.f.Lock()
   186  		if len(el) == 0 {
   187  			// clean up and shutdown the worker
   188  			atomic.AddInt32(&w.f.readyWorkers, 1)
   189  			w.f.Unlock()
   190  
   191  			break
   192  		}
   193  		w.f.Unlock()
   194  		w.f.flusherFunc(w.f.targetID, el)
   195  		for idx := range el {
   196  			el[idx].done()
   197  		}
   198  		el = el[:0]
   199  	}
   200  }