Maps are slow (grafana#3113)

* Remove maps due to slow iteration speed. Signed-off-by: Joe Elliott <[email protected]> * unrollin' Signed-off-by: Joe Elliott <[email protected]> * cleanup Signed-off-by: Joe Elliott <[email protected]> * improve AttributeFor Signed-off-by: Joe Elliott <[email protected]> * sort attrs for consistent testing Signed-off-by: Joe Elliott <[email protected]> * changelog Signed-off-by: Joe Elliott <[email protected]> --------- Signed-off-by: Joe Elliott <[email protected]>
VoIPGRID · Nov 17, 2023 · 548d693 · 548d693
1 parent 85c021b
commit 548d693
Show file tree

Hide file tree

Showing 12 changed files with 409 additions and 181 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,7 +10,7 @@
 * [ENHANCEMENT] Make the trace ID label name configurable for remote written exemplars [#3074](https://github.com/grafana/tempo/pull/3074)
 * [ENHANCEMENT] Update poller to make use of previous results and reduce backend load. [#2652](https://github.com/grafana/tempo/pull/2652) (@zalegrala)
 * [ENHANCEMENT] Improve TraceQL regex performance in certain queries. [#3139](https://github.com/grafana/tempo/pull/3139) (@joe-elliott)
-* [BUGFIX] Readd session token to s3 credentials. [#3144](https://github.com/grafana/tempo/pull/3144) (@farodin91)
+* [ENHANCEMENT] Improve TraceQL performance in complex queries. [#3113](https://github.com/grafana/tempo/pull/3113) (@joe-elliott)
 
 ## v2.3.0 / 2023-10-30
 

diff --git a/pkg/traceql/ast_execute.go b/pkg/traceql/ast_execute.go
@@ -475,28 +475,11 @@ func (s Static) execute(Span) (Static, error) {
 }
 
 func (a Attribute) execute(span Span) (Static, error) {
-	atts := span.Attributes()
-	static, ok := atts[a]
+	static, ok := span.AttributeFor(a)
 	if ok {
 		return static, nil
 	}
 
-	// if the requested attribute has a scope none then we will check first for span attributes matching
-	// then any attributes matching. we don't need to both if this is an intrinsic b/c those will always
-	// be caught above if they exist
-	if a.Scope == AttributeScopeNone && a.Intrinsic == IntrinsicNone {
-		for attribute, static := range atts {
-			if a.Name == attribute.Name && attribute.Scope == AttributeScopeSpan {
-				return static, nil
-			}
-		}
-		for attribute, static := range atts {
-			if a.Name == attribute.Name {
-				return static, nil
-			}
-		}
-	}
-
 	return NewStaticNil(), nil
 }
 

diff --git a/pkg/traceql/ast_test.go b/pkg/traceql/ast_test.go
@@ -371,7 +371,24 @@ func (m *mockSpan) WithAttrBool(key string, value bool) *mockSpan {
 	return m
 }
 
-func (m *mockSpan) Attributes() map[Attribute]Static {
+func (m *mockSpan) AttributeFor(a Attribute) (Static, bool) {
+	s, ok := m.attributes[a]
+	// if not found explicitly, check if it's a span attribute
+	if !ok && a.Scope == AttributeScopeNone {
+		aSpan := a
+		aSpan.Scope = AttributeScopeSpan
+		s, ok = m.attributes[aSpan]
+	}
+	// if not found explicitly, check if it's a resource attribute
+	if !ok && a.Scope == AttributeScopeNone {
+		aRes := a
+		aRes.Scope = AttributeScopeResource
+		s, ok = m.attributes[aRes]
+	}
+	return s, ok
+}
+
+func (m *mockSpan) AllAttributes() map[Attribute]Static {
 	return m.attributes
 }
 

diff --git a/pkg/traceql/engine.go b/pkg/traceql/engine.go
@@ -179,7 +179,7 @@ func (e *Engine) ExecuteTagValues(
 	case AttributeScopeResource,
 		AttributeScopeSpan: // If tag is scoped, we can check the map directly
 		collectAttributeValue = func(s Span) bool {
-			if v, ok := s.Attributes()[tag]; ok {
+			if v, ok := s.AttributeFor(tag); ok {
 				return cb(v)
 			}
 			return false
@@ -193,13 +193,13 @@ func (e *Engine) ExecuteTagValues(
 		// If the tag is unscoped, we need to check resource and span scoped manually by building a new Attribute with each scope.
 		collectAttributeValue = func(s Span) bool {
 			if tag.Intrinsic != IntrinsicNone { // it's intrinsic
-				if v, ok := s.Attributes()[tag]; ok {
+				if v, ok := s.AttributeFor(tag); ok {
 					return cb(v)
 				}
 			} else { // it's unscoped
 				for _, scope := range []AttributeScope{AttributeScopeResource, AttributeScopeSpan} {
 					scopedAttr := Attribute{Scope: scope, Parent: tag.Parent, Name: tag.Name}
-					if v, ok := s.Attributes()[scopedAttr]; ok {
+					if v, ok := s.AttributeFor(scopedAttr); ok {
 						return cb(v)
 					}
 				}
@@ -296,7 +296,7 @@ func (e *Engine) asTraceSearchMetadata(spanset *Spanset) *tempopb.TraceSearchMet
 			Attributes:        nil,
 		}
 
-		atts := span.Attributes()
+		atts := span.AllAttributes()
 
 		if name, ok := atts[NewIntrinsic(IntrinsicName)]; ok {
 			tempopbSpan.Name = name.S

diff --git a/pkg/traceql/storage.go b/pkg/traceql/storage.go
@@ -77,9 +77,12 @@ func (f *FetchSpansRequest) appendCondition(c ...Condition) {
 }
 
 type Span interface {
-	// these are the actual fields used by the engine to evaluate queries
-	// if a Filter parameter is passed the spans returned will only have this field populated
-	Attributes() map[Attribute]Static
+	// AttributeFor returns the attribute for the given key. If the attribute is not found then
+	// the second return value will be false.
+	AttributeFor(Attribute) (Static, bool)
+	// AllAttributes returns a map of all attributes for this span. AllAttributes should be used sparingly
+	// and is expected to be significantly slower than AttributeFor.
+	AllAttributes() map[Attribute]Static
 
 	ID() []byte
 	StartTimeUnixNanos() uint64

diff --git a/pkg/traceqlmetrics/metrics.go b/pkg/traceqlmetrics/metrics.go
@@ -281,13 +281,13 @@ func GetMetrics(ctx context.Context, query, groupBy string, spanLimit int, start
 			}
 
 			var (
-				attrs  = s.Attributes()
-				series = MetricSeries{}
-				err    = attrs[status] == statusErr
+				series    = MetricSeries{}
+				status, _ = s.AttributeFor(status)
+				err       = status == statusErr
 			)
 
 			for i, g := range groupBys {
-				series[i] = KeyValue{Key: groupByKeys[i], Value: lookup(g, attrs)}
+				series[i] = KeyValue{Key: groupByKeys[i], Value: lookup(g, s)}
 			}
 
 			results.Record(series, s.DurationNanos(), err)
@@ -307,9 +307,9 @@ func GetMetrics(ctx context.Context, query, groupBy string, spanLimit int, start
 	return results, nil
 }
 
-func lookup(needles []traceql.Attribute, haystack map[traceql.Attribute]traceql.Static) traceql.Static {
+func lookup(needles []traceql.Attribute, span traceql.Span) traceql.Static {
 	for _, n := range needles {
-		if v, ok := haystack[n]; ok {
+		if v, ok := span.AttributeFor(n); ok {
 			return v
 		}
 	}

diff --git a/pkg/traceqlmetrics/mocks.go b/pkg/traceqlmetrics/mocks.go
@@ -45,14 +45,19 @@ func (m *mockSpan) WithErr() *mockSpan {
 	return m
 }
 
-func (m *mockSpan) Attributes() map[traceql.Attribute]traceql.Static { return m.attrs }
-func (m *mockSpan) ID() []byte                                       { return nil }
-func (m *mockSpan) StartTimeUnixNanos() uint64                       { return m.start }
-func (m *mockSpan) DurationNanos() uint64                            { return m.duration }
+func (m *mockSpan) AllAttributes() map[traceql.Attribute]traceql.Static { return m.attrs }
+func (m *mockSpan) ID() []byte                                          { return nil }
+func (m *mockSpan) StartTimeUnixNanos() uint64                          { return m.start }
+func (m *mockSpan) DurationNanos() uint64                               { return m.duration }
 func (m *mockSpan) DescendantOf([]traceql.Span, []traceql.Span, bool, bool, []traceql.Span) []traceql.Span {
 	return nil
 }
 
+func (m *mockSpan) AttributeFor(a traceql.Attribute) (traceql.Static, bool) {
+	s, ok := m.attrs[a]
+	return s, ok
+}
+
 func (m *mockSpan) SiblingOf([]traceql.Span, []traceql.Span, bool, []traceql.Span) []traceql.Span {
 	return nil
 }

diff --git a/tempodb/encoding/vparquet/block_traceql.go b/tempodb/encoding/vparquet/block_traceql.go
@@ -35,10 +35,36 @@ type span struct {
 	cbSpanset      *traceql.Spanset
 }
 
-func (s *span) Attributes() map[traceql.Attribute]traceql.Static {
+func (s *span) AllAttributes() map[traceql.Attribute]traceql.Static {
 	return s.attributes
 }
 
+func (s *span) AttributeFor(a traceql.Attribute) (traceql.Static, bool) {
+	atts := s.attributes
+	static, ok := atts[a]
+	if ok {
+		return static, ok
+	}
+
+	// if the requested attribute has a scope none then we will check first for span attributes matching
+	// then any attributes matching. we don't need to both if this is an intrinsic b/c those will always
+	// be caught above if they exist
+	if a.Scope == traceql.AttributeScopeNone && a.Intrinsic == traceql.IntrinsicNone {
+		for attribute, static := range atts {
+			if a.Name == attribute.Name && attribute.Scope == traceql.AttributeScopeSpan {
+				return static, true
+			}
+		}
+		for attribute, static := range atts {
+			if a.Name == attribute.Name {
+				return static, true
+			}
+		}
+	}
+
+	return traceql.NewStaticNil(), false
+}
+
 func (s *span) ID() []byte {
 	return s.id
 }

diff --git a/tempodb/encoding/vparquet2/block_traceql.go b/tempodb/encoding/vparquet2/block_traceql.go
@@ -37,10 +37,36 @@ type span struct {
 	cbSpanset      *traceql.Spanset
 }
 
-func (s *span) Attributes() map[traceql.Attribute]traceql.Static {
+func (s *span) AllAttributes() map[traceql.Attribute]traceql.Static {
 	return s.attributes
 }
 
+func (s *span) AttributeFor(a traceql.Attribute) (traceql.Static, bool) {
+	atts := s.attributes
+	static, ok := atts[a]
+	if ok {
+		return static, ok
+	}
+
+	// if the requested attribute has a scope none then we will check first for span attributes matching
+	// then any attributes matching. we don't need to both if this is an intrinsic b/c those will always
+	// be caught above if they exist
+	if a.Scope == traceql.AttributeScopeNone && a.Intrinsic == traceql.IntrinsicNone {
+		for attribute, static := range atts {
+			if a.Name == attribute.Name && attribute.Scope == traceql.AttributeScopeSpan {
+				return static, true
+			}
+		}
+		for attribute, static := range atts {
+			if a.Name == attribute.Name {
+				return static, true
+			}
+		}
+	}
+
+	return traceql.NewStaticNil(), false
+}
+
 func (s *span) ID() []byte {
 	return s.id
 }