MB-28562: Support search query callbacks before and after execution
+ SearchQueryStartCallback + SearchQueryEndCallback
This commit is contained in:
parent
676c85c935
commit
40f63baeb9
|
@ -50,9 +50,11 @@ const storePath = "store"
|
||||||
|
|
||||||
var mappingInternalKey = []byte("_mapping")
|
var mappingInternalKey = []byte("_mapping")
|
||||||
|
|
||||||
const SearchMemCheckCallbackKey = "_search_mem_callback_key"
|
const SearchQueryStartCallbackKey = "_search_query_start_callback_key"
|
||||||
|
const SearchQueryEndCallbackKey = "_search_query_end_callback_key"
|
||||||
|
|
||||||
type SearchMemCheckCallbackFn func(size uint64) error
|
type SearchQueryStartCallbackFn func(size uint64) error
|
||||||
|
type SearchQueryEndCallbackFn func(size uint64) error
|
||||||
|
|
||||||
func indexStorePath(path string) string {
|
func indexStorePath(path string) string {
|
||||||
return path + string(os.PathSeparator) + storePath
|
return path + string(os.PathSeparator) + storePath
|
||||||
|
@ -483,15 +485,24 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr
|
||||||
collector.SetFacetsBuilder(facetsBuilder)
|
collector.SetFacetsBuilder(facetsBuilder)
|
||||||
}
|
}
|
||||||
|
|
||||||
if memCb := ctx.Value(SearchMemCheckCallbackKey); memCb != nil {
|
memNeeded := memNeededForSearch(req, searcher, collector)
|
||||||
if memCbFn, ok := memCb.(SearchMemCheckCallbackFn); ok {
|
if cb := ctx.Value(SearchQueryStartCallbackKey); cb != nil {
|
||||||
err = memCbFn(memNeededForSearch(req, searcher, collector))
|
if cbF, ok := cb.(SearchQueryStartCallbackFn); ok {
|
||||||
|
err = cbF(memNeeded)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if cb := ctx.Value(SearchQueryEndCallbackKey); cb != nil {
|
||||||
|
if cbF, ok := cb.(SearchQueryEndCallbackFn); ok {
|
||||||
|
defer func() {
|
||||||
|
_ = cbF(memNeeded)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
err = collector.Collect(ctx, searcher, indexReader)
|
err = collector.Collect(ctx, searcher, indexReader)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
|
|
|
@ -1871,7 +1871,7 @@ func BenchmarkScorchSearchOverhead(b *testing.B) {
|
||||||
benchmarkSearchOverhead(scorch.Name, b)
|
benchmarkSearchOverhead(scorch.Name, b)
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSearchMemCheckCallback(t *testing.T) {
|
func TestSearchQueryCallback(t *testing.T) {
|
||||||
defer func() {
|
defer func() {
|
||||||
err := os.RemoveAll("testidx")
|
err := os.RemoveAll("testidx")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
@ -1910,8 +1910,8 @@ func TestSearchMemCheckCallback(t *testing.T) {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx := context.WithValue(context.Background(), SearchMemCheckCallbackKey,
|
ctx := context.WithValue(context.Background(), SearchQueryStartCallbackKey,
|
||||||
SearchMemCheckCallbackFn(f))
|
SearchQueryStartCallbackFn(f))
|
||||||
_, err = index.SearchInContext(ctx, req)
|
_, err = index.SearchInContext(ctx, req)
|
||||||
if err != expErr {
|
if err != expErr {
|
||||||
t.Fatalf("Expected: %v, Got: %v", expErr, err)
|
t.Fatalf("Expected: %v, Got: %v", expErr, err)
|
||||||
|
|
Loading…
Reference in New Issue