Skip to content

fix(embedding): guard lazy dimension cache against concurrent first access - #2749

Open
skyzhao1223 wants to merge 1 commit into
coze-dev:mainfrom
skyzhao1223:fix/embedding-dims-race
Open

skyzhao1223 wants to merge 1 commit into
coze-dev:mainfrom
skyzhao1223:fix/embedding-dims-race

Conversation

@skyzhao1223

Copy link
Copy Markdown

Fixes a data race in the lazy embedding dimension cache.

Problem

Dimensions() lazily embeds a probe string and caches the result in dims/dim with no synchronization:

func (e *embedder) Dimensions() int64 {
    if e.dim <= 0 {
        embeddings, _ := e.EmbedStrings(...)
        e.dim = int64(len(embeddings[0]))
    }
    return e.dim
}

When several retrieval/indexing goroutines first call Dimensions() concurrently, they race on the shared field (detected by go test -race) and issue duplicate embed requests.

Fix

Guard the lazy initialization with a sync.Mutex in both embedder (infra/embedding/impl/http/http.go) and denseOnlyWrap (infra/embedding/impl/wrap/dense_only.go). A failed probe is not cached, so transient embed failures can be retried on the next call.

Test

dimensions_test.go runs Dimensions() from 32 goroutines against an httptest mock. It passes with the fix and reports DATA RACE on the old code.

go test -race ./infra/embedding/impl/http/ -run TestDimensionsConcurrent

…ccess

Dimensions() lazily embeds a probe string and caches the result in dims/dim
without synchronization. Concurrent first calls raced on the shared field
(go test -race: DATA RACE) and issued duplicate embed requests. Guard the
lazy init with a mutex so only one goroutine embeds and the result is
published safely; a failed probe is not cached, so it can be retried.

Add a concurrent Dimensions() test that fails with -race on the old code.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant