github.com/coyove/sdss@v0.0.0-20231129015646-c2ec58cca6a2/contrib/bitmap/bitmap_agg.go (about)

     1  package bitmap
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  
     7  	"github.com/coyove/sdss/contrib/clock"
     8  )
     9  
    10  var ErrBitmapFull = fmt.Errorf("bitmap full (%d)", Capcity)
    11  
    12  type aggTask struct {
    13  	key    Key
    14  	values []uint64
    15  	out    chan error
    16  }
    17  
    18  type SaveAggregator struct {
    19  	cb        func(*Range) error
    20  	tasks     chan *aggTask
    21  	workerOut chan bool
    22  	current   *Range
    23  
    24  	survey struct {
    25  		c, r int
    26  	}
    27  
    28  	window time.Duration
    29  }
    30  
    31  func (r *Range) AggregateSaves(callback func(*Range) error) *SaveAggregator {
    32  	fts := &SaveAggregator{}
    33  	fts.tasks = make(chan *aggTask, 10000)
    34  	fts.workerOut = make(chan bool, 1)
    35  	fts.cb = callback
    36  	fts.current = r
    37  	fts.window = 100 * time.Millisecond
    38  
    39  	go func() {
    40  		for fts.worker() {
    41  		}
    42  		fts.workerOut <- true
    43  	}()
    44  	return fts
    45  }
    46  
    47  func (sa *SaveAggregator) SetWindow(w time.Duration) *SaveAggregator {
    48  	sa.window = w
    49  	return sa
    50  }
    51  
    52  func (sa *SaveAggregator) Range() *Range {
    53  	return sa.current
    54  }
    55  
    56  func (sa *SaveAggregator) Close() {
    57  	close(sa.tasks)
    58  	<-sa.workerOut
    59  }
    60  
    61  func (sa *SaveAggregator) worker() bool {
    62  	start := clock.UnixNano()
    63  	tm := time.NewTimer(sa.window)
    64  
    65  	var tasks []*aggTask
    66  
    67  MORE:
    68  	select {
    69  	case t, ok := <-sa.tasks:
    70  		if !ok {
    71  			if len(tasks) == 0 {
    72  				return false
    73  			}
    74  		} else {
    75  			tasks = append(tasks, t)
    76  			if clock.UnixNano()-start < sa.window.Nanoseconds() {
    77  				goto MORE
    78  			}
    79  		}
    80  	case <-tm.C:
    81  	}
    82  
    83  	tm.Stop()
    84  	if len(tasks) == 0 {
    85  		return true
    86  	}
    87  
    88  	sa.survey.c += len(tasks)
    89  	sa.survey.r += 1
    90  	if sa.survey.r > 100 {
    91  		sa.survey.c = int(sa.Metrics())
    92  		sa.survey.r = 1
    93  	}
    94  
    95  	for i, t := range tasks {
    96  		if !sa.current.Add(t.key, t.values) {
    97  			for j := i; j < len(tasks); j++ {
    98  				tasks[j].out <- ErrBitmapFull
    99  			}
   100  			tasks = tasks[:i]
   101  			break
   102  		}
   103  	}
   104  
   105  	err := sa.cb(sa.current)
   106  	for _, t := range tasks {
   107  		t.out <- err
   108  	}
   109  	return true
   110  }
   111  
   112  func (sa *SaveAggregator) AddAsync(key Key, values []uint64) chan error {
   113  	t := &aggTask{
   114  		key:    key,
   115  		values: values,
   116  		out:    make(chan error, 1),
   117  	}
   118  	sa.tasks <- t
   119  	return t.out
   120  }
   121  
   122  func (sa *SaveAggregator) Add(key Key, values []uint64) error {
   123  	return <-sa.AddAsync(key, values)
   124  }
   125  
   126  func (sa *SaveAggregator) Metrics() float64 {
   127  	return float64(sa.survey.c) / float64(sa.survey.r)
   128  }