From c15c3e972febc244b7c919c2e4588ef29bc503db Mon Sep 17 00:00:00 2001 From: "ochan.kwon" Date: Mon, 27 Jul 2026 20:25:46 +0900 Subject: [PATCH] fix(retriever): rebuild a missing vector store engine on demand The engine registry is per-process, so a store's engine is missing whenever this process did not create it. That happens three ways: - Startup skips a store whose engine fails to construct and never retries it, so that instance cannot serve the store again -- not even after a restart, if the backend is slow to come up. - A store registered on one instance is absent on every other until they restart, because registration is not broadcast. - Registration can fail on the instance handling the create, leaving the store missing even there. Each one surfaces as "vector store is currently unavailable" and is cleared today only by restarting the process. Add GetOrLoadByStoreID, which falls back to building the engine from the store row when the lookup misses, and route the two bound-store lookups through it. The method goes on RetrieveEngineRegistry rather than StoreRegistry because that is the interface the factory functions receive; the two declare GetByStoreID separately and neither embeds the other. Concurrent callers collapse onto one build. That build runs on a context detached from the caller that started it: callers share the result, so letting the first one's cancellation abort it would fail everyone else waiting. A panic inside the build is contained, because singleflight re-raises it on a goroutine the HTTP recovery middleware cannot reach and the build calls third-party client constructors. Publishing is guarded by a per-store generation sampled before the build, so an engine that finishes after the entry was registered or removed by someone else does not overwrite it and leave that engine orphaned with its connections open. A failed build starts a cooldown: collapsing only helps concurrent callers, so without it a backend that stays down would cost a full build timeout on every sequential request. Separate "the store is gone" from "its engine cannot be built right now". Async delete handlers answer ErrVectorStoreNotFound with asynq.SkipRetry, so reporting a database blip or a backend outage that way would permanently discard knowledge-base and index delete tasks and orphan the vector data they exist to remove. ErrVectorStoreUnavailable covers the retryable cases; the repository already distinguishes them, returning (nil, nil) for a missing row and an error for a failure. For the same reason a caller's own cancellation is never folded into a store sentinel. The ownership check runs first and queries the database with that context, so on a shutdown it is usually where cancellation is noticed. A search resolves one engine per distinct store, sequentially, and this server sets no HTTP read or write timeout, so the resolution loop gets a budget above a single build. An exhausted budget still leaves the engines warming, since the build is detached, and is reported as an unavailable store because retrying is likely to succeed. Build failures are logged where they happen, with the panic stack. The cause cannot travel back to the caller -- factory errors embed endpoints and credentials -- so without this the one thing an operator needs when self-healing fails to heal would be discarded with the error. Note a deliberate contract change: an ownership lookup failure used to be reported as not-found, and a test asserted that on purpose. A database failure says nothing about whether the store exists, and treating it as permanent is what discards the task, so that assertion is inverted rather than preserved. --- internal/application/service/knowledgebase.go | 17 +- .../service/knowledgebase_pr3_test.go | 7 + .../knowledgebase_search_budget_test.go | 97 +++++ .../knowledgebase_search_fanout_test.go | 32 ++ .../knowledgebase_search_storegroup.go | 28 +- .../application/service/retriever/factory.go | 66 ++- .../service/retriever/factory_test.go | 7 +- .../application/service/retriever/registry.go | 256 ++++++++++- .../retriever/registry_rehydrate_test.go | 410 ++++++++++++++++++ .../service/retriever/registry_test.go | 18 +- internal/application/service/tag.go | 2 + internal/container/container.go | 6 +- .../engine_factory_opensearch_test.go | 4 +- .../retrieve_registry_wiring_test.go | 59 +++ internal/types/interfaces/retriever.go | 15 + 15 files changed, 996 insertions(+), 28 deletions(-) create mode 100644 internal/application/service/knowledgebase_search_budget_test.go create mode 100644 internal/application/service/retriever/registry_rehydrate_test.go create mode 100644 internal/container/retrieve_registry_wiring_test.go diff --git a/internal/application/service/knowledgebase.go b/internal/application/service/knowledgebase.go index 931de9fa0..ed7bf9c50 100644 --- a/internal/application/service/knowledgebase.go +++ b/internal/application/service/knowledgebase.go @@ -261,14 +261,20 @@ func (s *knowledgeBaseService) validateVectorStoreBinding( "reason": "cross-tenant or unknown store", }, "[kb.create] vector store not owned by tenant") return apperrors.NewVectorStoreBindingInvalidError("vector store not found") - case errors.Is(err, retriever.ErrVectorStoreNotFound): + case errors.Is(err, retriever.ErrVectorStoreNotFound), + errors.Is(err, retriever.ErrVectorStoreUnavailable): logger.WarnWithFields(ctx, logger.Fields{ "tenant_id": tenantID, "store_id": sanitized, - "reason": "store registered in DB but missing in registry", + "reason": "store recorded in DB but no engine could be resolved", }, "[kb.create] vector store currently unavailable") return apperrors.NewVectorStoreUnavailableError( "vector store is currently unavailable; check its connection configuration") + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + // The caller went away or ran out of time while the binding was being + // verified, which can now include rebuilding the store's engine. That + // is not a server fault, so it must not be logged and answered as one. + return err default: logger.ErrorWithFields(ctx, err, map[string]interface{}{ "tenant_id": tenantID, @@ -862,6 +868,13 @@ func (s *knowledgeBaseService) ProcessKBDelete(ctx context.Context, t *asynq.Tas logger.Errorf(ctx, "KB delete task aborted: %v (tenant=%d, kb=%s)", err, payload.TenantID, payload.KnowledgeBaseID) return asynq.SkipRetry } + if errors.Is(err, retriever.ErrVectorStoreUnavailable) { + // The store is there but its engine could not be built right now. + // Falling through would drop the embeddings this task exists to + // remove and report success, so ask for another attempt instead. + logger.Errorf(ctx, "KB delete task deferred: %v (tenant=%d, kb=%s)", err, payload.TenantID, payload.KnowledgeBaseID) + return err + } if err != nil { logger.Warnf(ctx, "Failed to create retrieve engine: %v", err) } else { diff --git a/internal/application/service/knowledgebase_pr3_test.go b/internal/application/service/knowledgebase_pr3_test.go index 468ec0072..5b75bbbfc 100644 --- a/internal/application/service/knowledgebase_pr3_test.go +++ b/internal/application/service/knowledgebase_pr3_test.go @@ -62,6 +62,13 @@ func (f *fakeRegistry) GetByStoreID(storeID string) (interfaces.RetrieveEngineSe return nil, stderrors.New("not registered") } +// This fake never rebuilds a missing engine, so a miss stays a miss. +func (f *fakeRegistry) GetOrLoadByStoreID( + _ context.Context, _ uint64, storeID string, +) (interfaces.RetrieveEngineService, error) { + return f.GetByStoreID(storeID) +} + // fakeKBRepo is the smallest KnowledgeBaseRepository needed by // CreateKnowledgeBase + CopyKnowledgeBase tests. It stores rows in a map // keyed by ID. No tenant scoping is applied — the tested paths already diff --git a/internal/application/service/knowledgebase_search_budget_test.go b/internal/application/service/knowledgebase_search_budget_test.go new file mode 100644 index 000000000..6b149ad4e --- /dev/null +++ b/internal/application/service/knowledgebase_search_budget_test.go @@ -0,0 +1,97 @@ +package service + +import ( + "context" + stderrors "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/Tencent/WeKnora/internal/application/service/retriever" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +// ctxRecordingRegistry captures the context handed to the store lookup and +// then fails, so the caller stops before building retrieval params. +type ctxRecordingRegistry struct { + interfaces.RetrieveEngineRegistry + seen context.Context +} + +func (r *ctxRecordingRegistry) GetByStoreID(string) (interfaces.RetrieveEngineService, error) { + return nil, stderrors.New("store not in registry") +} + +func (r *ctxRecordingRegistry) GetOrLoadByStoreID( + ctx context.Context, _ uint64, _ string, +) (interfaces.RetrieveEngineService, error) { + r.seen = ctx + return nil, stderrors.New("build refused") +} + +// TestResolveStoreGroups_BoundsEngineResolution pins the ceiling on engine +// resolution. +// +// Groups are resolved one at a time and each may rebuild a missing engine by +// dialing its backend. This server configures no HTTP read or write timeout, +// so without a budget here a search over several cold stores would run for as +// long as the sum of those builds with nothing to stop it. +func TestResolveStoreGroups_BoundsEngineResolution(t *testing.T) { + storeID := "00000000-0000-0000-0000-0000000000bb" + registry := &ctxRecordingRegistry{} + svc := &knowledgeBaseService{ + retrieveEngine: registry, + ownership: &fakeOwnership{owned: map[string]uint64{storeID: 1}}, + } + kb := &types.KnowledgeBase{ID: "kb-1", TenantID: 1, VectorStoreID: &storeID} + + _, err := svc.resolveStoreGroups( + context.Background(), kb, []*types.KnowledgeBase{kb}, types.SearchParams{}, 5) + + require.Error(t, err, "the refused build must surface as an error") + require.NotNil(t, registry.seen, "the lookup must have been reached") + + deadline, ok := registry.seen.Deadline() + require.True(t, ok, "engine resolution must run under a deadline") + // Compared against a literal rather than the constant itself: checking the + // constant against its own value passes no matter what it is changed to. + assert.InDelta(t, 12.0, time.Until(deadline).Seconds(), 0.5, + "the deadline must come from the resolve budget") +} + +// TestStoreResolveBudgetExceedsOneBuild pins the relationship the budget +// depends on. Resolution can rebuild an engine, so a budget at or below a +// single build timeout would cut off the first rebuild it is meant to allow +// and turn every cold store into an error. +func TestStoreResolveBudgetExceedsOneBuild(t *testing.T) { + t.Parallel() + assert.Greater(t, storeResolveBudget, retriever.EngineBuildTimeout, + "the search budget must leave room for at least one engine build") +} + +// TestResolveStoreGroups_BudgetDoesNotOutliveTheCaller checks that the budget +// only ever shortens the caller's context, never extends it. +func TestResolveStoreGroups_BudgetDoesNotOutliveTheCaller(t *testing.T) { + storeID := "00000000-0000-0000-0000-0000000000cc" + registry := &ctxRecordingRegistry{} + svc := &knowledgeBaseService{ + retrieveEngine: registry, + ownership: &fakeOwnership{owned: map[string]uint64{storeID: 1}}, + } + kb := &types.KnowledgeBase{ID: "kb-1", TenantID: 1, VectorStoreID: &storeID} + + callerDeadline := time.Now().Add(2 * time.Second) + ctx, cancel := context.WithDeadline(context.Background(), callerDeadline) + defer cancel() + + _, err := svc.resolveStoreGroups(ctx, kb, []*types.KnowledgeBase{kb}, types.SearchParams{}, 5) + require.Error(t, err) + + deadline, ok := registry.seen.Deadline() + require.True(t, ok) + assert.False(t, deadline.After(callerDeadline), + "a caller that allows less time than the budget must still win") +} diff --git a/internal/application/service/knowledgebase_search_fanout_test.go b/internal/application/service/knowledgebase_search_fanout_test.go index d2130dd2c..12043a9df 100644 --- a/internal/application/service/knowledgebase_search_fanout_test.go +++ b/internal/application/service/knowledgebase_search_fanout_test.go @@ -223,6 +223,31 @@ func TestClassifyFactoryError_TenantInfoMissingMaps(t *testing.T) { assert.Equal(t, apperrors.ErrVectorStoreBindingInvalid, app.Code) } +// A rebuild that runs out of time is reported as an unavailable store rather +// than falling through as a raw error, which the handler would surface as an +// internal failure. The store exists and the caller may succeed on a retry. +func TestClassifyFactoryError_DeadlineExceededMapsTo2201(t *testing.T) { + t.Parallel() + err := classifyFactoryError(context.Background(), + context.DeadlineExceeded, 7, "store-id-y") + app, ok := apperrors.IsAppError(err) + require.True(t, ok, "expected AppError, got %T", err) + assert.Equal(t, apperrors.ErrVectorStoreUnavailable, app.Code) + if strings.Contains(app.Message, "store-id-y") { + t.Fatalf("AppError message leaked store UUID: %q", app.Message) + } +} + +// A caller that walked away is left as-is: there is no response to shape, and +// turning it into a store verdict would misreport why the work stopped. +func TestClassifyFactoryError_CancellationPassesThrough(t *testing.T) { + t.Parallel() + err := classifyFactoryError(context.Background(), context.Canceled, 1, "x") + require.ErrorIs(t, err, context.Canceled) + _, ok := apperrors.IsAppError(err) + assert.False(t, ok, "a cancellation is not a store verdict") +} + func TestClassifyFactoryError_GenericErrorPassesThrough(t *testing.T) { t.Parallel() raw := stderrors.New("generic infra failure") @@ -440,6 +465,13 @@ func (r *fakeFanoutRegistry) GetByStoreID(id string) (interfaces.RetrieveEngineS return nil, stderrors.New("store not registered") } +// This fake never rebuilds a missing engine, so a miss stays a miss. +func (r *fakeFanoutRegistry) GetOrLoadByStoreID( + _ context.Context, _ uint64, id string, +) (interfaces.RetrieveEngineService, error) { + return r.GetByStoreID(id) +} + func buildBoundComposite(t *testing.T, svc interfaces.RetrieveEngineService) *retriever.CompositeRetrieveEngine { t.Helper() const storeID = "00000000-0000-0000-0000-000000000001" diff --git a/internal/application/service/knowledgebase_search_storegroup.go b/internal/application/service/knowledgebase_search_storegroup.go index d1c0fcb0d..8e194c94a 100644 --- a/internal/application/service/knowledgebase_search_storegroup.go +++ b/internal/application/service/knowledgebase_search_storegroup.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "time" "github.com/Tencent/WeKnora/internal/application/service/retriever" apperrors "github.com/Tencent/WeKnora/internal/errors" @@ -12,6 +13,12 @@ import ( secutils "github.com/Tencent/WeKnora/internal/utils" ) +// storeResolveBudget caps the time a single search spends resolving the +// engines for its store groups. Resolution is sequential and can rebuild a +// missing engine, so the worst case is one build timeout per distinct store; +// this bounds the total rather than the individual attempt. +const storeResolveBudget = 12 * time.Second + // storeGroup is one fan-out unit of HybridSearch: a set of KB IDs that share // the same (VectorStore, owning tenant) pair. // @@ -97,6 +104,15 @@ func (s *knowledgeBaseService) resolveStoreGroups( buckets[key] = append(buckets[key], kb) } + // Resolving a group can rebuild a missing store engine, which dials a + // backend. Those rebuilds happen one after another here, and this server + // sets no read or write timeout, so a search across several cold stores + // would otherwise have nothing bounding it. The rebuild itself is detached + // from this context, so an exhausted budget still leaves the engines + // warming and the next search finds them ready. + resolveCtx, cancelResolve := context.WithTimeout(ctx, storeResolveBudget) + defer cancelResolve() + groups := make([]*storeGroup, 0, len(buckets)) for key, groupKBs := range buckets { var storeIDPtr *string @@ -105,7 +121,7 @@ func (s *knowledgeBaseService) resolveStoreGroups( storeIDPtr = &sid } engine, err := retriever.CreateRetrieveEngineForKB( - ctx, s.retrieveEngine, s.ownership, key.tenantID, storeIDPtr) + resolveCtx, s.retrieveEngine, s.ownership, key.tenantID, storeIDPtr) if err != nil { return nil, classifyFactoryError(ctx, err, key.tenantID, key.storeID) } @@ -146,9 +162,19 @@ func classifyFactoryError( case errors.Is(err, retriever.ErrVectorStoreForbidden): return apperrors.NewVectorStoreBindingInvalidError( "vector store bound to the knowledge base is not available") + case errors.Is(err, retriever.ErrVectorStoreUnavailable): + return apperrors.NewVectorStoreUnavailableError( + "vector store is currently unavailable") case errors.Is(err, retriever.ErrVectorStoreNotFound): return apperrors.NewVectorStoreUnavailableError( "vector store is currently unavailable") + case errors.Is(err, context.DeadlineExceeded): + // Resolving the store ran out of time, which can happen while its + // engine is being rebuilt. The binding is fine and a retry may work, + // so report it as unavailable rather than letting it fall through as + // an internal error. + return apperrors.NewVectorStoreUnavailableError( + "vector store is currently unavailable") case errors.Is(err, retriever.ErrTenantInfoMissing): return apperrors.NewVectorStoreBindingInvalidError( "tenant information missing in context") diff --git a/internal/application/service/retriever/factory.go b/internal/application/service/retriever/factory.go index ce8fc0211..c143ac198 100644 --- a/internal/application/service/retriever/factory.go +++ b/internal/application/service/retriever/factory.go @@ -19,11 +19,21 @@ var ( // context (synchronous, unbound KB path) and none is present. ErrTenantInfoMissing = errors.New("tenant info not found in context") - // ErrVectorStoreNotFound is returned when the store ID is not registered - // (or an internal lookup for ownership failed). Async workers should treat - // this as non-retryable. + // ErrVectorStoreNotFound is returned when the store does not exist for the + // tenant. Async workers should treat this as non-retryable: no amount of + // waiting brings back a store that is not in the database. ErrVectorStoreNotFound = errors.New("vector store not available") + // ErrVectorStoreUnavailable is returned when the store exists but its + // engine could not be produced right now — the metadata database was + // unreachable, or building the engine failed against a backend that may + // simply be down. Async workers should retry rather than discard the task, + // which is why this is a separate sentinel: reporting it as not-found + // would turn a passing outage into permanently dropped work. It carries no + // detail because the underlying errors embed endpoints and credentials; + // the cause is logged where it happens. + ErrVectorStoreUnavailable = errors.New("vector store engine unavailable") + // ErrVectorStoreForbidden is returned when the resolved store is not // owned by the given tenant. This guards against cross-tenant access // in case the upstream validation layer has a gap. Async workers should @@ -83,12 +93,37 @@ func VerifyBinding( if !owned { return ErrVectorStoreForbidden } - if _, err := registry.GetByStoreID(storeID); err != nil { - return ErrVectorStoreNotFound + if _, err := registry.GetOrLoadByStoreID(ctx, tenantID, storeID); err != nil { + return classifyLookupError(err) } return nil } +// classifyLookupError narrows an engine-lookup failure to what the caller is +// allowed to see, while preserving the distinction that decides whether work +// gets retried or discarded. Context errors and the store sentinels pass +// through; anything unexpected is reported as retryable, because treating an +// unknown failure as permanent is what silently drops work. +func classifyLookupError(err error) error { + switch { + case isContextError(err), + errors.Is(err, ErrVectorStoreNotFound), + errors.Is(err, ErrVectorStoreUnavailable), + errors.Is(err, ErrVectorStoreForbidden): + return err + default: + return ErrVectorStoreUnavailable + } +} + +// isContextError reports whether err is the caller giving up rather than a +// verdict about the store. The distinction matters because async workers treat +// the store sentinels as permanent and stop retrying, so a cancelled or +// timed-out request must not be reported as one. +func isContextError(err error) bool { + return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) +} + // CreateRetrieveEngineForKB returns a CompositeRetrieveEngine resolved from // a KB's VectorStore binding. // @@ -165,14 +200,22 @@ func resolveBoundEngine( ) (*CompositeRetrieveEngine, error) { owned, err := ownership.StoreOwnedBy(ctx, storeID, tenantID) if err != nil { + // This lookup queries the database with the caller's context, so it is + // where a shutdown or a disconnect is usually noticed first. Reporting + // that as a store verdict would let async workers discard work that + // only needs running again. + if isContextError(err) { + return nil, err + } // Infrastructure failure — record the raw error for operators but - // do not leak internals to the caller. + // do not leak internals to the caller. The store itself may be fine, + // so this is retryable rather than not-found. logger.ErrorWithFields(ctx, err, map[string]interface{}{ "tenant_id": tenantID, "store_id": storeID, "reason": "ownership lookup failed", }) - return nil, ErrVectorStoreNotFound + return nil, ErrVectorStoreUnavailable } if !owned { // Cross-tenant attempt (or the store has been deleted in the @@ -183,14 +226,17 @@ func resolveBoundEngine( return nil, ErrVectorStoreForbidden } - svc, err := registry.GetByStoreID(storeID) + svc, err := registry.GetOrLoadByStoreID(ctx, tenantID, storeID) if err != nil { + if isContextError(err) { + return nil, err + } logger.ErrorWithFields(ctx, err, map[string]interface{}{ "tenant_id": tenantID, "store_id": storeID, - "reason": "store not registered", + "reason": "store engine could not be resolved", }) - return nil, ErrVectorStoreNotFound + return nil, classifyLookupError(err) } // Build the composite directly from the resolved service. diff --git a/internal/application/service/retriever/factory_test.go b/internal/application/service/retriever/factory_test.go index fadc1bea3..827aaed2f 100644 --- a/internal/application/service/retriever/factory_test.go +++ b/internal/application/service/retriever/factory_test.go @@ -259,8 +259,11 @@ func TestCreateRetrieveEngineForKB_OwnershipLookupError(t *testing.T) { _, err := CreateRetrieveEngineForKB(context.Background(), registry, ownership, 1, &storeID) require.Error(t, err) - assert.True(t, errors.Is(err, ErrVectorStoreNotFound), - "ownership infrastructure error must collapse to ErrVectorStoreNotFound; got %q", err) + assert.True(t, errors.Is(err, ErrVectorStoreUnavailable), + "a database failure says nothing about whether the store exists, so it must be "+ + "retryable rather than the permanent not-found sentinel; got %q", err) + assert.False(t, errors.Is(err, ErrVectorStoreNotFound), + "async workers stop retrying on not-found, which would discard the task") } // ----- CreateRetrieveEngineFromPayload ----- diff --git a/internal/application/service/retriever/registry.go b/internal/application/service/retriever/registry.go index 38de8f76b..d28015540 100644 --- a/internal/application/service/retriever/registry.go +++ b/internal/application/service/retriever/registry.go @@ -1,13 +1,34 @@ package retriever import ( + "context" "fmt" + "runtime/debug" "sync" + "time" + "golang.org/x/sync/singleflight" + + "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) +// EngineBuildTimeout bounds a single on-demand engine construction. It is +// exported because callers that budget a sequence of resolutions need to size +// their own ceiling above one build; leaving the two coupled only by a comment +// invites them to drift apart. +// +// Not every engine constructor observes the context it is handed, so this +// bounds the ones that dial eagerly rather than every possible backend. +const EngineBuildTimeout = 10 * time.Second + +// rebuildCooldown throttles rebuild attempts for a store whose engine just +// failed to build. Without it a backend that stays down costs a full build +// timeout on every request, because collapsing only helps concurrent callers — +// sequential ones each open a new attempt. +const rebuildCooldown = 30 * time.Second + // RetrieveEngineRegistry implements the retrieval engine registry. // It maintains two maps: // - byEngineType: env stores registered via RETRIEVE_DRIVER (backward compatible) @@ -18,13 +39,96 @@ type RetrieveEngineRegistry struct { byEngineType map[types.RetrieverEngineType]interfaces.RetrieveEngineService byStoreID map[string]interfaces.RetrieveEngineService mu sync.RWMutex + + // repo and factory let the registry rebuild an engine that is missing from + // byStoreID. Both are optional: when either is nil the registry cannot + // rebuild anything and GetOrLoadByStoreID stays a plain lookup. + repo interfaces.VectorStoreRepository + factory interfaces.EngineFactory + sf singleflight.Group + + // storeGen counts every mutation of a store's entry. An on-demand build + // samples it before starting and publishes only if it has not moved, so a + // build cannot undo a registration or a removal that landed while it ran. + storeGen map[string]uint64 + // failedUntil holds the cooldown deadline per store, keyed only by stores + // that reached a build attempt. + failedUntil map[string]time.Time + // onFlightJoin, when set, fires once a caller is attached to a build. + // Tests need it to order a second caller against a build already running: + // until a caller is attached there is nothing to observe from outside, and + // releasing the build too early lets that caller miss it and read the + // finished engine instead, which looks identical from the results alone. + onFlightJoin func() + // flightObserver, when set, reports whether a caller shared its build with + // other callers. Collapsing is the point of this path but leaves no trace a + // test can check from the outside: a caller that misses the flight and then + // finds the finished engine is indistinguishable from one that waited on + // it. Production leaves this nil. + flightObserver func(shared bool) } -// NewRetrieveEngineRegistry creates a new retrieval engine registry -func NewRetrieveEngineRegistry() interfaces.RetrieveEngineRegistry { +// storeGeneration samples the mutation counter for a store. +func (r *RetrieveEngineRegistry) storeGeneration(storeID string) uint64 { + r.mu.RLock() + defer r.mu.RUnlock() + + return r.storeGen[storeID] +} + +// registerIfGenUnchanged publishes svc only when the store's entry has not +// been touched since gen was sampled. Reports whether the engine was published. +func (r *RetrieveEngineRegistry) registerIfGenUnchanged( + storeID string, gen uint64, svc interfaces.RetrieveEngineService, +) bool { + r.mu.Lock() + defer r.mu.Unlock() + + if r.storeGen[storeID] != gen { + return false + } + r.byStoreID[storeID] = svc + delete(r.failedUntil, storeID) + return true +} + +// inFailureCooldown reports whether a recent build failure should short-circuit +// another attempt. +func (r *RetrieveEngineRegistry) inFailureCooldown(storeID string) bool { + r.mu.RLock() + defer r.mu.RUnlock() + + until, exists := r.failedUntil[storeID] + return exists && time.Now().Before(until) +} + +// markBuildFailed starts the cooldown for a store whose engine build failed. +func (r *RetrieveEngineRegistry) markBuildFailed(storeID string) { + r.mu.Lock() + defer r.mu.Unlock() + + if r.failedUntil == nil { + r.failedUntil = make(map[string]time.Time) + } + r.failedUntil[storeID] = time.Now().Add(rebuildCooldown) +} + +// NewRetrieveEngineRegistry creates a new retrieval engine registry. +// +// repo and factory let the registry rebuild an engine that is missing from its +// store map. Passing nil for either disables that, leaving GetOrLoadByStoreID +// equivalent to GetByStoreID; both are required arguments rather than an +// optional extra so that every construction site has to say which it wants. +func NewRetrieveEngineRegistry( + repo interfaces.VectorStoreRepository, factory interfaces.EngineFactory, +) interfaces.RetrieveEngineRegistry { return &RetrieveEngineRegistry{ byEngineType: make(map[types.RetrieverEngineType]interfaces.RetrieveEngineService), byStoreID: make(map[string]interfaces.RetrieveEngineService), + storeGen: make(map[string]uint64), + failedUntil: make(map[string]time.Time), + repo: repo, + factory: factory, } } @@ -85,6 +189,19 @@ func (r *RetrieveEngineRegistry) RegisterWithStoreID(storeID string, svc interfa defer r.mu.Unlock() r.byStoreID[storeID] = svc + // Count this alongside unregistrations: an on-demand build that started + // earlier must not overwrite the entry published here, which would leave + // the engine this call installed orphaned with its connections open. + r.bumpGenerationLocked(storeID) +} + +// bumpGenerationLocked invalidates any on-demand build already in flight for +// this store. Callers must hold the write lock. +func (r *RetrieveEngineRegistry) bumpGenerationLocked(storeID string) { + if r.storeGen == nil { + r.storeGen = make(map[string]uint64) + } + r.storeGen[storeID]++ } // GetByStoreID retrieves an engine service by VectorStore ID. @@ -100,6 +217,117 @@ func (r *RetrieveEngineRegistry) GetByStoreID(storeID string) (interfaces.Retrie return svc, nil } +// GetOrLoadByStoreID returns the engine for storeID, rebuilding it from the +// database when this process has no entry for it. +// +// When either repo or factory is nil the registry cannot rebuild anything and +// a miss is reported exactly as GetByStoreID reports it. +func (r *RetrieveEngineRegistry) GetOrLoadByStoreID( + ctx context.Context, tenantID uint64, storeID string, +) (interfaces.RetrieveEngineService, error) { + if svc, err := r.GetByStoreID(storeID); err == nil { + return svc, nil + } + if r.repo == nil || r.factory == nil { + return nil, ErrVectorStoreNotFound + } + if r.inFailureCooldown(storeID) { + // A recent build failed; do not spend another timeout on it yet. + return nil, ErrVectorStoreUnavailable + } + // Sampled before the build starts: an unregistration landing after this + // point must prevent the finished engine from being published. + gen := r.storeGeneration(storeID) + + // Key by tenant as well as store so that a caller reaching this method + // without an ownership check cannot join another tenant's flight. + key := fmt.Sprintf("%d:%s", tenantID, storeID) + + results := r.doChanJoin(key, func() (res interface{}, err error) { + // singleflight re-raises a panic from this function on a goroutine of + // its own, out of reach of the HTTP recovery middleware. The build + // below calls third-party client constructors, so without this a + // broken one would take down the process instead of failing a request. + defer func() { + if recovered := recover(); recovered != nil { + logger.GetLogger(ctx).Errorf( + "[retriever.registry] engine build panicked for store %s: %v\n%s", + storeID, recovered, debug.Stack()) + res, err = nil, ErrVectorStoreUnavailable + } + }() + + // Detach from the initiating request. Callers collapse onto one build, + // so letting the first caller's cancellation abort it would fail every + // other caller waiting on the same engine. + buildCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), EngineBuildTimeout) + defer cancel() + + // An earlier flight for this key may have finished after the miss above. + if svc, lookupErr := r.GetByStoreID(storeID); lookupErr == nil { + return svc, nil + } + store, err := r.repo.GetByID(buildCtx, tenantID, storeID) + if err != nil { + // The store may well exist; the metadata database just could not + // answer. Saying "not found" here would make async workers discard + // their task over a passing outage. + logger.GetLogger(ctx).Errorf( + "[retriever.registry] loading store %s for rebuild failed: %v", storeID, err) + return nil, ErrVectorStoreUnavailable + } + if store == nil { + return nil, ErrVectorStoreNotFound + } + svc, err := r.factory(buildCtx, *store) + if err != nil { + // The cause dies with this log: it names the backend endpoint, so + // it must not travel back to the caller. + logger.GetLogger(ctx).Errorf( + "[retriever.registry] rebuilding engine for store %s failed, "+ + "retrying no sooner than %s: %v", storeID, rebuildCooldown, err) + r.markBuildFailed(storeID) + return nil, ErrVectorStoreUnavailable + } + if svc == nil { + // Publishing nil would panic every later reader of this store. + logger.GetLogger(ctx).Errorf( + "[retriever.registry] engine factory returned no engine for store %s", storeID) + return nil, ErrVectorStoreUnavailable + } + if !r.registerIfGenUnchanged(storeID, gen, svc) { + // The entry changed while this engine was being built, so this one + // is stale before it is published. Whatever landed instead is + // authoritative; the caller retries and picks it up. + return nil, ErrVectorStoreUnavailable + } + return svc, nil + }) + + select { + case <-ctx.Done(): + // The shared build keeps running for the other callers; this caller + // simply stops waiting. Reporting the context error rather than the + // not-found sentinel matters: callers map the sentinel to a permanent + // failure and would drop retryable work on a shutdown. + return nil, ctx.Err() + case result := <-results: + if r.flightObserver != nil { + r.flightObserver(result.Shared) + } + if result.Err != nil { + // Already a sentinel: the build logged its own cause, and the raw + // error is deliberately not carried back to the caller. + return nil, result.Err + } + svc, ok := result.Val.(interfaces.RetrieveEngineService) + if !ok { + return nil, ErrVectorStoreUnavailable + } + return svc, nil + } +} + // UnregisterByStoreID removes an engine service from the byStoreID map. // Idempotent: returns silently if the storeID is not found. // @@ -111,8 +339,32 @@ func (r *RetrieveEngineRegistry) UnregisterByStoreID(storeID string) { defer r.mu.Unlock() delete(r.byStoreID, storeID) + r.bumpGenerationLocked(storeID) + // Let an operator retry immediately after removing a store rather than + // waiting out a cooldown left over from the previous configuration. + delete(r.failedUntil, storeID) } // Compile-time assertion: *RetrieveEngineRegistry satisfies the // interfaces.RetrieveEngineRegistry contract, including GetByStoreID. var _ interfaces.RetrieveEngineRegistry = (*RetrieveEngineRegistry)(nil) + +// CanRebuildStores reports whether the registry was given what it needs to +// rebuild a store engine on demand. Exposed so that wiring can be asserted: +// a registry without those dependencies still serves lookups, so nothing else +// would reveal that rebuilding was silently left off. +func (r *RetrieveEngineRegistry) CanRebuildStores() bool { + return r.repo != nil && r.factory != nil +} + +// doChanJoin attaches the caller to a build for key, starting one if none is +// running, and reports the attachment to onFlightJoin. +func (r *RetrieveEngineRegistry) doChanJoin( + key string, build func() (interface{}, error), +) <-chan singleflight.Result { + results := r.sf.DoChan(key, build) + if r.onFlightJoin != nil { + r.onFlightJoin() + } + return results +} diff --git a/internal/application/service/retriever/registry_rehydrate_test.go b/internal/application/service/retriever/registry_rehydrate_test.go new file mode 100644 index 000000000..7969842ed --- /dev/null +++ b/internal/application/service/retriever/registry_rehydrate_test.go @@ -0,0 +1,410 @@ +package retriever + +import ( + "context" + stderrors "errors" + "sync" + "sync/atomic" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const rehydrateStoreID = "00000000-0000-0000-0000-0000000000aa" + +// fakeStoreRepo serves a single store by ID. Embedding the interface keeps the +// fake to the one method the rehydration path uses; any other call would +// nil-panic loudly rather than silently returning a zero value. +type fakeStoreRepo struct { + interfaces.VectorStoreRepository + store *types.VectorStore + err error +} + +func (f *fakeStoreRepo) GetByID(_ context.Context, _ uint64, _ string) (*types.VectorStore, error) { + return f.store, f.err +} + +// blockingFactory models a real engine constructor: it dials, which means it +// observes the context it is handed and aborts when that context is cancelled. +type blockingFactory struct { + entered chan struct{} + release chan struct{} + once sync.Once + calls atomic.Int32 +} + +func newBlockingFactory() *blockingFactory { + return &blockingFactory{entered: make(chan struct{}), release: make(chan struct{})} +} + +func (b *blockingFactory) build(ctx context.Context, _ types.VectorStore) ( + interfaces.RetrieveEngineService, error, +) { + b.calls.Add(1) + b.once.Do(func() { close(b.entered) }) + select { + case <-b.release: + return &mockEngineService{engineType: types.ElasticsearchRetrieverEngineType}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func newRehydratingRegistry(factory interfaces.EngineFactory) *RetrieveEngineRegistry { + return &RetrieveEngineRegistry{ + byEngineType: make(map[types.RetrieverEngineType]interfaces.RetrieveEngineService), + byStoreID: make(map[string]interfaces.RetrieveEngineService), + repo: &fakeStoreRepo{store: &types.VectorStore{ID: rehydrateStoreID}}, + factory: factory, + } +} + +// TestGetOrLoadByStoreID_LeaderCancelDoesNotPoisonWaiters pins the reason this +// path uses DoChan with a detached build context. +// +// Two requests miss the registry for the same store, so they collapse into one +// engine build. If that shared build inherits the first caller's context, the +// first caller going away — a closed browser tab, a shutting-down worker — +// aborts the build and hands its cancellation error to every other caller +// waiting on it. Those callers are alive and would then be told the store is +// unavailable, which is the very failure this rehydration exists to remove. +func TestGetOrLoadByStoreID_LeaderCancelDoesNotPoisonWaiters(t *testing.T) { + factory := newBlockingFactory() + registry := newRehydratingRegistry(factory.build) + // Installed before any caller starts: the hooks are plain fields, and a + // caller attaches to its build before the build itself reports progress, + // so setting them later both races and misses the first attachment. + shared := make(chan bool, 4) + registry.flightObserver = func(s bool) { shared <- s } + joined := make(chan struct{}, 4) + registry.onFlightJoin = func() { joined <- struct{}{} } + + leaderCtx, cancelLeader := context.WithCancel(context.Background()) + defer cancelLeader() + + var leaderErr error + leaderDone := make(chan struct{}) + go func() { + defer close(leaderDone) + _, leaderErr = registry.GetOrLoadByStoreID(leaderCtx, 1, rehydrateStoreID) + }() + + <-factory.entered // the build is running + <-joined // ...and the leader is attached to it + + var ( + waiterSvc interfaces.RetrieveEngineService + waiterErr error + ) + waiterDone := make(chan struct{}) + go func() { + defer close(waiterDone) + waiterSvc, waiterErr = registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + }() + + // Wait for the waiter to attach before letting the build finish. Released + // any earlier, the build would complete and register, and the waiter would + // read that engine directly instead of ever sharing the build — which the + // result assertions alone cannot tell apart. + <-joined + + cancelLeader() + <-leaderDone + + close(factory.release) + <-waiterDone + + require.NoError(t, waiterErr, "waiter must not inherit the leader's cancellation") + require.NotNil(t, waiterSvc, "waiter must receive the engine the shared build produced") + assert.ErrorIs(t, leaderErr, context.Canceled, "the cancelled leader still reports its own cancellation") + assert.Equal(t, int32(1), factory.calls.Load(), "the waiter must join the flight, not start a second build") + require.Len(t, shared, 1, "only the waiter reaches the result branch; the leader left on its context") + assert.Empty(t, joined, "exactly two attachments, both accounted for above") + assert.True(t, <-shared, "the waiter must have waited on the leader's build, not run its own") + + svc, err := registry.GetByStoreID(rehydrateStoreID) + require.NoError(t, err, "a completed build must leave the engine registered") + assert.NotNil(t, svc) +} + +// TestGetOrLoadByStoreID_ConcurrentMissesBuildOnce checks that a burst of +// requests for the same missing store produces one engine, not one per caller. +// Engine construction dials a backend, so duplicating it per request would turn +// a cold store into a connection storm. +func TestGetOrLoadByStoreID_ConcurrentMissesBuildOnce(t *testing.T) { + factory := newBlockingFactory() + registry := newRehydratingRegistry(factory.build) + close(factory.release) // let every build finish immediately + + const callers = 16 + var wg sync.WaitGroup + services := make([]interfaces.RetrieveEngineService, callers) + errs := make([]error, callers) + + wg.Add(callers) + for i := range callers { + go func() { + defer wg.Done() + services[i], errs[i] = registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + }() + } + wg.Wait() + + for i := range callers { + require.NoError(t, errs[i], "caller %d", i) + require.NotNil(t, services[i], "caller %d", i) + assert.Same(t, services[0], services[i], "every caller must get the same engine") + } + assert.Equal(t, int32(1), factory.calls.Load(), "one build for the whole burst") +} + +// TestGetOrLoadByStoreID_FactoryPanicDoesNotCrashProcess pins the panic guard. +// singleflight re-raises a panic from the shared build on its own goroutine, +// where the HTTP recovery middleware cannot catch it, so an unguarded panic +// from a third-party client constructor would end the process rather than the +// request. Reaching the assertions below at all is the real assertion. +func TestGetOrLoadByStoreID_FactoryPanicDoesNotCrashProcess(t *testing.T) { + registry := newRehydratingRegistry( + func(context.Context, types.VectorStore) (interfaces.RetrieveEngineService, error) { + panic("client constructor exploded") + }, + ) + + svc, err := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + + require.ErrorIs(t, err, ErrVectorStoreUnavailable, "a panicking build is reported as retryable") + assert.Nil(t, svc) + _, lookupErr := registry.GetByStoreID(rehydrateStoreID) + assert.Error(t, lookupErr, "a panicking build must leave nothing registered") +} + +// TestGetOrLoadByStoreID_UnregisterDuringBuildIsNotUndone covers the delete +// race: a store removed while its engine is still being built must stay +// removed, otherwise the finishing build silently resurrects it. +func TestGetOrLoadByStoreID_UnregisterDuringBuildIsNotUndone(t *testing.T) { + factory := newBlockingFactory() + registry := newRehydratingRegistry(factory.build) + + var ( + svc interfaces.RetrieveEngineService + err error + ) + done := make(chan struct{}) + go func() { + defer close(done) + svc, err = registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + }() + + <-factory.entered + registry.UnregisterByStoreID(rehydrateStoreID) // delete lands mid-build + close(factory.release) + <-done + + require.ErrorIs(t, err, ErrVectorStoreUnavailable, "a build overtaken by a delete must not succeed") + assert.Nil(t, svc) + _, lookupErr := registry.GetByStoreID(rehydrateStoreID) + assert.Error(t, lookupErr, "the deleted store must not be back in the registry") +} + +// TestGetOrLoadByStoreID_FailedBuildEntersCooldown checks that a store whose +// build just failed is not retried on the next request. A backend that stays +// down would otherwise cost a full build timeout per request, in exactly the +// outage this rehydration is meant to soften. +func TestGetOrLoadByStoreID_FailedBuildEntersCooldown(t *testing.T) { + var calls atomic.Int32 + registry := newRehydratingRegistry( + func(context.Context, types.VectorStore) (interfaces.RetrieveEngineService, error) { + calls.Add(1) + return nil, assert.AnError + }, + ) + + _, firstErr := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + _, secondErr := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + + require.ErrorIs(t, firstErr, ErrVectorStoreUnavailable) + require.ErrorIs(t, secondErr, ErrVectorStoreUnavailable, "raw build errors must not leak past the sentinel") + assert.Equal(t, int32(1), calls.Load(), "the second request is served by the cooldown, not a new build") + + // Removing the store clears the cooldown so an operator can retry at once. + registry.UnregisterByStoreID(rehydrateStoreID) + _, thirdErr := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + require.Error(t, thirdErr) + assert.Equal(t, int32(2), calls.Load(), "unregistering must clear the cooldown") +} + +// TestGetOrLoadByStoreID_WithoutRepoOrFactoryIsPlainLookup pins the degraded +// mode: a registry built without the rebuild dependencies must behave exactly +// as it did before rehydration existed. +func TestGetOrLoadByStoreID_WithoutRepoOrFactoryIsPlainLookup(t *testing.T) { + registry := &RetrieveEngineRegistry{ + byEngineType: make(map[types.RetrieverEngineType]interfaces.RetrieveEngineService), + byStoreID: make(map[string]interfaces.RetrieveEngineService), + } + + _, err := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + require.ErrorIs(t, err, ErrVectorStoreNotFound) + + registered := &mockEngineService{engineType: types.ElasticsearchRetrieverEngineType} + registry.RegisterWithStoreID(rehydrateStoreID, registered) + svc, err := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + require.NoError(t, err) + assert.Same(t, registered, svc, "a hit must not consult the database") +} + +// cancellingRegistry reports a caller-side cancellation from the rebuild path, +// the way the real registry does when the request that was waiting goes away. +// The lookup still misses, so the factory must reach the rebuild path first. +type cancellingRegistry struct { + interfaces.RetrieveEngineRegistry +} + +func (cancellingRegistry) GetByStoreID(string) (interfaces.RetrieveEngineService, error) { + return nil, stderrors.New("store not in registry") +} + +func (cancellingRegistry) GetOrLoadByStoreID( + context.Context, uint64, string, +) (interfaces.RetrieveEngineService, error) { + return nil, context.Canceled +} + +// TestBoundLookup_CancellationIsNotReportedAsMissingStore keeps a cancellation +// from being reported as a missing store. +// +// The async delete handlers treat ErrVectorStoreNotFound as permanent and +// answer it with asynq.SkipRetry. A shutdown cancels in-flight handlers, so +// folding that cancellation into the sentinel would make a rolling restart +// discard knowledge-base and index delete tasks that should have been retried, +// leaving vector data behind with nothing scheduled to remove it. +func TestBoundLookup_CancellationIsNotReportedAsMissingStore(t *testing.T) { + ownership := &fakeOwnership{owned: map[string]uint64{rehydrateStoreID: 1}} + storeID := rehydrateStoreID + + t.Run("CreateRetrieveEngineForKB", func(t *testing.T) { + _, err := CreateRetrieveEngineForKB( + context.Background(), cancellingRegistry{}, ownership, 1, &storeID) + + require.ErrorIs(t, err, context.Canceled, "the cancellation must reach the caller") + require.NotErrorIs(t, err, ErrVectorStoreNotFound, + "async handlers map this sentinel to SkipRetry and would drop the task") + }) + + t.Run("VerifyBinding", func(t *testing.T) { + err := VerifyBinding(context.Background(), cancellingRegistry{}, ownership, 1, storeID) + + require.ErrorIs(t, err, context.Canceled, "the cancellation must reach the caller") + require.NotErrorIs(t, err, ErrVectorStoreNotFound, + "async handlers map this sentinel to SkipRetry and would drop the task") + }) +} + +// TestGetOrLoadByStoreID_DatabaseFailureIsRetryable separates "the database +// could not answer" from "the store does not exist". +// +// The repository reports a missing row as (nil, nil) and an outage as an +// error, and only the first is permanent. Reporting an outage as not-found +// would make the async delete handlers answer a passing blip with SkipRetry +// and discard work that had nothing wrong with it. +func TestGetOrLoadByStoreID_DatabaseFailureIsRetryable(t *testing.T) { + registry := newRehydratingRegistry( + func(context.Context, types.VectorStore) (interfaces.RetrieveEngineService, error) { + t.Fatal("the engine factory must not run when the store could not be loaded") + return nil, nil + }, + ) + registry.repo = &fakeStoreRepo{err: stderrors.New("connection refused")} + + _, err := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + + require.ErrorIs(t, err, ErrVectorStoreUnavailable, "an outage is retryable") + require.NotErrorIs(t, err, ErrVectorStoreNotFound, + "async handlers stop retrying on not-found and would drop the task") +} + +// TestGetOrLoadByStoreID_AbsentStoreIsNotFound is the other half: a store that +// genuinely is not there must stay permanent, so workers stop instead of +// retrying something that can never succeed. +func TestGetOrLoadByStoreID_AbsentStoreIsNotFound(t *testing.T) { + registry := newRehydratingRegistry( + func(context.Context, types.VectorStore) (interfaces.RetrieveEngineService, error) { + t.Fatal("the engine factory must not run for a store that does not exist") + return nil, nil + }, + ) + registry.repo = &fakeStoreRepo{store: nil} + + _, err := registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + + require.ErrorIs(t, err, ErrVectorStoreNotFound) + require.NotErrorIs(t, err, ErrVectorStoreUnavailable) +} + +// TestBoundLookup_OwnershipCancellationIsNotAStoreVerdict covers the branch +// that runs before the registry is ever consulted. +// +// The ownership check queries the database with the caller's context, so on a +// shutdown it is normally the first place the cancellation is seen. Reporting +// it as a store verdict there would let the async handlers discard the task +// just as surely as reporting it later would. +func TestBoundLookup_OwnershipCancellationIsNotAStoreVerdict(t *testing.T) { + ownership := &fakeOwnership{err: context.Canceled} + storeID := rehydrateStoreID + + _, err := CreateRetrieveEngineForKB( + context.Background(), newRehydratingRegistry(nil), ownership, 1, &storeID) + + require.ErrorIs(t, err, context.Canceled, "the cancellation must reach the caller") + require.NotErrorIs(t, err, ErrVectorStoreNotFound, + "async handlers map this sentinel to SkipRetry and would drop the task") + require.NotErrorIs(t, err, ErrVectorStoreUnavailable, + "a cancellation is not a statement about the store") +} + +// TestBoundLookup_OwnershipOutageIsRetryable is the neighbouring case: the +// database failed for a reason of its own, which says nothing about whether +// the store exists, so the task deserves another attempt. +func TestBoundLookup_OwnershipOutageIsRetryable(t *testing.T) { + ownership := &fakeOwnership{err: stderrors.New("connection refused")} + storeID := rehydrateStoreID + + _, err := CreateRetrieveEngineForKB( + context.Background(), newRehydratingRegistry(nil), ownership, 1, &storeID) + + require.ErrorIs(t, err, ErrVectorStoreUnavailable) + require.NotErrorIs(t, err, ErrVectorStoreNotFound) +} + +// TestGetOrLoadByStoreID_BuildDoesNotOverwriteAConcurrentRegistration guards +// the create path against the rebuild path. +// +// Registering a store publishes its engine directly. If a rebuild that started +// earlier finished afterwards and published its own, the engine callers had +// just been given would be replaced by a second one built from the same row — +// and the first would be left running with its connections open and nobody +// holding it. +func TestGetOrLoadByStoreID_BuildDoesNotOverwriteAConcurrentRegistration(t *testing.T) { + factory := newBlockingFactory() + registry := newRehydratingRegistry(factory.build) + + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = registry.GetOrLoadByStoreID(context.Background(), 1, rehydrateStoreID) + }() + + <-factory.entered + registered := &mockEngineService{engineType: types.PostgresRetrieverEngineType} + registry.RegisterWithStoreID(rehydrateStoreID, registered) + close(factory.release) + <-done + + live, err := registry.GetByStoreID(rehydrateStoreID) + require.NoError(t, err) + assert.Same(t, registered, live, + "the engine published by the registration must survive the finishing build") +} diff --git a/internal/application/service/retriever/registry_test.go b/internal/application/service/retriever/registry_test.go index f8d2aa0ec..3dee21ce6 100644 --- a/internal/application/service/retriever/registry_test.go +++ b/internal/application/service/retriever/registry_test.go @@ -59,7 +59,7 @@ func newMock(engineType types.RetrieverEngineType) interfaces.RetrieveEngineServ // --- Register (byEngineType) tests --- func TestRegistry_Register(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) t.Run("success", func(t *testing.T) { err := reg.Register(newMock(types.PostgresRetrieverEngineType)) @@ -74,7 +74,7 @@ func TestRegistry_Register(t *testing.T) { } func TestRegistry_GetRetrieveEngineService(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) _ = reg.Register(newMock(types.PostgresRetrieverEngineType)) t.Run("found", func(t *testing.T) { @@ -91,7 +91,7 @@ func TestRegistry_GetRetrieveEngineService(t *testing.T) { } func TestRegistry_GetAllRetrieveEngineServices(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) _ = reg.Register(newMock(types.PostgresRetrieverEngineType)) _ = reg.Register(newMock(types.ElasticsearchRetrieverEngineType)) @@ -110,7 +110,7 @@ func TestRegistry_GetAllRetrieveEngineServices(t *testing.T) { // --- RegisterWithStoreID (byStoreID) tests --- func TestRegistry_RegisterWithStoreID(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) t.Run("success", func(t *testing.T) { reg.RegisterWithStoreID("store-1", newMock(types.PostgresRetrieverEngineType)) @@ -140,7 +140,7 @@ func TestRegistry_RegisterWithStoreID(t *testing.T) { } func TestRegistry_GetByStoreID(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) reg.RegisterWithStoreID("store-1", newMock(types.PostgresRetrieverEngineType)) t.Run("found", func(t *testing.T) { @@ -157,7 +157,7 @@ func TestRegistry_GetByStoreID(t *testing.T) { } func TestRegistry_UnregisterByStoreID(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) reg.RegisterWithStoreID("store-1", newMock(types.PostgresRetrieverEngineType)) t.Run("removes registered store", func(t *testing.T) { @@ -174,7 +174,7 @@ func TestRegistry_UnregisterByStoreID(t *testing.T) { // --- Dual map isolation tests --- func TestRegistry_DualMapIsolation(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) _ = reg.Register(newMock(types.PostgresRetrieverEngineType)) reg.RegisterWithStoreID("store-pg", newMock(types.PostgresRetrieverEngineType)) @@ -201,7 +201,7 @@ func TestRegistry_DualMapIsolation(t *testing.T) { // --- Concurrency test --- func TestRegistry_ConcurrentAccess(t *testing.T) { - reg := NewRetrieveEngineRegistry().(*RetrieveEngineRegistry) + reg := NewRetrieveEngineRegistry(nil, nil).(*RetrieveEngineRegistry) const goroutines = 10 var wg sync.WaitGroup @@ -229,7 +229,7 @@ func TestRegistry_ConcurrentAccess(t *testing.T) { // --- Interface compliance --- func TestRegistry_ImplementsStoreRegistry(t *testing.T) { - reg := NewRetrieveEngineRegistry() + reg := NewRetrieveEngineRegistry(nil, nil) concreteReg, ok := reg.(*RetrieveEngineRegistry) require.True(t, ok) diff --git a/internal/application/service/tag.go b/internal/application/service/tag.go index be2b889e5..e0f4f6b82 100644 --- a/internal/application/service/tag.go +++ b/internal/application/service/tag.go @@ -428,6 +428,8 @@ func (s *knowledgeTagService) ProcessIndexDelete(ctx context.Context, t *asynq.T logger.Errorf(ctx, "Index delete task aborted: %v (tenant=%d, kb=%s)", err, payload.TenantID, payload.KnowledgeBaseID) return asynq.SkipRetry } + // ErrVectorStoreUnavailable deliberately falls through to the retry path + // below: the store exists and its engine may build on a later attempt. if err != nil { logger.Warnf(ctx, "Failed to create retrieve engine for index cleanup: %v", err) return err diff --git a/internal/container/container.go b/internal/container/container.go index c1c258508..8a9265f5e 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -1028,8 +1028,12 @@ func initRawFileService(_ *config.Config) (interfaces.FileService, error) { // - Error if initialization fails func initRetrieveEngineRegistry( db *gorm.DB, cfg *config.Config, auditSvc interfaces.AuditLogService, + storeRepo interfaces.VectorStoreRepository, engineFactory interfaces.EngineFactory, ) (interfaces.RetrieveEngineRegistry, error) { - registry := retriever.NewRetrieveEngineRegistry() + // storeRepo and engineFactory let the registry rebuild a store engine that + // is absent from this process, which happens when startup skipped it after + // a construction failure or when another instance registered it. + registry := retriever.NewRetrieveEngineRegistry(storeRepo, engineFactory) retrieveDriver := strings.Split(os.Getenv("RETRIEVE_DRIVER"), ",") log := logger.GetLogger(context.Background()) // Audit sink for OpenSearch driver events (index created / reindex). Driver diff --git a/internal/container/engine_factory_opensearch_test.go b/internal/container/engine_factory_opensearch_test.go index d21c1c810..8599d312c 100644 --- a/internal/container/engine_factory_opensearch_test.go +++ b/internal/container/engine_factory_opensearch_test.go @@ -94,7 +94,9 @@ func TestInitRetrieveEngineRegistry_OpenSearchEnvPath(t *testing.T) { t.Setenv("RETRIEVE_DRIVER", "opensearch") t.Setenv("OPENSEARCH_ADDR", ts.URL) - registry, err := initRetrieveEngineRegistry(db, &config.Config{}, &fakeAuditSvc{}) + // nil store repository and engine factory: this exercises the env-driver + // path, which never rebuilds a database-backed store. + registry, err := initRetrieveEngineRegistry(db, &config.Config{}, &fakeAuditSvc{}, nil, nil) if err != nil { t.Fatalf("initRetrieveEngineRegistry: %v", err) } diff --git a/internal/container/retrieve_registry_wiring_test.go b/internal/container/retrieve_registry_wiring_test.go new file mode 100644 index 000000000..61635d2ae --- /dev/null +++ b/internal/container/retrieve_registry_wiring_test.go @@ -0,0 +1,59 @@ +package container + +import ( + "testing" + + "go.uber.org/dig" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/Tencent/WeKnora/internal/application/repository" + "github.com/Tencent/WeKnora/internal/application/service/retriever" + "github.com/Tencent/WeKnora/internal/config" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +// TestRetrieveEngineRegistryWiring checks that the container can still build +// the retrieval engine registry, and that what it builds can rebuild a missing +// store engine. +// +// The registry depends on the vector store repository and the engine factory, +// and the engine factory must not in turn depend on the registry. A cycle or a +// missing provider surfaces only when the process starts, so it is worth +// pinning here rather than discovering it at deploy time. The nil check at the +// end is the part that matters: a registry resolved without those two +// dependencies still satisfies the interface and still serves lookups, so it +// would pass every other test while silently never rebuilding anything. +func TestRetrieveEngineRegistryWiring(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open in-mem db: %v", err) + } + + c := dig.New() + provide := func(constructor interface{}) { + t.Helper() + if err := c.Provide(constructor); err != nil { + t.Fatalf("provide: %v", err) + } + } + provide(func() *gorm.DB { return db }) + provide(func() *config.Config { return &config.Config{} }) + provide(func() interfaces.AuditLogService { return &fakeAuditSvc{} }) + provide(repository.NewVectorStoreRepository) + provide(NewEngineFactory) + provide(initRetrieveEngineRegistry) + + err = c.Invoke(func(registry interfaces.RetrieveEngineRegistry) { + concrete, ok := registry.(*retriever.RetrieveEngineRegistry) + if !ok { + t.Fatalf("expected *retriever.RetrieveEngineRegistry, got %T", registry) + } + if !concrete.CanRebuildStores() { + t.Error("registry was built without the dependencies it needs to rebuild a store engine") + } + }) + if err != nil { + t.Fatalf("container could not build the registry: %v", err) + } +} diff --git a/internal/types/interfaces/retriever.go b/internal/types/interfaces/retriever.go index 2036808ff..581810557 100644 --- a/internal/types/interfaces/retriever.go +++ b/internal/types/interfaces/retriever.go @@ -81,6 +81,21 @@ type RetrieveEngineRegistry interface { // rather than calling this directly. The factories wrap GetByStoreID with // tenant ownership verification (defense-in-depth against cross-tenant IDOR). GetByStoreID(storeID string) (RetrieveEngineService, error) + + // GetOrLoadByStoreID returns the engine for storeID, rebuilding it from the + // database when this process has no entry for it. The registry is per-process: + // an engine registered on one instance is missing on every other until that + // instance restarts, and an engine whose creation failed during startup stays + // missing even across restarts. Rebuilding on demand lets both cases recover + // without an operator-driven rollout. + // + // Unlike GetByStoreID, this method scopes its database lookup to tenantID, so + // it cannot hydrate a store belonging to another tenant. Callers should still + // verify ownership first: the tenant scope is defense-in-depth, not a + // replacement for the ownership check. + GetOrLoadByStoreID( + ctx context.Context, tenantID uint64, storeID string, + ) (RetrieveEngineService, error) } // RetrieveEngineService defines the retrieve engine service interface