github.com/netdata/go.d.plugin@v0.58.1/pkg/metrics/unique_counter.go (about)

     1  // SPDX-License-Identifier: GPL-3.0-or-later
     2  
     3  package metrics
     4  
     5  import (
     6  	"github.com/axiomhq/hyperloglog"
     7  	"github.com/netdata/go.d.plugin/pkg/stm"
     8  )
     9  
    10  type (
    11  	UniqueCounter interface {
    12  		stm.Value
    13  		Insert(s string)
    14  		Value() int
    15  		Reset()
    16  	}
    17  
    18  	mapUniqueCounter struct {
    19  		m map[string]bool
    20  	}
    21  
    22  	hyperLogLogUniqueCounter struct {
    23  		sketch *hyperloglog.Sketch
    24  	}
    25  
    26  	UniqueCounterVec struct {
    27  		useHyperLogLog bool
    28  		Items          map[string]UniqueCounter
    29  	}
    30  )
    31  
    32  var (
    33  	_ stm.Value = mapUniqueCounter{}
    34  	_ stm.Value = hyperLogLogUniqueCounter{}
    35  	_ stm.Value = UniqueCounterVec{}
    36  )
    37  
    38  func NewUniqueCounter(useHyperLogLog bool) UniqueCounter {
    39  	if useHyperLogLog {
    40  		return &hyperLogLogUniqueCounter{hyperloglog.New()}
    41  	}
    42  	return mapUniqueCounter{map[string]bool{}}
    43  }
    44  
    45  func (c mapUniqueCounter) WriteTo(rv map[string]int64, key string, mul, div int) {
    46  	rv[key] = int64(float64(c.Value()*mul) / float64(div))
    47  }
    48  
    49  func (c mapUniqueCounter) Insert(s string) {
    50  	c.m[s] = true
    51  }
    52  
    53  func (c mapUniqueCounter) Value() int {
    54  	return len(c.m)
    55  }
    56  
    57  func (c mapUniqueCounter) Reset() {
    58  	for key := range c.m {
    59  		delete(c.m, key)
    60  	}
    61  }
    62  
    63  // WriteTo writes its value into given map.
    64  func (c hyperLogLogUniqueCounter) WriteTo(rv map[string]int64, key string, mul, div int) {
    65  	rv[key] = int64(float64(c.Value()*mul) / float64(div))
    66  }
    67  
    68  func (c *hyperLogLogUniqueCounter) Insert(s string) {
    69  	c.sketch.Insert([]byte(s))
    70  }
    71  
    72  func (c *hyperLogLogUniqueCounter) Value() int {
    73  	return int(c.sketch.Estimate())
    74  }
    75  
    76  func (c *hyperLogLogUniqueCounter) Reset() {
    77  	c.sketch = hyperloglog.New()
    78  }
    79  
    80  func NewUniqueCounterVec(useHyperLogLog bool) UniqueCounterVec {
    81  	return UniqueCounterVec{
    82  		Items:          map[string]UniqueCounter{},
    83  		useHyperLogLog: useHyperLogLog,
    84  	}
    85  }
    86  
    87  // WriteTo writes its value into given map.
    88  func (c UniqueCounterVec) WriteTo(rv map[string]int64, key string, mul, div int) {
    89  	for name, value := range c.Items {
    90  		value.WriteTo(rv, key+"_"+name, mul, div)
    91  	}
    92  }
    93  
    94  // Get gets UniqueCounter instance by name
    95  func (c UniqueCounterVec) Get(name string) UniqueCounter {
    96  	item, ok := c.Items[name]
    97  	if ok {
    98  		return item
    99  	}
   100  	item = NewUniqueCounter(c.useHyperLogLog)
   101  	c.Items[name] = item
   102  	return item
   103  }
   104  
   105  func (c UniqueCounterVec) Reset() {
   106  	for _, value := range c.Items {
   107  		value.Reset()
   108  	}
   109  }