tstor/pkg/kvtrace/kvmetrics.go

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)