Propagate context to Search operation. (#2117)

* add search context to agent service Signed-off-by: hlts2 <hiroto.funakoshi.hiroto@gmail.com> * fix test execution error Signed-off-by: hlts2 <hiroto.funakoshi.hiroto@gmail.com> * fix build error of ngt stateful test Signed-off-by: hlts2 <hiroto.funakoshi.hiroto@gmail.com> --------- Signed-off-by: hlts2 <hiroto.funakoshi.hiroto@gmail.com>
vdaas · Jul 25, 2023 · a195594 · a195594
1 parent 5be892c
commit a195594
Show file tree

Hide file tree

Showing 5 changed files with 169 additions and 89 deletions.
diff --git a/internal/core/algorithm/ngt/ngt.go b/internal/core/algorithm/ngt/ngt.go
@@ -25,6 +25,7 @@ package ngt
 import "C"
 
 import (
+	"context"
 	"reflect"
 	"sync"
 	"unsafe"
@@ -40,7 +41,7 @@ type (
 	// NGT is core interface.
 	NGT interface {
 		// Search returns search result as []SearchResult
-		Search(vec []float32, size int, epsilon, radius float32) ([]SearchResult, error)
+		Search(ctx context.Context, vec []float32, size int, epsilon, radius float32) ([]SearchResult, error)
 
 		// Linear Search returns linear search result as []SearchResult
 		LinearSearch(vec []float32, size int) ([]SearchResult, error)
@@ -366,7 +367,7 @@ func (n *ngt) loadObjectSpace() error {
 }
 
 // Search returns search result as []SearchResult.
-func (n *ngt) Search(vec []float32, size int, epsilon, radius float32) (result []SearchResult, err error) {
+func (n *ngt) Search(ctx context.Context, vec []float32, size int, epsilon, radius float32) (result []SearchResult, err error) {
 	if len(vec) != int(n.dimension) {
 		return nil, errors.ErrIncompatibleDimensionSize(len(vec), int(n.dimension))
 	}
@@ -415,6 +416,12 @@ func (n *ngt) Search(vec []float32, size int, epsilon, radius float32) (result [
 	result = make([]SearchResult, rsize)
 
 	for i := range result {
+		select {
+		case <-ctx.Done():
+			n.PutErrorBuffer(ebuf)
+			return result[:i], nil
+		default:
+		}
 		d := C.ngt_get_result(results, C.uint32_t(i), ebuf)
 		if d.id == 0 && d.distance == 0 {
 			result[i] = SearchResult{0, 0, n.newGoError(ebuf)}