gibheer
/
bleve
0
0
Fork 0
Code

MB-28562: Support search query callbacks before and after execution 

+ SearchQueryStartCallback
+ SearchQueryEndCallback
This commit is contained in:
abhinavdangeti 2018-03-08 10:38:09 -08:00
parent 676c85c935
commit 40f63baeb9
2 changed files with 19 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
index_impl.go
View File

@ -50,9 +50,11 @@ const storePath = "store"
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 {
return path + string(os.PathSeparator) + storePath
@ -483,15 +485,24 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr
collector.SetFacetsBuilder(facetsBuilder)
}
if memCb := ctx.Value(SearchMemCheckCallbackKey); memCb != nil {
if memCbFn, ok := memCb.(SearchMemCheckCallbackFn); ok {
err = memCbFn(memNeededForSearch(req, searcher, collector))
memNeeded := memNeededForSearch(req, searcher, collector)
if cb := ctx.Value(SearchQueryStartCallbackKey); cb != nil {
if cbF, ok := cb.(SearchQueryStartCallbackFn); ok {
err = cbF(memNeeded)
}
}
if err != nil {
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)
if err != nil {
return nil, err

6
index_test.go
View File

@ -1871,7 +1871,7 @@ func BenchmarkScorchSearchOverhead(b *testing.B) {
benchmarkSearchOverhead(scorch.Name, b)
}
func TestSearchMemCheckCallback(t *testing.T) {
func TestSearchQueryCallback(t *testing.T) {
defer func() {
err := os.RemoveAll("testidx")
if err != nil {
@ -1910,8 +1910,8 @@ func TestSearchMemCheckCallback(t *testing.T) {
return nil
}
ctx := context.WithValue(context.Background(), SearchMemCheckCallbackKey,
SearchMemCheckCallbackFn(f))
ctx := context.WithValue(context.Background(), SearchQueryStartCallbackKey,
SearchQueryStartCallbackFn(f))
_, err = index.SearchInContext(ctx, req)
if err != expErr {
t.Fatalf("Expected: %v, Got: %v", expErr, err)