Skip to content

Commit 34fe775

Browse files
authored
fix(hnsw): restore multi-threaded build; vendor USearch #735 fix (#25163)
patch the fix with usearch PR: unum-cloud/USearch#772 add unit test to test the same bvt case flaky test to make sure no orphan in final index. revert the change build thread = 1 Approved by: @fengttt, @heni02, @XuPeng-SH
1 parent d538101 commit 34fe775

9 files changed

Lines changed: 276 additions & 92 deletions

File tree

pkg/vectorindex/hnsw/build.go

Lines changed: 98 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -36,11 +36,37 @@ type HnswBuild[T types.RealNumbers] struct {
3636
indexes []*HnswModel[T]
3737
nthread int
3838
add_chan chan AddItem[T]
39-
err_chan chan error
4039
wg sync.WaitGroup
4140
once sync.Once
4241
mutex sync.Mutex
4342
count atomic.Int64
43+
44+
// Worker-error propagation for the multi-threaded build. `stopped` is closed
45+
// once the first worker fails (or the context is cancelled); producers select on
46+
// it so an enqueue never blocks forever after the workers are gone, and finalizers
47+
// surface the recorded error instead of finishing a build as if it succeeded.
48+
stopOnce sync.Once
49+
stopped chan struct{}
50+
errMu sync.Mutex
51+
workerErr error
52+
}
53+
54+
// recordWorkerErr stores the first worker error and wakes any blocked producer /
55+
// finalizer. First-error-wins: the root failure is the most useful to report.
56+
func (h *HnswBuild[T]) recordWorkerErr(err error) {
57+
h.stopOnce.Do(func() {
58+
h.errMu.Lock()
59+
h.workerErr = err
60+
h.errMu.Unlock()
61+
close(h.stopped)
62+
})
63+
}
64+
65+
// WorkerErr returns the recorded worker error (nil if none). Safe to call any time.
66+
func (h *HnswBuild[T]) WorkerErr() error {
67+
h.errMu.Lock()
68+
defer h.errMu.Unlock()
69+
return h.workerErr
4470
}
4571

