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:
Max 2025-12-20 16:46:54 +08:00
parent 580da82033
commit aacb81fb52
2 changed files with 62 additions and 19 deletions

View file

@ -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"}},

View file

@ -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
} }