Enhance Test Cases and Improve Parallel Search Logic
- Updated multiple test cases in `jsapi_test.go` to utilize `testutils.Prepare` and `testutils.Clean` for better test setup and teardown, ensuring a consistent testing environment. - Refactored the `parallelAny` and `parallelRace` methods in `search.go` to improve goroutine management and result handling, reducing unnecessary locking and enhancing performance. - Implemented checks to prevent goroutines from executing after a successful result is found, optimizing resource usage during parallel searches.
This commit is contained in:
parent
580da82033
commit
aacb81fb52
2 changed files with 62 additions and 19 deletions
|
|
@ -8,6 +8,7 @@ import (
|
||||||
"github.com/yaoapp/yao/agent/context"
|
"github.com/yaoapp/yao/agent/context"
|
||||||
"github.com/yaoapp/yao/agent/search"
|
"github.com/yaoapp/yao/agent/search"
|
||||||
"github.com/yaoapp/yao/agent/search/types"
|
"github.com/yaoapp/yao/agent/search/types"
|
||||||
|
"github.com/yaoapp/yao/agent/testutils"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestNewJSAPI(t *testing.T) {
|
func TestNewJSAPI(t *testing.T) {
|
||||||
|
|
@ -51,6 +52,9 @@ func TestJSAPI_Web_WithOptions(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestJSAPI_KB(t *testing.T) {
|
func TestJSAPI_KB(t *testing.T) {
|
||||||
|
testutils.Prepare(t)
|
||||||
|
defer testutils.Clean(t)
|
||||||
|
|
||||||
api := search.NewJSAPI(nil, &types.Config{
|
api := search.NewJSAPI(nil, &types.Config{
|
||||||
KB: &types.KBConfig{Collections: []string{"docs"}},
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
||||||
}, nil)
|
}, nil)
|
||||||
|
|
@ -66,6 +70,9 @@ func TestJSAPI_KB(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestJSAPI_KB_WithOptions(t *testing.T) {
|
func TestJSAPI_KB_WithOptions(t *testing.T) {
|
||||||
|
testutils.Prepare(t)
|
||||||
|
defer testutils.Clean(t)
|
||||||
|
|
||||||
api := search.NewJSAPI(nil, &types.Config{
|
api := search.NewJSAPI(nil, &types.Config{
|
||||||
KB: &types.KBConfig{Collections: []string{"docs"}},
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
||||||
}, nil)
|
}, nil)
|
||||||
|
|
@ -87,6 +94,9 @@ func TestJSAPI_KB_WithOptions(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestJSAPI_DB(t *testing.T) {
|
func TestJSAPI_DB(t *testing.T) {
|
||||||
|
testutils.Prepare(t)
|
||||||
|
defer testutils.Clean(t)
|
||||||
|
|
||||||
api := search.NewJSAPI(nil, &types.Config{
|
api := search.NewJSAPI(nil, &types.Config{
|
||||||
DB: &types.DBConfig{Models: []string{"product"}},
|
DB: &types.DBConfig{Models: []string{"product"}},
|
||||||
}, &search.Uses{QueryDSL: "builtin"})
|
}, &search.Uses{QueryDSL: "builtin"})
|
||||||
|
|
@ -102,6 +112,9 @@ func TestJSAPI_DB(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestJSAPI_DB_WithOptions(t *testing.T) {
|
func TestJSAPI_DB_WithOptions(t *testing.T) {
|
||||||
|
testutils.Prepare(t)
|
||||||
|
defer testutils.Clean(t)
|
||||||
|
|
||||||
api := search.NewJSAPI(nil, &types.Config{
|
api := search.NewJSAPI(nil, &types.Config{
|
||||||
DB: &types.DBConfig{Models: []string{"product"}},
|
DB: &types.DBConfig{Models: []string{"product"}},
|
||||||
}, &search.Uses{QueryDSL: "builtin"})
|
}, &search.Uses{QueryDSL: "builtin"})
|
||||||
|
|
@ -122,6 +135,9 @@ func TestJSAPI_DB_WithOptions(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestJSAPI_All(t *testing.T) {
|
func TestJSAPI_All(t *testing.T) {
|
||||||
|
testutils.Prepare(t)
|
||||||
|
defer testutils.Clean(t)
|
||||||
|
|
||||||
api := search.NewJSAPI(nil, &types.Config{
|
api := search.NewJSAPI(nil, &types.Config{
|
||||||
KB: &types.KBConfig{Collections: []string{"docs"}},
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
||||||
DB: &types.DBConfig{Models: []string{"product"}},
|
DB: &types.DBConfig{Models: []string{"product"}},
|
||||||
|
|
@ -155,6 +171,9 @@ func TestJSAPI_All(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestJSAPI_Any(t *testing.T) {
|
func TestJSAPI_Any(t *testing.T) {
|
||||||
|
testutils.Prepare(t)
|
||||||
|
defer testutils.Clean(t)
|
||||||
|
|
||||||
api := search.NewJSAPI(nil, &types.Config{
|
api := search.NewJSAPI(nil, &types.Config{
|
||||||
KB: &types.KBConfig{Collections: []string{"docs"}},
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
||||||
DB: &types.DBConfig{Models: []string{"product"}},
|
DB: &types.DBConfig{Models: []string{"product"}},
|
||||||
|
|
@ -186,6 +205,9 @@ func TestJSAPI_Any(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestJSAPI_Race(t *testing.T) {
|
func TestJSAPI_Race(t *testing.T) {
|
||||||
|
testutils.Prepare(t)
|
||||||
|
defer testutils.Clean(t)
|
||||||
|
|
||||||
api := search.NewJSAPI(nil, &types.Config{
|
api := search.NewJSAPI(nil, &types.Config{
|
||||||
KB: &types.KBConfig{Collections: []string{"docs"}},
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
||||||
DB: &types.DBConfig{Models: []string{"product"}},
|
DB: &types.DBConfig{Models: []string{"product"}},
|
||||||
|
|
|
||||||
|
|
@ -151,14 +151,24 @@ func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]*
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func(idx int, r *types.Request) {
|
go func(idx int, r *types.Request) {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
result, _ := s.Search(ctx, r)
|
|
||||||
|
// Check if done before starting
|
||||||
select {
|
select {
|
||||||
|
case <-done:
|
||||||
|
return
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
result, _ := s.Search(ctx, r)
|
||||||
|
|
||||||
|
// Try to send result
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
// Already found a successful result
|
||||||
case resultChan <- struct {
|
case resultChan <- struct {
|
||||||
idx int
|
idx int
|
||||||
result *types.Result
|
result *types.Result
|
||||||
}{idx, result}:
|
}{idx, result}:
|
||||||
case <-done:
|
|
||||||
// Already found a successful result, discard this one
|
|
||||||
}
|
}
|
||||||
}(i, req)
|
}(i, req)
|
||||||
}
|
}
|
||||||
|
|
@ -170,25 +180,23 @@ func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]*
|
||||||
}()
|
}()
|
||||||
|
|
||||||
// Collect results until we find one with items (success)
|
// Collect results until we find one with items (success)
|
||||||
var mu sync.Mutex
|
var foundSuccess bool
|
||||||
for res := range resultChan {
|
for res := range resultChan {
|
||||||
mu.Lock()
|
|
||||||
results[res.idx] = res.result
|
results[res.idx] = res.result
|
||||||
// Check if this result has items (success = has results and no error)
|
// Check if this result has items (success = has results and no error)
|
||||||
if res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" {
|
if !foundSuccess && res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" {
|
||||||
mu.Unlock()
|
foundSuccess = true
|
||||||
close(done) // Signal other goroutines to stop sending
|
close(done) // Signal other goroutines to stop
|
||||||
return results, nil
|
|
||||||
}
|
}
|
||||||
mu.Unlock()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// No successful result found, return all results
|
// All goroutines have completed (resultChan is closed)
|
||||||
return results, nil
|
return results, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// parallelRace returns as soon as any search completes (like Promise.race)
|
// parallelRace returns as soon as any search completes (like Promise.race)
|
||||||
// Returns immediately when first result arrives, regardless of success/failure
|
// Returns immediately when first result arrives, regardless of success/failure
|
||||||
|
// Note: Still waits for all goroutines to complete before returning to avoid resource leaks
|
||||||
func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
|
func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
|
||||||
results := make([]*types.Result, len(reqs))
|
results := make([]*types.Result, len(reqs))
|
||||||
resultChan := make(chan struct {
|
resultChan := make(chan struct {
|
||||||
|
|
@ -203,14 +211,24 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func(idx int, r *types.Request) {
|
go func(idx int, r *types.Request) {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
result, _ := s.Search(ctx, r)
|
|
||||||
|
// Check if done before starting
|
||||||
select {
|
select {
|
||||||
|
case <-done:
|
||||||
|
return
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
result, _ := s.Search(ctx, r)
|
||||||
|
|
||||||
|
// Try to send result
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
// Already got first result
|
||||||
case resultChan <- struct {
|
case resultChan <- struct {
|
||||||
idx int
|
idx int
|
||||||
result *types.Result
|
result *types.Result
|
||||||
}{idx, result}:
|
}{idx, result}:
|
||||||
case <-done:
|
|
||||||
// Already got first result, discard this one
|
|
||||||
}
|
}
|
||||||
}(i, req)
|
}(i, req)
|
||||||
}
|
}
|
||||||
|
|
@ -221,14 +239,17 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]
|
||||||
close(resultChan)
|
close(resultChan)
|
||||||
}()
|
}()
|
||||||
|
|
||||||
// Return immediately when first result arrives
|
// Get first result and signal others to stop
|
||||||
if res, ok := <-resultChan; ok {
|
var gotFirst bool
|
||||||
|
for res := range resultChan {
|
||||||
results[res.idx] = res.result
|
results[res.idx] = res.result
|
||||||
close(done) // Signal other goroutines to stop sending
|
if !gotFirst {
|
||||||
return results, nil
|
gotFirst = true
|
||||||
|
close(done) // Signal other goroutines to stop
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// No results (shouldn't happen with valid requests)
|
// All goroutines have completed (resultChan is closed)
|
||||||
return results, nil
|
return results, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue