diff --git a/search/searcher/search_fuzzy.go b/search/searcher/search_fuzzy.go index 608e04874..8ebbbf0e9 100644 --- a/search/searcher/search_fuzzy.go +++ b/search/searcher/search_fuzzy.go @@ -60,7 +60,7 @@ func NewFuzzySearcher(ctx context.Context, indexReader index.IndexReader, term s if ctx != nil { reportIOStats(dictBytesRead, ctx) - aggregateBytesRead(ctx, dictBytesRead) + search.RecordSearchCost(ctx, "add", dictBytesRead) } return NewMultiTermSearcher(ctx, indexReader, candidates, field, diff --git a/search/searcher/search_geoboundingbox.go b/search/searcher/search_geoboundingbox.go index d58a0645a..b309af38a 100644 --- a/search/searcher/search_geoboundingbox.go +++ b/search/searcher/search_geoboundingbox.go @@ -222,6 +222,10 @@ func buildRectFilter(ctx context.Context, dvReader index.DocValueReader, field s } }) if err == nil && found { + bytes := dvReader.BytesRead() + if bytes > 0 { + search.RecordSearchCost(ctx, "add", bytes) + } for i := range lons { if geo.BoundingBoxContains(lons[i], lats[i], minLon, minLat, maxLon, maxLat) { diff --git a/search/searcher/search_geopointdistance.go b/search/searcher/search_geopointdistance.go index 7e5cc0337..510a3a773 100644 --- a/search/searcher/search_geopointdistance.go +++ b/search/searcher/search_geopointdistance.go @@ -134,6 +134,10 @@ func buildDistFilter(ctx context.Context, dvReader index.DocValueReader, field s } }) if err == nil && found { + bytes := dvReader.BytesRead() + if bytes > 0 { + search.RecordSearchCost(ctx, "add", bytes) + } for i := range lons { dist := geo.Haversin(lons[i], lats[i], centerLon, centerLat) if dist <= maxDist/1000 { diff --git a/search/searcher/search_geopolygon.go b/search/searcher/search_geopolygon.go index 3bc6b35dd..33dee6c8a 100644 --- a/search/searcher/search_geopolygon.go +++ b/search/searcher/search_geopolygon.go @@ -107,6 +107,10 @@ func buildPolygonFilter(ctx context.Context, dvReader index.DocValueReader, fiel // Note: this approach works for points which are strictly inside // the polygon. ie it might fail for certain points on the polygon boundaries. if err == nil && found { + bytes := dvReader.BytesRead() + if bytes > 0 { + search.RecordSearchCost(ctx, "add", bytes) + } nVertices := len(coordinates) if len(coordinates) < 3 { return false diff --git a/search/searcher/search_geoshape.go b/search/searcher/search_geoshape.go index 23907acbe..adf949238 100644 --- a/search/searcher/search_geoshape.go +++ b/search/searcher/search_geoshape.go @@ -119,7 +119,7 @@ func buildRelationFilterOnShapes(ctx context.Context, dvReader index.DocValueRea if err == nil && found { bytes := dvReader.BytesRead() if bytes > 0 { - aggregateBytesRead(ctx, bytes) + search.RecordSearchCost(ctx, "add", bytes) } return found } diff --git a/search/searcher/search_numeric_range.go b/search/searcher/search_numeric_range.go index c22101714..0928611a0 100644 --- a/search/searcher/search_numeric_range.go +++ b/search/searcher/search_numeric_range.go @@ -89,7 +89,7 @@ func NewNumericRangeSearcher(ctx context.Context, indexReader index.IndexReader, // loaded, using the context if ctx != nil { reportIOStats(dictBytesRead, ctx) - aggregateBytesRead(ctx, dictBytesRead) + search.RecordSearchCost(ctx, "add", dictBytesRead) } // cannot return MatchNoneSearcher because of interaction with @@ -112,7 +112,7 @@ func NewNumericRangeSearcher(ctx context.Context, indexReader index.IndexReader, if ctx != nil { reportIOStats(dictBytesRead, ctx) - aggregateBytesRead(ctx, dictBytesRead) + search.RecordSearchCost(ctx, "add", dictBytesRead) } return NewMultiTermSearcherBytes(ctx, indexReader, terms, field, diff --git a/search/searcher/search_regexp.go b/search/searcher/search_regexp.go index 0e75bc6de..6d7cd9f65 100644 --- a/search/searcher/search_regexp.go +++ b/search/searcher/search_regexp.go @@ -103,7 +103,7 @@ func NewRegexpSearcher(ctx context.Context, indexReader index.IndexReader, patte if ctx != nil { reportIOStats(dictBytesRead, ctx) - aggregateBytesRead(ctx, dictBytesRead) + search.RecordSearchCost(ctx, "add", dictBytesRead) } return NewMultiTermSearcher(ctx, indexReader, candidateTerms, field, boost, diff --git a/search/searcher/search_term.go b/search/searcher/search_term.go index c74566a68..b7b268d8b 100644 --- a/search/searcher/search_term.go +++ b/search/searcher/search_term.go @@ -151,15 +151,3 @@ func getQueryType(ctx context.Context) bool { } return true } - -func aggregateBytesRead(ctx context.Context, bytes uint64) { - if ctx != nil { - queryType, ok := ctx.Value(search.QueryTypeKey).(string) - if ok { - aggCallbackFn := ctx.Value(search.SearchCostAggregatorKey) - if aggCallbackFn != nil { - aggCallbackFn.(search.SearchCostAggregatorCallbackFn)("add", queryType, bytes) - } - } - } -} diff --git a/search/searcher/search_term_prefix.go b/search/searcher/search_term_prefix.go index 5fddcfb93..d0c613dc8 100644 --- a/search/searcher/search_term_prefix.go +++ b/search/searcher/search_term_prefix.go @@ -50,7 +50,7 @@ func NewTermPrefixSearcher(ctx context.Context, indexReader index.IndexReader, p if ctx != nil { reportIOStats(fieldDict.BytesRead(), ctx) - aggregateBytesRead(ctx, fieldDict.BytesRead()) + search.RecordSearchCost(ctx, "add", fieldDict.BytesRead()) } return NewMultiTermSearcher(ctx, indexReader, terms, field, boost, options, true) diff --git a/search/searcher/search_term_range.go b/search/searcher/search_term_range.go index 247bbc11f..c9a1cfc21 100644 --- a/search/searcher/search_term_range.go +++ b/search/searcher/search_term_range.go @@ -85,7 +85,7 @@ func NewTermRangeSearcher(ctx context.Context, indexReader index.IndexReader, if ctx != nil { reportIOStats(fieldDict.BytesRead(), ctx) - aggregateBytesRead(ctx, fieldDict.BytesRead()) + search.RecordSearchCost(ctx, "add", fieldDict.BytesRead()) } return NewMultiTermSearcher(ctx, indexReader, terms, field, boost, options, true) diff --git a/search/util.go b/search/util.go index c8a2c75b8..43870eab9 100644 --- a/search/util.go +++ b/search/util.go @@ -73,9 +73,14 @@ func MergeFieldTermLocations(dest []FieldTermLocation, matches []*DocumentMatch) const SearchIOStatsCallbackKey = "_search_io_stats_callback_key" type SearchIOStatsCallbackFunc func(uint64) -type SearchCostAggregatorCallbackFn func(string, string, uint64) -const SearchCostAggregatorKey = "_search_cost_aggregator_key" +// The callback signature is (message, queryType, cost) which allows +// the caller to act on a particular query type and what its the associated +// cost of an operation. "add" indicates to increment the cost for the query +// "done" indicates a finish of the accounting of the costs. +type SearchIncrementalCostCallbackFn func(string, string, uint64) + +const SearchIncrementalCostKey = "_search_incremental_cost_key" const QueryTypeKey = "_query_type_key" func RecordSearchCost(ctx context.Context, msg string, bytes uint64) { @@ -86,8 +91,8 @@ func RecordSearchCost(ctx context.Context, msg string, bytes uint64) { queryType = "" } - aggCallbackFn := ctx.Value(SearchCostAggregatorKey) + aggCallbackFn := ctx.Value(SearchIncrementalCostKey) if aggCallbackFn != nil { - aggCallbackFn.(SearchCostAggregatorCallbackFn)(msg, queryType, bytes) + aggCallbackFn.(SearchIncrementalCostCallbackFn)(msg, queryType, bytes) } }