91 lines
2.2 KiB
Go
91 lines
2.2 KiB
Go
|
package kvtrace
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
|
||
|
|
"github.com/royalcat/kv"
|
||
|
|
"go.opentelemetry.io/otel"
|
||
|
|
"go.opentelemetry.io/otel/attribute"
|
||
|
|
"go.opentelemetry.io/otel/trace"
|
||
|
|
)
|
||
|
|
|
||
|
|
var tracer = otel.Tracer("github.com/royalcat/kv/tracer")
|
||
|
|
|
||
|
|
type traceSrtore[K, V any] struct {
|
||
|
|
kv kv.Store[K, V]
|
||
|
|
attrs []attribute.KeyValue
|
||
|
|
}
|
||
|
|
|
||
|
|
func WrapTracing[K, V any](kv kv.Store[K, V], attrs ...attribute.KeyValue) kv.Store[K, V] {
|
||
|
|
return &traceSrtore[K, V]{
|
||
|
|
kv: kv,
|
||
|
|
attrs: attrs,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Close implements kv.Store.
|
||
|
|
func (m *traceSrtore[K, V]) Close(ctx context.Context) error {
|
||
|
|
ctx, span := tracer.Start(ctx, "Close", trace.WithAttributes(m.attrs...))
|
||
|
|
defer span.End()
|
||
|
|
|
||
|
|
return m.kv.Close(ctx)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Delete implements kv.Store.
|
||
|
|
func (m *traceSrtore[K, V]) Delete(ctx context.Context, k K) error {
|
||
|
|
ctx, span := tracer.Start(ctx, "Delete", trace.WithAttributes(m.attrs...))
|
||
|
|
defer span.End()
|
||
|
|
|
||
|
|
return m.kv.Delete(ctx, k)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Get implements kv.Store.
|
||
|
|
func (m *traceSrtore[K, V]) Get(ctx context.Context, k K) (v V, found bool, err error) {
|
||
|
|
ctx, span := tracer.Start(ctx, "Get", trace.WithAttributes(m.attrs...))
|
||
|
|
defer span.End()
|
||
|
|
|
||
|
|
return m.kv.Get(ctx, k)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Range implements kv.Store.
|
||
|
|
func (m *traceSrtore[K, V]) Range(ctx context.Context, iter kv.Iter[K, V]) error {
|
||
|
|
ctx, span := tracer.Start(ctx, "Range", trace.WithAttributes(m.attrs...))
|
||
|
|
defer span.End()
|
||
|
|
|
||
|
|
count := 0
|
||
|
|
iterCount := func(k K, v V) bool {
|
||
|
|
count++
|
||
|
|
return iter(k, v)
|
||
|
|
}
|
||
|
|
|
||
|
|
err := m.kv.Range(ctx, iterCount)
|
||
|
|
span.SetAttributes(attribute.Int("count", count))
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
// RangeWithPrefix implements kv.Store.
|
||
|
|
func (m *traceSrtore[K, V]) RangeWithPrefix(ctx context.Context, k K, iter kv.Iter[K, V]) error {
|
||
|
|
ctx, span := tracer.Start(ctx, "RangeWithPrefix", trace.WithAttributes(m.attrs...))
|
||
|
|
defer span.End()
|
||
|
|
|
||
|
|
count := 0
|
||
|
|
iterCount := func(k K, v V) bool {
|
||
|
|
count++
|
||
|
|
return iter(k, v)
|
||
|
|
}
|
||
|
|
|
||
|
|
err := m.kv.Range(ctx, iterCount)
|
||
|
|
span.SetAttributes(attribute.Int("count", count))
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
// Set implements kv.Store.
|
||
|
|
func (m *traceSrtore[K, V]) Set(ctx context.Context, k K, v V) error {
|
||
|
|
ctx, span := tracer.Start(ctx, "Set", trace.WithAttributes(m.attrs...))
|
||
|
|
defer span.End()
|
||
|
|
|
||
|
|
return m.kv.Set(ctx, k, v)
|
||
|
|
}
|
||
|
|
|
||
|
|
var _ kv.Store[any, any] = (*traceSrtore[any, any])(nil)
|