4672
type AddItem[T types.RealNumbers] struct {
@@ -52,30 +78,25 @@ type AddItem[T types.RealNumbers] struct {
5278
func NewHnswBuild[T types.RealNumbers](sqlproc *sqlexec.SqlProcess, uid string, nworker int32,
5379
cfg vectorindex.IndexConfig, tblcfg vectorindex.IndexTableConfig) (info *HnswBuild[T], err error) {
5480

55-
/*
56-
// estimate the number of worker threads
57-
nthread := 0
58-
if nworker <= 1 {
59-
// single database thread and set nthread to ThreadsBuild
60-
nthread = int(vectorindex.GetConcurrency(tblcfg.ThreadsBuild))
61-
} else {
62-
// multiple database worker threads
63-
threadsbuild := vectorindex.GetConcurrencyForBuild(tblcfg.ThreadsBuild)
64-
nthread = int(float64(threadsbuild) / float64(nworker))
65-
}
66-
if nthread < 1 {
67-
nthread = 1
68-
}
69-
*/
70-
71-
// MatrixOne #24849 / USearch #735 (open): concurrent add() can orphan nodes —
72-
// the vector is stored (contains() returns true) but the HNSW graph never links
73-
// it, so search() can never reach it, producing flaky recall@1 (an exact match
74-
// is intermittently missed). This is a real build race, not just HNSW
75-
// approximation. Reproduced in pkg/vectorindex/hnsw/zz_orphan_test.go:
76-
// multi-threaded build orphans ~1/30, single-threaded 0/30. Until the upstream
77-
// race is fixed, force a single build thread for correctness.
78-
nthread := 1
81+
// estimate the number of worker threads
82+
//
83+
// MatrixOne #24849 / USearch #735: concurrent add() used to orphan nodes (a
84+
// vector stored but never linked into the HNSW graph, so search() could not
85+
// reach it — flaky recall@1). That race is fixed in our usearch build (the
86+
// two-pass add: all forward links before any reverse link), so concurrent
87+
// builds now match single-threaded reachability. Multi-threaded build restored.
88+
nthread := 0
89+
if nworker <= 1 {
90+
// single database thread and set nthread to ThreadsBuild
91+
nthread = int(vectorindex.GetConcurrency(tblcfg.ThreadsBuild))
92+
} else {
93+
// multiple database worker threads
94+
threadsbuild := vectorindex.GetConcurrencyForBuild(tblcfg.ThreadsBuild)
95+
nthread = int(float64(threadsbuild) / float64(nworker))
96+
}
97+
if nthread < 1 {
98+
nthread = 1
99+
}
79100

80101
info = &HnswBuild[T]{
81102
uid: uid,
@@ -87,20 +108,21 @@ func NewHnswBuild[T types.RealNumbers](sqlproc *sqlexec.SqlProcess, uid string,
87108

88109
if nthread > 1 {
89110
info.add_chan = make(chan AddItem[T], nthread*4)
90-
info.err_chan = make(chan error, nthread)
111+
info.stopped = make(chan struct{})
91112

92113
// create multi-threads worker for add
93114
for i := 0; i < info.nthread; i++ {
94115

95116
info.wg.Add(1)
96117
go func() {
97118
defer info.wg.Done()
98-
var err0 error
99-
closed := false
100-
for !closed {
101-
closed, err0 = info.addFromChannel(sqlproc)
119+
for {
120+
closed, err0 := info.addFromChannel(sqlproc)
102121
if err0 != nil {
103-
info.err_chan <- err0
122+
info.recordWorkerErr(err0)
123+
return
124+
}
125+
if closed {
104126
return
105127
}
106128
}
@@ -134,21 +156,27 @@ func (h *HnswBuild[T]) addFromChannel(sqlproc *sqlexec.SqlProcess) (stream_close
134156
return false, nil
135157
}
136158

137-
func (h *HnswBuild[T]) CloseAndWait() {
159+
// CloseAndWait closes the work queue, waits for all workers to drain it, and
160+
// returns the first worker error (nil on success). It is idempotent; later calls
161+
// return the same recorded error.
162+
func (h *HnswBuild[T]) CloseAndWait() error {
138163
if h.nthread > 1 {
139164
h.once.Do(func() {
140165
close(h.add_chan)
141166
h.wg.Wait()
142167
})
143168
}
169+
return h.WorkerErr()
144170
}
145171

146172
// destroy
147173
func (h *HnswBuild[T]) Destroy() error {
148174

149175
var errs error
150176

151-
h.CloseAndWait()
177+
if err := h.CloseAndWait(); err != nil {
178+
errs = errors.Join(errs, err)
179+
}
152180

153181
for _, idx := range h.indexes {
154182
err := idx.Destroy()
@@ -162,18 +190,20 @@ func (h *HnswBuild[T]) Destroy() error {
162190

163191
func (h *HnswBuild[T]) Add(key int64, vec []T) error {
164192
if h.nthread > 1 {
165-
193+
// copy the []T slice.
194+
item := AddItem[T]{key, append(make([]T, 0, len(vec)), vec...)}
166195
select {
167-
case err := <-h.err_chan:
168-
return err
169-
default:
196+
case h.add_chan <- item:
197+
return nil
198+
case <-h.stopped:
199+
// A worker failed or the context was cancelled. Stop feeding the queue
200+
// (the send would otherwise block forever once workers are gone) and
201+
// surface the recorded error. recordWorkerErr stores the error before
202+
// closing `stopped`, so WorkerErr() is non-nil here.
203+
return h.WorkerErr()
170204
}
171-
// copy the []float32 slice.
172-
h.add_chan <- AddItem[T]{key, append(make([]T, 0, len(vec)), vec...)}
173-
return nil
174-
} else {
175-
return h.addVector(key, vec)
176205
}
206+
return h.addVector(key, vec)
177207
}
178208

179209
func (h *HnswBuild[T]) createIndexUniqueKey(id int64) string {
@@ -217,26 +247,32 @@ func (h *HnswBuild[T]) getIndexForAdd() (idx *HnswModel[T], save_idx *HnswModel[
217247
}
218248
h.count.Add(1)
219249

250+
// Reserve an in-flight slot on the index this add will go to, under the same lock
251+
// that decides rollover. A later rollover that hands this index back as save_idx
252+
// will wait for these to drain before SaveToFile() saves+destroys it.
253+
idx.inflight.Add(1)
254+
220255
return idx, save_idx, nil
221256
}
222257

223258
// add vector to the build
224259
// it will check the current index is full and add the vector to available index
225260
// sync version for multi-thread
226261
func (h *HnswBuild[T]) addVectorSync(key int64, vec []T) error {
227-
var err error
228-
var idx *HnswModel[T]
229-
var save_idx *HnswModel[T]
230-
231-
idx, save_idx, err = h.getIndexForAddSync()
262+
idx, save_idx, err := h.getIndexForAddSync()
232263
if err != nil {
233264
return err
234265
}
266+
defer idx.inflight.Done()
235267

236268
if save_idx != nil {
237-
// save the current index to file
238-
err = save_idx.SaveToFile()
239-
if err != nil {
269+
// Wait for every add already assigned to the rolled-over index to finish before
270+
// saving+destroying it. Otherwise SaveToFile() could persist a partial index or
271+
// free the usearch index while a peer worker is still calling idx.Add() on it.
272+
// This index receives no new adds (rollover already swapped in the next index
273+
// under the lock), so the wait converges.
274+
save_idx.inflight.Wait()
275+
if err = save_idx.SaveToFile(); err != nil {
240276
return err
241277
}
242278
}
@@ -248,21 +284,19 @@ func (h *HnswBuild[T]) addVectorSync(key int64, vec []T) error {
248284
// it will check the current index is full and add the vector to available index
249285
// single-threaded version.
250286
func (h *HnswBuild[T]) addVector(key int64, vec []T) error {
251-
var err error
252-
var idx *HnswModel[T]
253-
var save_idx *HnswModel[T]
254-
255287
h.mutex.Lock()
256288
defer h.mutex.Unlock()
257-
idx, save_idx, err = h.getIndexForAdd()
289+
idx, save_idx, err := h.getIndexForAdd()
258290
if err != nil {
259291
return err
260292
}
293+
defer idx.inflight.Done()
261294

262295
if save_idx != nil {
263-
// save the current index to file
264-
err = save_idx.SaveToFile()
265-
if err != nil {
296+
// Single-threaded: the rolled-over index has no in-flight adds (each add
297+
// completes before the next), so this is a no-op barrier kept for symmetry.
298+
save_idx.inflight.Wait()
299+
if err = save_idx.SaveToFile(); err != nil {
266300
return err
267301
}
268302
}
@@ -275,7 +309,12 @@ func (h *HnswBuild[T]) addVector(key int64, vec []T) error {
275309
// 2. sync the index file to index table
276310
func (h *HnswBuild[T]) ToInsertSql(ts int64) ([]string, error) {
277311

278-
h.CloseAndWait()
312+
// Surface any worker error from the multi-threaded build. Without this a worker
313+
// that failed on the last queued vector (after Add already returned nil) would be
314+
// silently dropped and the build finalized as if it succeeded.
315+
if err := h.CloseAndWait(); err != nil {
316+
return nil, err
317+
}
279318

280319
if len(h.indexes) == 0 {
281320
return []string{}, nil

0 commit comments

Comments
 (0)