1// Unless explicitly stated otherwise all files in this repository are licensed 2// under the Apache License Version 2.0. 3// This product includes software developed at Datadog (https://www.datadoghq.com/). 4// Copyright 2016 Datadog, Inc. 5 6package buntdb 7 8import ( 9 "context" 10 "math" 11 "time" 12 13 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" 14 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext" 15 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" 16 "gopkg.in/DataDog/dd-trace-go.v1/internal/log" 17 18 "github.com/tidwall/buntdb" 19) 20 21// A DB wraps a buntdb.DB, automatically tracing any transactions. 22type DB struct { 23 *buntdb.DB 24 opts []Option 25 ctx context.Context 26} 27 28// Open calls buntdb.Open and wraps the result. 29func Open(path string, opts ...Option) (*DB, error) { 30 db, err := buntdb.Open(path) 31 if err != nil { 32 return nil, err 33 } 34 return WrapDB(db, opts...), nil 35} 36 37// WrapDB wraps a buntdb.DB so it can be traced. 38func WrapDB(db *buntdb.DB, opts ...Option) *DB { 39 return &DB{ 40 DB: db, 41 opts: opts, 42 ctx: context.Background(), 43 } 44} 45 46// Begin calls the underlying DB.Begin and traces the transaction. 47func (db *DB) Begin(writable bool) (*Tx, error) { 48 tx, err := db.DB.Begin(writable) 49 if err != nil { 50 return nil, err 51 } 52 return WrapTx(tx, db.opts...), nil 53} 54 55// Update calls the underlying DB.Update and traces the transaction. 56func (db *DB) Update(fn func(tx *Tx) error) error { 57 return db.DB.Update(func(tx *buntdb.Tx) error { 58 return fn(WrapTx(tx, db.opts...)) 59 }) 60} 61 62// View calls the underlying DB.View and traces the transaction. 63func (db *DB) View(fn func(tx *Tx) error) error { 64 return db.DB.View(func(tx *buntdb.Tx) error { 65 return fn(WrapTx(tx, db.opts...)) 66 }) 67} 68 69// WithContext sets the context for the DB. 70func (db *DB) WithContext(ctx context.Context) *DB { 71 newdb := *db 72 newdb.opts = append(newdb.opts[:len(newdb.opts):len(newdb.opts)], WithContext(ctx)) 73 return &newdb 74} 75 76// A Tx wraps a buntdb.Tx, automatically tracing any queries. 77type Tx struct { 78 *buntdb.Tx 79 cfg *config 80} 81 82// WrapTx wraps a buntdb.Tx so it can be traced. 83func WrapTx(tx *buntdb.Tx, opts ...Option) *Tx { 84 cfg := new(config) 85 defaults(cfg) 86 for _, opt := range opts { 87 opt(cfg) 88 } 89 log.Debug("contrib/tidwall/buntdb: Wrapping Transaction: %#v", cfg) 90 return &Tx{ 91 Tx: tx, 92 cfg: cfg, 93 } 94} 95 96func (tx *Tx) startSpan(name string) ddtrace.Span { 97 opts := []ddtrace.StartSpanOption{ 98 tracer.SpanType(ext.AppTypeDB), 99 tracer.ServiceName(tx.cfg.serviceName), 100 tracer.ResourceName(name), 101 } 102 if !math.IsNaN(tx.cfg.analyticsRate) { 103 opts = append(opts, tracer.Tag(ext.EventSampleRate, tx.cfg.analyticsRate)) 104 } 105 span, _ := tracer.StartSpanFromContext(tx.cfg.ctx, "buntdb.query", opts...) 106 return span 107} 108 109// WithContext sets the context for the Tx. 110func (tx *Tx) WithContext(ctx context.Context) *Tx { 111 newcfg := *tx.cfg 112 newcfg.ctx = ctx 113 return &Tx{ 114 Tx: tx.Tx, 115 cfg: &newcfg, 116 } 117} 118 119// Ascend calls the underlying Tx.Ascend and traces the query. 120func (tx *Tx) Ascend(index string, iterator func(key, value string) bool) error { 121 span := tx.startSpan("Ascend") 122 err := tx.Tx.Ascend(index, iterator) 123 span.Finish(tracer.WithError(err)) 124 return err 125} 126 127// AscendEqual calls the underlying Tx.AscendEqual and traces the query. 128func (tx *Tx) AscendEqual(index, pivot string, iterator func(key, value string) bool) error { 129 span := tx.startSpan("AscendEqual") 130 err := tx.Tx.AscendEqual(index, pivot, iterator) 131 span.Finish(tracer.WithError(err)) 132 return err 133} 134 135// AscendGreaterOrEqual calls the underlying Tx.AscendGreaterOrEqual and traces the query. 136func (tx *Tx) AscendGreaterOrEqual(index, pivot string, iterator func(key, value string) bool) error { 137 span := tx.startSpan("AscendGreaterOrEqual") 138 err := tx.Tx.AscendGreaterOrEqual(index, pivot, iterator) 139 span.Finish(tracer.WithError(err)) 140 return err 141} 142 143// AscendKeys calls the underlying Tx.AscendKeys and traces the query. 144func (tx *Tx) AscendKeys(pattern string, iterator func(key, value string) bool) error { 145 span := tx.startSpan("AscendKeys") 146 err := tx.Tx.AscendKeys(pattern, iterator) 147 span.Finish(tracer.WithError(err)) 148 return err 149} 150 151// AscendLessThan calls the underlying Tx.AscendLessThan and traces the query. 152func (tx *Tx) AscendLessThan(index, pivot string, iterator func(key, value string) bool) error { 153 span := tx.startSpan("AscendLessThan") 154 err := tx.Tx.AscendLessThan(index, pivot, iterator) 155 span.Finish(tracer.WithError(err)) 156 return err 157} 158 159// AscendRange calls the underlying Tx.AscendRange and traces the query. 160func (tx *Tx) AscendRange(index, greaterOrEqual, lessThan string, iterator func(key, value string) bool) error { 161 span := tx.startSpan("AscendRange") 162 err := tx.Tx.AscendRange(index, greaterOrEqual, lessThan, iterator) 163 span.Finish(tracer.WithError(err)) 164 return err 165} 166 167// CreateIndex calls the underlying Tx.CreateIndex and traces the query. 168func (tx *Tx) CreateIndex(name, pattern string, less ...func(a, b string) bool) error { 169 span := tx.startSpan("CreateIndex") 170 err := tx.Tx.CreateIndex(name, pattern, less...) 171 span.Finish(tracer.WithError(err)) 172 return err 173} 174 175// CreateIndexOptions calls the underlying Tx.CreateIndexOptions and traces the query. 176func (tx *Tx) CreateIndexOptions(name, pattern string, opts *buntdb.IndexOptions, less ...func(a, b string) bool) error { 177 span := tx.startSpan("CreateIndexOptions") 178 err := tx.Tx.CreateIndexOptions(name, pattern, opts, less...) 179 span.Finish(tracer.WithError(err)) 180 return err 181} 182 183// CreateSpatialIndex calls the underlying Tx.CreateSpatialIndex and traces the query. 184func (tx *Tx) CreateSpatialIndex(name, pattern string, rect func(item string) (min, max []float64)) error { 185 span := tx.startSpan("CreateSpatialIndex") 186 err := tx.Tx.CreateSpatialIndex(name, pattern, rect) 187 span.Finish(tracer.WithError(err)) 188 return err 189} 190 191// CreateSpatialIndexOptions calls the underlying Tx.CreateSpatialIndexOptions and traces the query. 192func (tx *Tx) CreateSpatialIndexOptions(name, pattern string, opts *buntdb.IndexOptions, rect func(item string) (min, max []float64)) error { 193 span := tx.startSpan("CreateSpatialIndexOptions") 194 err := tx.Tx.CreateSpatialIndexOptions(name, pattern, opts, rect) 195 span.Finish(tracer.WithError(err)) 196 return err 197} 198 199// Delete calls the underlying Tx.Delete and traces the query. 200func (tx *Tx) Delete(key string) (val string, err error) { 201 span := tx.startSpan("Delete") 202 val, err = tx.Tx.Delete(key) 203 span.Finish(tracer.WithError(err)) 204 return val, err 205} 206 207// DeleteAll calls the underlying Tx.DeleteAll and traces the query. 208func (tx *Tx) DeleteAll() error { 209 span := tx.startSpan("DeleteAll") 210 err := tx.Tx.DeleteAll() 211 span.Finish(tracer.WithError(err)) 212 return err 213} 214 215// Descend calls the underlying Tx.Descend and traces the query. 216func (tx *Tx) Descend(index string, iterator func(key, value string) bool) error { 217 span := tx.startSpan("Descend") 218 err := tx.Tx.Descend(index, iterator) 219 span.Finish(tracer.WithError(err)) 220 return err 221} 222 223// DescendEqual calls the underlying Tx.DescendEqual and traces the query. 224func (tx *Tx) DescendEqual(index, pivot string, iterator func(key, value string) bool) error { 225 span := tx.startSpan("DescendEqual") 226 err := tx.Tx.DescendEqual(index, pivot, iterator) 227 span.Finish(tracer.WithError(err)) 228 return err 229} 230 231// DescendGreaterThan calls the underlying Tx.DescendGreaterThan and traces the query. 232func (tx *Tx) DescendGreaterThan(index, pivot string, iterator func(key, value string) bool) error { 233 span := tx.startSpan("DescendGreaterThan") 234 err := tx.Tx.DescendGreaterThan(index, pivot, iterator) 235 span.Finish(tracer.WithError(err)) 236 return err 237} 238 239// DescendKeys calls the underlying Tx.DescendKeys and traces the query. 240func (tx *Tx) DescendKeys(pattern string, iterator func(key, value string) bool) error { 241 span := tx.startSpan("DescendKeys") 242 err := tx.Tx.DescendKeys(pattern, iterator) 243 span.Finish(tracer.WithError(err)) 244 return err 245} 246 247// DescendLessOrEqual calls the underlying Tx.DescendLessOrEqual and traces the query. 248func (tx *Tx) DescendLessOrEqual(index, pivot string, iterator func(key, value string) bool) error { 249 span := tx.startSpan("DescendLessOrEqual") 250 err := tx.Tx.DescendLessOrEqual(index, pivot, iterator) 251 span.Finish(tracer.WithError(err)) 252 return err 253} 254 255// DescendRange calls the underlying Tx.DescendRange and traces the query. 256func (tx *Tx) DescendRange(index, lessOrEqual, greaterThan string, iterator func(key, value string) bool) error { 257 span := tx.startSpan("DescendRange") 258 err := tx.Tx.DescendRange(index, lessOrEqual, greaterThan, iterator) 259 span.Finish(tracer.WithError(err)) 260 return err 261} 262 263// DropIndex calls the underlying Tx.DropIndex and traces the query. 264func (tx *Tx) DropIndex(name string) error { 265 span := tx.startSpan("DropIndex") 266 err := tx.Tx.DropIndex(name) 267 span.Finish(tracer.WithError(err)) 268 return err 269} 270 271// Get calls the underlying Tx.Get and traces the query. 272func (tx *Tx) Get(key string, ignoreExpired ...bool) (val string, err error) { 273 span := tx.startSpan("Get") 274 val, err = tx.Tx.Get(key, ignoreExpired...) 275 span.Finish(tracer.WithError(err)) 276 return val, err 277} 278 279// Indexes calls the underlying Tx.Indexes and traces the query. 280func (tx *Tx) Indexes() ([]string, error) { 281 span := tx.startSpan("Indexes") 282 indexes, err := tx.Tx.Indexes() 283 span.Finish(tracer.WithError(err)) 284 return indexes, err 285} 286 287// Intersects calls the underlying Tx.Intersects and traces the query. 288func (tx *Tx) Intersects(index, bounds string, iterator func(key, value string) bool) error { 289 span := tx.startSpan("Intersects") 290 err := tx.Tx.Intersects(index, bounds, iterator) 291 span.Finish(tracer.WithError(err)) 292 return err 293} 294 295// Len calls the underlying Tx.Len and traces the query. 296func (tx *Tx) Len() (int, error) { 297 span := tx.startSpan("Len") 298 n, err := tx.Tx.Len() 299 span.Finish(tracer.WithError(err)) 300 return n, err 301} 302 303// Nearby calls the underlying Tx.Nearby and traces the query. 304func (tx *Tx) Nearby(index, bounds string, iterator func(key, value string, dist float64) bool) error { 305 span := tx.startSpan("Nearby") 306 err := tx.Tx.Nearby(index, bounds, iterator) 307 span.Finish(tracer.WithError(err)) 308 return err 309} 310 311// Set calls the underlying Tx.Set and traces the query. 312func (tx *Tx) Set(key, value string, opts *buntdb.SetOptions) (previousValue string, replaced bool, err error) { 313 span := tx.startSpan("Set") 314 previousValue, replaced, err = tx.Tx.Set(key, value, opts) 315 span.Finish(tracer.WithError(err)) 316 return previousValue, replaced, err 317} 318 319// TTL calls the underlying Tx.TTL and traces the query. 320func (tx *Tx) TTL(key string) (time.Duration, error) { 321 span := tx.startSpan("TTL") 322 duration, err := tx.Tx.TTL(key) 323 span.Finish(tracer.WithError(err)) 324 return duration, err 325} 326