From aacb81fb524b9d9d0f3ae941691c0e6e759da133 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 20 Dec 2025 16:46:54 +0800 Subject: [PATCH] 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. --- agent/search/jsapi_test.go | 22 ++++++++++++++ agent/search/search.go | 59 ++++++++++++++++++++++++++------------ 2 files changed, 62 insertions(+), 19 deletions(-) diff --git a/agent/search/jsapi_test.go b/agent/search/jsapi_test.go index 92be80d4..c7c34fa2 100644 --- a/agent/search/jsapi_test.go +++ b/agent/search/jsapi_test.go @@ -8,6 +8,7 @@ import ( "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search" "github.com/yaoapp/yao/agent/search/types" + "github.com/yaoapp/yao/agent/testutils" ) func TestNewJSAPI(t *testing.T) { @@ -51,6 +52,9 @@ func TestJSAPI_Web_WithOptions(t *testing.T) { } func TestJSAPI_KB(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + api := search.NewJSAPI(nil, &types.Config{ KB: &types.KBConfig{Collections: []string{"docs"}}, }, nil) @@ -66,6 +70,9 @@ func TestJSAPI_KB(t *testing.T) { } func TestJSAPI_KB_WithOptions(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + api := search.NewJSAPI(nil, &types.Config{ KB: &types.KBConfig{Collections: []string{"docs"}}, }, nil) @@ -87,6 +94,9 @@ func TestJSAPI_KB_WithOptions(t *testing.T) { } func TestJSAPI_DB(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + api := search.NewJSAPI(nil, &types.Config{ DB: &types.DBConfig{Models: []string{"product"}}, }, &search.Uses{QueryDSL: "builtin"}) @@ -102,6 +112,9 @@ func TestJSAPI_DB(t *testing.T) { } func TestJSAPI_DB_WithOptions(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + api := search.NewJSAPI(nil, &types.Config{ DB: &types.DBConfig{Models: []string{"product"}}, }, &search.Uses{QueryDSL: "builtin"}) @@ -122,6 +135,9 @@ func TestJSAPI_DB_WithOptions(t *testing.T) { } func TestJSAPI_All(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + api := search.NewJSAPI(nil, &types.Config{ KB: &types.KBConfig{Collections: []string{"docs"}}, DB: &types.DBConfig{Models: []string{"product"}}, @@ -155,6 +171,9 @@ func TestJSAPI_All(t *testing.T) { } func TestJSAPI_Any(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + api := search.NewJSAPI(nil, &types.Config{ KB: &types.KBConfig{Collections: []string{"docs"}}, DB: &types.DBConfig{Models: []string{"product"}}, @@ -186,6 +205,9 @@ func TestJSAPI_Any(t *testing.T) { } func TestJSAPI_Race(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + api := search.NewJSAPI(nil, &types.Config{ KB: &types.KBConfig{Collections: []string{"docs"}}, DB: &types.DBConfig{Models: []string{"product"}}, diff --git a/agent/search/search.go b/agent/search/search.go index a769832f..9ac04b3b 100644 --- a/agent/search/search.go +++ b/agent/search/search.go @@ -151,14 +151,24 @@ func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]* wg.Add(1) go func(idx int, r *types.Request) { defer wg.Done() - result, _ := s.Search(ctx, r) + + // Check if done before starting 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 { idx int result *types.Result }{idx, result}: - case <-done: - // Already found a successful result, discard this one } }(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) - var mu sync.Mutex + var foundSuccess bool for res := range resultChan { - mu.Lock() results[res.idx] = res.result // Check if this result has items (success = has results and no error) - if res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" { - mu.Unlock() - close(done) // Signal other goroutines to stop sending - return results, nil + if !foundSuccess && res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" { + foundSuccess = true + close(done) // Signal other goroutines to stop } - mu.Unlock() } - // No successful result found, return all results + // All goroutines have completed (resultChan is closed) return results, nil } // parallelRace returns as soon as any search completes (like Promise.race) // 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) { results := make([]*types.Result, len(reqs)) resultChan := make(chan struct { @@ -203,14 +211,24 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([] wg.Add(1) go func(idx int, r *types.Request) { defer wg.Done() - result, _ := s.Search(ctx, r) + + // Check if done before starting select { + case <-done: + return + default: + } + + result, _ := s.Search(ctx, r) + + // Try to send result + select { + case <-done: + // Already got first result case resultChan <- struct { idx int result *types.Result }{idx, result}: - case <-done: - // Already got first result, discard this one } }(i, req) } @@ -221,14 +239,17 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([] close(resultChan) }() - // Return immediately when first result arrives - if res, ok := <-resultChan; ok { + // Get first result and signal others to stop + var gotFirst bool + for res := range resultChan { results[res.idx] = res.result - close(done) // Signal other goroutines to stop sending - return results, nil + if !gotFirst { + 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 }