79 lines
2.3 KiB
Go
79 lines
2.3 KiB
Go
package telemetry
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"github.com/redis/go-redis/v9"
|
|
"go.opentelemetry.io/otel/codes"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
)
|
|
|
|
type spanBody struct {
|
|
io.ReadCloser
|
|
span trace.Span
|
|
once sync.Once
|
|
}
|
|
|
|
func (b *spanBody) finish(err error) {
|
|
b.once.Do(func() {
|
|
if err != nil && err != io.EOF {
|
|
b.span.SetStatus(codes.Error, "body read failed")
|
|
}
|
|
b.span.End()
|
|
})
|
|
}
|
|
func (b *spanBody) Read(p []byte) (int, error) {
|
|
n, err := b.ReadCloser.Read(p)
|
|
if err != nil {
|
|
b.finish(err)
|
|
}
|
|
return n, err
|
|
}
|
|
func (b *spanBody) Close() error { err := b.ReadCloser.Close(); b.finish(err); return err }
|
|
|
|
type RedisHook struct{ Runtime *Runtime }
|
|
|
|
func (h RedisHook) DialHook(next redis.DialHook) redis.DialHook {
|
|
return func(ctx context.Context, network, addr string) (conn net.Conn, err error) {
|
|
err = h.Runtime.Do(ctx, "redis.connect", func(ctx context.Context) error { conn, err = next(ctx, network, addr); return err })
|
|
return
|
|
}
|
|
}
|
|
func (h RedisHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
|
|
return func(ctx context.Context, cmd redis.Cmder) error {
|
|
var result error
|
|
_ = h.Runtime.Do(ctx, "redis.command", func(ctx context.Context) error {
|
|
result = next(ctx, cmd)
|
|
if result == redis.Nil {
|
|
return nil
|
|
}
|
|
return result
|
|
})
|
|
return result
|
|
}
|
|
}
|
|
func (h RedisHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
|
|
return func(ctx context.Context, cmds []redis.Cmder) error {
|
|
return h.Runtime.Do(ctx, "redis.pipeline", func(ctx context.Context) error { return next(ctx, cmds) })
|
|
}
|
|
}
|
|
|
|
// SQL records execution spans without SQL text or arguments. Query spans cover
|
|
// obtaining rows; iteration/Scan remain caller-owned. Use Do around transactions.
|
|
type SQL struct {
|
|
DB *sql.DB
|
|
Runtime *Runtime
|
|
}
|
|
|
|
func (d SQL) ExecContext(ctx context.Context, query string, args ...any) (result sql.Result, err error) {
|
|
err = d.Runtime.Do(ctx, "db.exec", func(ctx context.Context) error { result, err = d.DB.ExecContext(ctx, query, args...); return err })
|
|
return
|
|
}
|
|
func (d SQL) QueryContext(ctx context.Context, query string, args ...any) (rows *sql.Rows, err error) {
|
|
err = d.Runtime.Do(ctx, "db.query", func(ctx context.Context) error { rows, err = d.DB.QueryContext(ctx, query, args...); return err })
|
|
return
|
|
}
|