coverage — ~bigbes/sr-ht-spec — 3cb1c03d — search/index.go

Coverage
80.8% 118/146 statements
Δ
—
Blob
bd22689
Uncovered L107-L110L128-L129L142-L143L171-L172L184-L185L205-L206L209-L210L240-L241L248-L249L252-L253L260-L262L267-L269L271-L272L290-L292L297-L298L300-L301L303-L304L307-L308L311-L312L324-L325L328-L329L335-L336L353-L354L375-L376
1 package search
2
3 import (
4 "context"
5 "errors"
6 "fmt"
7 "os"
8 "path/filepath"
9 "sync"
10 "time"
11
12 "github.com/blevesearch/bleve/v2"
13
14 "sourcecraft.dev/bigbes/sr-ht-spec/core"
15 )
16
17 // Document is one unit of the index: a whole document, or one dated entry of an
18 // activity log. Extract produces these from a doc.Archive; the index stores
19 // them and hands back Hits shaped the same way.
20 type Document struct {
21 // Space is what makes the one global index filterable per project. Every
22 // document carries it, and a project query is a term filter over the set.
23 Space core.SpaceRef
24 // ID is the document's id within its space: doc.Page.ID, or
25 // "<page id>#<date>-<n>" for one entry of an activity log.
26 ID string
27 // Rev is the revision the document was read at, carried through so a hit
28 // can be turned into a pinned `?rev=` URL rather than a link to whatever
29 // the branch says now.
30 Rev string
31 // Path is the document's path in the git tree. For a log entry it is the
32 // path of the log document the entry came out of.
33 Path string
34 // Anchor is the heading anchor within Path a hit should land on. Empty for
35 // an ordinary document, set for a log entry.
36 Anchor string
37 // Section is the top-level directory the document lives under, or
38 // doc.LogSection for an activity log and its entries.
39 Section string
40 Title string
41 // Text is everything searchable: the frontmatter projected to "key: value"
42 // lines followed by the rendered plain text. It is split by language and
43 // stored in the analyzed fields; it is not stored verbatim under its own
44 // name.
45 Text string
46 }
47
48 // Index is the one global bleve index, shared by every space and every project.
49 //
50 // One index, not one per project and not one per space. Per-project indexes
51 // were specified and then retracted in the design for a concrete reason: with N
52 // projects every merge fans out to N rebuilds and adding a space to a project
53 // forces one, while the "everything" project is a second full copy of the
54 // corpus. Here a project is a filter — see Query.Spaces — so a merge touches
55 // one index and every project containing the space sees the change for free.
56 //
57 // bleve is single-writer. One process holds the index open, which is why the
58 // push hooks are RPC shims into the daemon rather than separate processes.
59 // Within that process an Index is safe for concurrent use: searches share a
60 // read lock, and a rebuild takes the write lock.
61 type Index struct {
62 path string
63
64 mu sync.RWMutex
65 idx bleve.Index
66 }
67
68 // batchSize is how many documents are buffered before a batch is flushed
69 // during a full rebuild. Carried over from warren.
70 const batchSize = 200
71
72 // Stats reports what a rebuild did and how long it took.
73 //
74 // The design absorbs warren's batch full rebuild deliberately — at tens of
75 // documents a day, incremental indexing is machinery bought against a cost
76 // nobody has measured — and asks for the duration to be instrumented so the
77 // decision to revisit is triggered by a number rather than a hunch. That is
78 // what Took is for: it is the trigger, and callers are expected to log it.
79 type Stats struct {
80 // Spaces is how many distinct spaces the rebuild wrote.
81 Spaces int
82 // Indexed is how many documents were written.
83 Indexed int
84 // Deleted is how many stale documents were removed: documents that were in
85 // the index for a rebuilt space and are not in the new document set.
86 Deleted int
87 Took time.Duration
88 }
89
90 3 func (s Stats) String() string {
91 3 return fmt.Sprintf("spaces=%d indexed=%d deleted=%d took=%s",
92 3 s.Spaces, s.Indexed, s.Deleted, s.Took.Round(time.Millisecond))
93 3 }
94
95 // Open opens the global index at path, creating an empty one if it is not
96 // there. The index is a pure cache: deleting the directory and letting Open
97 // recreate it, followed by RebuildAll, is always a valid repair.
98 31 func Open(path string) (*Index, error) {
99 31 if path == "" {
100 1 return nil, errors.New("search: index path is required")
101 1 }
102 30 idx, err := bleve.Open(path)
103 30 switch {
104 29 case errors.Is(err, bleve.ErrorIndexPathDoesNotExist):
105 29 idx, err = bleve.New(path, buildMapping())
106 29 if err != nil {
107 0 return nil, fmt.Errorf("search: create index at %s: %w", path, err)
108 0 }
109 0 case err != nil:
110 0 return nil, fmt.Errorf("search: open index at %s: %w", path, err)
111 }
112 30 return &Index{path: path, idx: idx}, nil
113 }
114
115 // Path is where the index lives on disk.
116 2 func (x *Index) Path() string { return x.path }
117
118 // Close releases the index.
119 31 func (x *Index) Close() error {
120 31 x.mu.Lock()
121 31 defer x.mu.Unlock()
122 31 if x.idx == nil {
123 1 return nil
124 1 }
125 30 err := x.idx.Close()
126 30 x.idx = nil
127 30 if err != nil {
128 0 return fmt.Errorf("search: close index: %w", err)
129 0 }
130 30 return nil
131 }
132
133 // Count is how many documents the index holds, across every space.
134 2 func (x *Index) Count() (uint64, error) {
135 2 x.mu.RLock()
136 2 defer x.mu.RUnlock()
137 2 if x.idx == nil {
138 1 return 0, errors.New("search: index is closed")
139 1 }
140 1 n, err := x.idx.DocCount()
141 1 if err != nil {
142 0 return 0, fmt.Errorf("search: count documents: %w", err)
143 0 }
144 1 return n, nil
145 }
146
147 // RebuildSpace replaces everything the index holds for one space.
148 //
149 // This is a rebuild, not an incremental update: whatever was indexed for sp is
150 // removed and docs are written in its place, with no diffing of individual
151 // documents and no per-document staleness bookkeeping. The unit of freshness is
152 // a space at a revision, which is exactly what the index_stamp row in Postgres
153 // records — this package is handed a revision's worth of documents and does not
154 // know or care which of them changed.
155 //
156 // docs may be empty, which empties the space. Every document must belong to sp;
157 // a document from another space is a caller bug and is refused rather than
158 // written somewhere surprising.
159 41 func (x *Index) RebuildSpace(ctx context.Context, sp core.SpaceRef, docs []Document) (Stats, error) {
160 41 start := time.Now()
161 41 if sp.Owner == "" || sp.Name == "" {
162 1 return Stats{}, errors.New("search: RebuildSpace needs a space")
163 1 }
164 40 fields := make([]map[string]any, len(docs))
165 448 for i, d := range docs {
166 448 if d.Space != sp {
167 1 return Stats{}, fmt.Errorf("search: document %q belongs to space %s, not %s", d.ID, d.Space, sp)
168 1 }
169 447 f, err := bleveDoc(d)
170 447 if err != nil {
171 0 return Stats{}, err
172 0 }
173 447 fields[i] = f
174 }
175
176 39 x.mu.Lock()
177 39 defer x.mu.Unlock()
178 39 if x.idx == nil {
179 1 return Stats{}, errors.New("search: index is closed")
180 1 }
181
182 38 stale, err := x.keysOf(sp)
183 38 if err != nil {
184 0 return Stats{}, err
185 0 }
186 38 if err := ctx.Err(); err != nil {
187 1 return Stats{}, err
188 1 }
189
190 // One batch for the whole space, so a space is never half-replaced. Deletes
191 // go in first: a batch is keyed by document id and the last operation on a
192 // key wins, so a document that survives the rebuild is re-indexed rather
193 // than dropped.
194 37 batch := x.idx.NewBatch()
195 37 kept := 0
196 353 for key := range stale {
197 353 batch.Delete(key)
198 353 }
199 446 for i, d := range docs {
200 446 key := Key(d.Space, d.ID)
201 446 if _, ok := stale[key]; ok {
202 351 kept++
203 351 }
204 446 if err := batch.Index(key, fields[i]); err != nil {
205 0 return Stats{}, fmt.Errorf("search: stage %s: %w", key, err)
206 0 }
207 }
208 37 if err := x.idx.Batch(batch); err != nil {
209 0 return Stats{}, fmt.Errorf("search: rebuild space %s: %w", sp, err)
210 0 }
211
212 37 st := Stats{Spaces: 1, Indexed: len(docs), Deleted: len(stale) - kept, Took: time.Since(start)}
213 37 if len(docs) == 0 {
214 1 st.Spaces = 0
215 1 }
216 37 return st, nil
217 }
218
219 // DeleteSpace removes every document of a space from the index. It is what a
220 // deleted space calls; RebuildSpace with no documents does the same thing.
221 1 func (x *Index) DeleteSpace(ctx context.Context, sp core.SpaceRef) (Stats, error) {
222 1 return x.RebuildSpace(ctx, sp, nil)
223 1 }
224
225 // RebuildAll replaces the entire index with docs, which may span any number of
226 // spaces.
227 //
228 // The rebuild runs into a fresh index beside the live one and the two are
229 // swapped at the end, so a failure part-way through leaves the old index intact
230 // and serving. That is the one thing warren's `bleve.New` over the live path
231 // did not give, and it matters here because the index is open in a daemon that
232 // is answering queries while the rebuild runs.
233 4 func (x *Index) RebuildAll(ctx context.Context, docs []Document) (Stats, error) {
234 4 start := time.Now()
235 4 spaces := make(map[core.SpaceRef]struct{})
236 4 fields := make([]map[string]any, len(docs))
237 1004 for i, d := range docs {
238 1004 f, err := bleveDoc(d)
239 1004 if err != nil {
240 0 return Stats{}, err
241 0 }
242 1004 fields[i] = f
243 1004 spaces[d.Space] = struct{}{}
244 }
245
246 4 tmp := x.path + ".rebuilding"
247 4 if err := os.RemoveAll(tmp); err != nil {
248 0 return Stats{}, fmt.Errorf("search: clear %s: %w", tmp, err)
249 0 }
250 4 fresh, err := bleve.New(tmp, buildMapping())
251 4 if err != nil {
252 0 return Stats{}, fmt.Errorf("search: create index at %s: %w", tmp, err)
253 0 }
254 4 if err := indexAll(ctx, fresh, docs, fields); err != nil {
255 2 _ = fresh.Close()
256 2 _ = os.RemoveAll(tmp)
257 2 return Stats{}, err
258 2 }
259 2 if err := fresh.Close(); err != nil {
260 0 _ = os.RemoveAll(tmp)
261 0 return Stats{}, fmt.Errorf("search: close rebuilt index: %w", err)
262 0 }
263
264 2 x.mu.Lock()
265 2 defer x.mu.Unlock()
266 2 if x.idx == nil {
267 0 _ = os.RemoveAll(tmp)
268 0 return Stats{}, errors.New("search: index is closed")
269 0 }
270 2 if err := x.swap(tmp); err != nil {
271 0 return Stats{}, err
272 0 }
273 2 return Stats{
274 2 Spaces: len(spaces),
275 2 Indexed: len(docs),
276 2 Took: time.Since(start),
277 2 }, nil
278 }
279
280 // swap puts the freshly built index at tmp in place of the live one. The caller
281 // holds the write lock.
282 //
283 // The old directory is renamed aside rather than deleted first, so the window in
284 // which neither exists is a rename rather than a recursive delete. If reopening
285 // the new index fails the Index is left closed and the error is returned: the
286 // index is a cache, and a caller that cannot open it must rebuild it, not
287 // silently serve an empty one.
288 2 func (x *Index) swap(tmp string) error {
289 2 if err := x.idx.Close(); err != nil {
290 0 x.idx = nil
291 0 return fmt.Errorf("search: close live index: %w", err)
292 0 }
293 2 x.idx = nil
294 2
295 2 old := x.path + ".old"
296 2 if err := os.RemoveAll(old); err != nil {
297 0 return fmt.Errorf("search: clear %s: %w", old, err)
298 0 }
299 2 if err := os.Rename(x.path, old); err != nil && !errors.Is(err, os.ErrNotExist) {
300 0 return fmt.Errorf("search: move live index aside: %w", err)
301 0 }
302 2 if err := os.Rename(tmp, x.path); err != nil {
303 0 return fmt.Errorf("search: move rebuilt index into place: %w", err)
304 0 }
305 2 idx, err := bleve.Open(x.path)
306 2 if err != nil {
307 0 return fmt.Errorf("search: reopen index at %s: %w", x.path, err)
308 0 }
309 2 x.idx = idx
310 2 if err := os.RemoveAll(old); err != nil {
311 0 return fmt.Errorf("search: remove %s: %w", old, err)
312 0 }
313 2 return nil
314 }
315
316 4 func indexAll(ctx context.Context, idx bleve.Index, docs []Document, fields []map[string]any) error {
317 4 batch := idx.NewBatch()
318 1004 for i, d := range docs {
319 1004 if err := ctx.Err(); err != nil {
320 2 return err
321 2 }
322 1002 key := Key(d.Space, d.ID)
323 1002 if err := batch.Index(key, fields[i]); err != nil {
324 0 return fmt.Errorf("search: stage %s: %w", key, err)
325 0 }
326 1002 if batch.Size() >= batchSize {
327 5 if err := idx.Batch(batch); err != nil {
328 0 return fmt.Errorf("search: flush batch: %w", err)
329 0 }
330 5 batch = idx.NewBatch()
331 }
332 }
333 2 if batch.Size() > 0 {
334 1 if err := idx.Batch(batch); err != nil {
335 0 return fmt.Errorf("search: flush final batch: %w", err)
336 0 }
337 }
338 2 return nil
339 }
340
341 // keysOf lists the index keys currently held for a space. The caller holds a
342 // lock.
343 38 func (x *Index) keysOf(sp core.SpaceRef) (map[string]struct{}, error) {
344 38 q := bleve.NewTermQuery(sp.String())
345 38 q.SetField(fieldSpace)
346 38
347 38 const page = 1000
348 38 keys := make(map[string]struct{})
349 38 for from := 0; ; from += page {
350 38 req := bleve.NewSearchRequestOptions(q, page, from, false)
351 38 res, err := x.idx.Search(req)
352 38 if err != nil {
353 0 return nil, fmt.Errorf("search: list documents of %s: %w", sp, err)
354 0 }
355 353 for _, h := range res.Hits {
356 353 keys[h.ID] = struct{}{}
357 353 }
358 38 if len(res.Hits) < page {
359 38 return keys, nil
360 38 }
361 }
362 }
363
364 // tmpPaths are the working directories a rebuild uses, exposed only so a caller
365 // cleaning up after a crash knows what to look for.
366 5 func tmpPaths(path string) []string {
367 5 return []string{path + ".rebuilding", path + ".old"}
368 5 }
369
370 // CleanStale removes the working directories a crashed rebuild may have left
371 // behind next to the index. Safe to call at startup, before Open.
372 1 func CleanStale(path string) error {
373 2 for _, p := range tmpPaths(path) {
374 2 if err := os.RemoveAll(p); err != nil {
375 0 return fmt.Errorf("search: remove %s: %w", filepath.Base(p), err)
376 0 }
377 }
378 1 return nil
379 }