github.com/ipni/storetheindex@v0.8.30/rate/map.go (about)

     1  package rate
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  type Rate struct {
     9  	Count   uint64
    10  	Elapsed time.Duration
    11  	Samples int
    12  
    13  	observed bool
    14  }
    15  
    16  // Map records the cumulative rate, since the last call to Get or GetAll,
    17  // for each ID.
    18  type Map struct {
    19  	mutex sync.Mutex
    20  	rates map[string]Rate
    21  }
    22  
    23  func NewMap() *Map {
    24  	return &Map{
    25  		rates: map[string]Rate{},
    26  	}
    27  }
    28  
    29  // Update adds the new rate information to the existing information. If the
    30  // existing rate information is marked as observed, then it is reset before
    31  // applying the new information.
    32  func (m *Map) Update(id string, count uint64, elapsed time.Duration) {
    33  	if m == nil || count == 0 || elapsed == 0 {
    34  		return
    35  	}
    36  	m.mutex.Lock()
    37  	rate := m.rates[id]
    38  	if rate.observed {
    39  		rate = Rate{}
    40  	}
    41  	rate.Count += count
    42  	rate.Elapsed += elapsed
    43  	rate.Samples++
    44  	m.rates[id] = rate
    45  	m.mutex.Unlock()
    46  }
    47  
    48  // Get reads the accumulated rate information for the specified ID. The
    49  // information is marked as observed and is reset at next write.
    50  func (m *Map) Get(id string) (Rate, bool) {
    51  	if m == nil {
    52  		return Rate{}, false
    53  	}
    54  	m.mutex.Lock()
    55  	r, ok := m.rates[id]
    56  	if ok && !r.observed {
    57  		r.observed = true
    58  		m.rates[id] = r
    59  	}
    60  	m.mutex.Unlock()
    61  	return r, ok
    62  }
    63  
    64  // GetAll reads and removes all accumulated rate information. The
    65  // information is marked as observed and is reset at next write.
    66  func (m *Map) GetAll() map[string]Rate {
    67  	if m == nil {
    68  		return nil
    69  	}
    70  	m.mutex.Lock()
    71  	if len(m.rates) == 0 {
    72  		m.mutex.Unlock()
    73  		return nil
    74  	}
    75  	out := make(map[string]Rate, len(m.rates))
    76  	for id, r := range m.rates {
    77  		if !r.observed {
    78  			r.observed = true
    79  			m.rates[id] = r
    80  		}
    81  		out[id] = r
    82  	}
    83  	m.mutex.Unlock()
    84  	return out
    85  }