update

pkg/kvtrace/kvmetrics.go

@@ -0,0 +1,90 @@
+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)