5cc68e27dc
uses the count bus in the kopia backup package. This currently duplicates counts that we're getting from the kopia stats. A later pr will remove the old stats entirely in favor of the counter. --- #### Does this PR need a docs update or release note? - [x] ⛔ No #### Type of change - [x] 🧹 Tech Debt/Cleanup #### Test Plan - [x] ⚡ Unit test - [x] 💚 E2E
149 lines
2.7 KiB
Go
149 lines
2.7 KiB
Go
package count
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/puzpuzpuz/xsync/v2"
|
|
)
|
|
|
|
// Bus handles threadsafe counting of arbitrarily keyed metrics.
|
|
type Bus struct {
|
|
parent *Bus
|
|
stats *xsync.MapOf[string, *xsync.Counter]
|
|
}
|
|
|
|
func New() *Bus {
|
|
return &Bus{
|
|
stats: xsync.NewMapOf[*xsync.Counter](),
|
|
}
|
|
}
|
|
|
|
// Local generates a bus with a parent link. Any value added to
|
|
// the local instance also updates the parent by the same increment.
|
|
// This allows you to maintain an isolated set of counts for a
|
|
// bounded context while automatically tallying the global total.
|
|
func (b *Bus) Local() *Bus {
|
|
bus := New()
|
|
bus.parent = b
|
|
|
|
return bus
|
|
}
|
|
|
|
func (b *Bus) getCounter(k key) *xsync.Counter {
|
|
xc, _ := b.stats.LoadOrStore(string(k), xsync.NewCounter())
|
|
return xc
|
|
}
|
|
|
|
// Inc increases the count by 1.
|
|
func (b *Bus) Inc(k key) {
|
|
if b == nil {
|
|
return
|
|
}
|
|
|
|
b.Add(k, 1)
|
|
}
|
|
|
|
// Inc increases the count by n.
|
|
func (b *Bus) Add(k key, n int64) {
|
|
if b == nil {
|
|
return
|
|
}
|
|
|
|
b.getCounter(k).Add(n)
|
|
|
|
if b.parent != nil {
|
|
b.parent.Add(k, n)
|
|
}
|
|
}
|
|
|
|
// AdderFor returns a func that adds any value of i
|
|
// to the bus using the given key.
|
|
func (b *Bus) AdderFor(k key) func(i int64) {
|
|
return func(i int64) {
|
|
b.Add(k, i)
|
|
}
|
|
}
|
|
|
|
// Get returns the local count.
|
|
func (b *Bus) Get(k key) int64 {
|
|
if b == nil {
|
|
return -1
|
|
}
|
|
|
|
return b.getCounter(k).Value()
|
|
}
|
|
|
|
// Total returns the global count.
|
|
func (b *Bus) Total(k key) int64 {
|
|
if b == nil {
|
|
return -1
|
|
}
|
|
|
|
if b.parent != nil {
|
|
return b.parent.Total(k)
|
|
}
|
|
|
|
return b.Get(k)
|
|
}
|
|
|
|
// Values returns a map of all local values.
|
|
// Not a snapshot, and therefore not threadsafe.
|
|
func (b *Bus) Values() map[string]int64 {
|
|
if b == nil {
|
|
return map[string]int64{}
|
|
}
|
|
|
|
m := make(map[string]int64, b.stats.Size())
|
|
|
|
b.stats.Range(func(k string, v *xsync.Counter) bool {
|
|
m[k] = v.Value()
|
|
return true
|
|
})
|
|
|
|
return m
|
|
}
|
|
|
|
// TotalValues returns a map of all global values.
|
|
// Not a snapshot, and therefore not threadsafe.
|
|
func (b *Bus) TotalValues() map[string]int64 {
|
|
if b == nil {
|
|
return map[string]int64{}
|
|
}
|
|
|
|
if b.parent != nil {
|
|
return b.parent.TotalValues()
|
|
}
|
|
|
|
return b.Values()
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// context embedding
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type countKey string
|
|
|
|
const ctxKey countKey = "corsoCounter"
|
|
|
|
// Ctx retrieves the count.Bus embedded in the context.
|
|
func Ctx(ctx context.Context) *Bus {
|
|
l := ctx.Value(ctxKey)
|
|
if l == nil {
|
|
return New()
|
|
}
|
|
|
|
return l.(*Bus)
|
|
}
|
|
|
|
// Embed allows users to embed their own count.Bus within the context.
|
|
func Embed(
|
|
ctx context.Context,
|
|
bus *Bus,
|
|
) context.Context {
|
|
if bus == nil {
|
|
bus = New()
|
|
}
|
|
|
|
return context.WithValue(ctx, ctxKey, bus)
|
|
}
|