nodedb_vector/
builder.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Background HNSW builder thread.
4//!
5//! Each Data Plane core owns one builder thread. The core sends build
6//! requests with `try_send` and drains finished builds with `try_recv` once
7//! per tick, so a build never blocks the core's reactor. Both channels are
8//! bounded:
9//!
10//! - The request queue holds `capacity` requests. When it is full the core
11//!   keeps the job in its own backlog and sends it on a later tick. The
12//!   segment stays searchable by brute force meanwhile.
13//! - The completion queue holds `capacity` results. When it is full the
14//!   builder thread waits for the core to drain it.
15//!
16//! The thread builds requests in FIFO order and stops when the core drops
17//! its request sender.
18
19use std::sync::mpsc;
20use std::thread::JoinHandle;
21
22use tracing::{debug, info, warn};
23
24use crate::collection::{BuildComplete, BuildRequest};
25use crate::error::VectorError;
26use crate::hnsw::HnswIndex;
27
28/// Sender half: TPC core sends build requests to the builder thread.
29pub type BuildSender = mpsc::SyncSender<BuildRequest>;
30
31/// Receiver half: TPC core receives completed builds.
32pub type CompleteReceiver = mpsc::Receiver<BuildComplete>;
33
34/// Build requests a core's builder queue holds. Each request carries a whole
35/// segment's vectors, so the bound caps the memory in flight.
36pub const BUILD_QUEUE_CAPACITY: usize = 4;
37
38/// Spawn the HNSW builder thread for Data Plane core `core_id`, with request
39/// and completion queues of `capacity` entries each. Fails when the OS
40/// refuses the thread.
41pub fn spawn_builder(
42    core_id: usize,
43    capacity: usize,
44) -> std::io::Result<(BuildSender, CompleteReceiver, JoinHandle<()>)> {
45    let (request_tx, request_rx) = mpsc::sync_channel::<BuildRequest>(capacity);
46    let (complete_tx, complete_rx) = mpsc::sync_channel::<BuildComplete>(capacity);
47
48    let handle = std::thread::Builder::new()
49        .name(format!("hnsw-builder-{core_id}"))
50        .spawn(move || {
51            info!(core_id, "HNSW builder thread started");
52            builder_loop(core_id, request_rx, complete_tx);
53            info!(core_id, "HNSW builder thread stopped");
54        })?;
55
56    Ok((request_tx, complete_rx, handle))
57}
58
59fn builder_loop(
60    core_id: usize,
61    rx: mpsc::Receiver<BuildRequest>,
62    tx: mpsc::SyncSender<BuildComplete>,
63) {
64    while let Ok(req) = rx.recv() {
65        debug!(
66            core_id,
67            key = %req.key,
68            segment_id = req.segment_id,
69            vectors = req.vectors.len(),
70            dim = req.dim,
71            "building HNSW index"
72        );
73        let start = std::time::Instant::now();
74        let (key, segment_id, kind) = (req.key.clone(), req.segment_id, req.kind);
75        let result = build(core_id, req);
76        match &result {
77            Ok(index) => info!(
78                core_id,
79                key = %key,
80                segment_id,
81                vectors = index.len(),
82                elapsed_ms = start.elapsed().as_millis() as u64,
83                "HNSW index built"
84            ),
85            Err(e) => warn!(core_id, key = %key, segment_id, error = %e, "HNSW build failed"),
86        }
87        let complete = BuildComplete {
88            key,
89            segment_id,
90            kind,
91            result,
92        };
93        if tx.send(complete).is_err() {
94            warn!(core_id, "builder: core channel closed, stopping");
95            break;
96        }
97    }
98}
99
100/// Build one graph, inserting every vector in local-id order so node `i` of
101/// the graph is vector `i` of the request. The first insert error fails the
102/// build: a skipped vector would shift every later id.
103fn build(core_id: usize, req: BuildRequest) -> Result<HnswIndex, VectorError> {
104    let seed = (core_id as u64 + 1) * 1000 + u64::from(req.segment_id);
105    let mut index = HnswIndex::with_seed(req.dim, req.params, seed);
106    for vector in req.vectors {
107        index.insert(vector)?;
108    }
109    Ok(index)
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::collection::BuildKind;
116    use crate::hnsw::HnswParams;
117
118    fn request(segment_id: u32, vectors: Vec<Vec<f32>>) -> BuildRequest {
119        BuildRequest {
120            key: "k".into(),
121            segment_id,
122            kind: BuildKind::Seal,
123            vectors,
124            dim: 2,
125            params: HnswParams::default(),
126        }
127    }
128
129    #[test]
130    fn builds_keep_one_node_per_vector_and_report_errors() {
131        let (tx, rx, handle) = spawn_builder(0, 1).unwrap();
132        tx.send(request(1, vec![vec![1.0, 0.0], vec![0.0, 1.0]]))
133            .unwrap();
134        tx.send(request(2, vec![vec![1.0, 0.0], vec![1.0]]))
135            .unwrap();
136
137        let first = rx.recv().unwrap();
138        assert_eq!(first.segment_id, 1);
139        assert_eq!(first.result.unwrap().len(), 2);
140        let second = rx.recv().unwrap();
141        assert!(matches!(
142            second.result,
143            Err(VectorError::DimensionMismatch { .. })
144        ));
145
146        drop(tx);
147        handle.join().unwrap();
148    }
149
150    #[test]
151    fn a_full_request_queue_refuses_without_blocking() {
152        let (tx, _rx, _handle) = spawn_builder(0, 1).unwrap();
153        // The thread holds at most one request in hand and one queued; with
154        // the completion side undrained, later sends find the queue full.
155        let mut refused = false;
156        for id in 0..8 {
157            if let Err(mpsc::TrySendError::Full(_)) = tx.try_send(request(id, vec![vec![1.0, 0.0]]))
158            {
159                refused = true;
160                break;
161            }
162        }
163        assert!(refused, "a bounded queue must refuse once full");
164    }
165}
nodedb_vector/builder.rs

nodedb_vector/
builder.rs
