small refactor*

server/pkg/kvtrace/kvmetrics.go

@@ -0,0 +1,150 @@
+package kvtrace
+
+import (
+	"context"
+
+	"github.com/royalcat/kv"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/trace"
+)
+
+var tracer = otel.Tracer("github.com/royalcat/kv/tracer")
+
+type traceStore[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 &traceStore[K, V]{
+		kv:    kv,
+		attrs: attrs,
+	}
+}
+
+// Close implements kv.Store.
+func (m *traceStore[K, V]) Close(ctx context.Context) (err error) {
+	ctx, span := tracer.Start(ctx, "Close", trace.WithAttributes(m.attrs...))
+	defer span.End()
+	defer func() {
+		if err != nil && err != kv.ErrKeyNotFound {
+			span.SetStatus(codes.Error, err.Error())
+		} else {
+			span.SetStatus(codes.Ok, "")
+		}
+	}()
+
+	return m.kv.Close(ctx)
+}
+
+// Delete implements kv.Store.
+func (m *traceStore[K, V]) Delete(ctx context.Context, k K) (err error) {
+	ctx, span := tracer.Start(ctx, "Delete", trace.WithAttributes(m.attrs...))
+	defer span.End()
+	defer func() {
+		if err != nil && err != kv.ErrKeyNotFound {
+			span.SetStatus(codes.Error, err.Error())
+		} else {
+			span.SetStatus(codes.Ok, "")
+		}
+	}()
+
+	return m.kv.Delete(ctx, k)
+}
+
+// Get implements kv.Store.
+func (m *traceStore[K, V]) Get(ctx context.Context, k K) (v V, err error) {
+	ctx, span := tracer.Start(ctx, "Get", trace.WithAttributes(m.attrs...))
+	defer span.End()
+	defer func() {
+		if err != nil && err != kv.ErrKeyNotFound {
+			span.SetStatus(codes.Error, err.Error())
+		} else {
+			span.SetStatus(codes.Ok, "")
+		}
+	}()
+
+	return m.kv.Get(ctx, k)
+}
+
+// Get implements kv.Store.
+func (m *traceStore[K, V]) Edit(ctx context.Context, k K, edit kv.Edit[V]) (err error) {
+	ctx, span := tracer.Start(ctx, "Get", trace.WithAttributes(m.attrs...))
+	defer span.End()
+	defer func() {
+		if err != nil && err != kv.ErrKeyNotFound {
+			span.SetStatus(codes.Error, err.Error())
+		} else {
+			span.SetStatus(codes.Ok, "")
+		}
+	}()
+
+	return m.kv.Edit(ctx, k, edit)
+}
+
+// Range implements kv.Store.
+func (m *traceStore[K, V]) Range(ctx context.Context, iter kv.Iter[K, V]) (err error) {
+	ctx, span := tracer.Start(ctx, "Range", trace.WithAttributes(m.attrs...))
+	defer span.End()
+	defer func() {
+		if err != nil {
+			span.SetStatus(codes.Error, err.Error())
+		} else {
+			span.SetStatus(codes.Ok, "")
+		}
+	}()
+
+	count := 0
+	iterCount := func(k K, v V) error {
+		count++
+		return iter(k, v)
+	}
+	defer func() {
+		span.SetAttributes(attribute.Int("count", count))
+	}()
+
+	return m.kv.Range(ctx, iterCount)
+}
+
+// RangeWithPrefix implements kv.Store.
+func (m *traceStore[K, V]) RangeWithPrefix(ctx context.Context, k K, iter kv.Iter[K, V]) (err error) {
+	ctx, span := tracer.Start(ctx, "RangeWithPrefix", trace.WithAttributes(m.attrs...))
+	defer span.End()
+	defer func() {
+		if err != nil {
+			span.SetStatus(codes.Error, err.Error())
+		} else {
+			span.SetStatus(codes.Ok, "")
+		}
+	}()
+
+	count := 0
+	iterCount := func(k K, v V) error {
+		count++
+		return iter(k, v)
+	}
+	defer func() {
+		span.SetAttributes(attribute.Int("count", count))
+	}()
+
+	return m.kv.Range(ctx, iterCount)
+}
+
+// Set implements kv.Store.
+func (m *traceStore[K, V]) Set(ctx context.Context, k K, v V) (err error) {
+	ctx, span := tracer.Start(ctx, "Set", trace.WithAttributes(m.attrs...))
+	defer span.End()
+	defer func() {
+		if err != nil {
+			span.SetStatus(codes.Error, err.Error())
+		} else {
+			span.SetStatus(codes.Ok, "")
+		}
+	}()
+
+	return m.kv.Set(ctx, k, v)
+}
+
+var _ kv.Store[any, any] = (*traceStore[any, any])(nil)