sorted.go raw
1 package store
2
3 import (
4 "bufio"
5 "bytes"
6 "fmt"
7 "io"
8 "os"
9 "path/filepath"
10 "sort"
11 "strings"
12 )
13
14 // sortedFile is a fixed-width, sorted flat file with binary search.
15 //
16 // Records reach it through three sources:
17 //
18 // memory - appended by Put since the last flush;
19 // runs - immutable, individually sorted files (path.buf.NNNNNN);
20 // main - the .dat file, produced by merging all runs into it.
21 //
22 // Each flush creates a new run instead of appending to one shared sidecar.
23 // That distinction matters: two independently sorted batches appended to one
24 // file are not themselves sorted, so binary search over that file silently
25 // misses records. Separate run files keep every source sorted, and a merge
26 // turns all of them into the main file. This is the LSM shape pogreb-style
27 // engines use, and the layout morly's index files follow.
28 type sortedFile struct {
29 path string
30 recLen int
31 cmpLen int
32
33 f *os.File
34 count int64 // records in the main file
35
36 buf []byte
37 bufN int
38 sorted bool
39
40 runs []runRef
41 runData [][]byte
42 runSeq int
43 del map[string]struct{}
44
45 mergeThresh int64
46 bulk bool
47 serOff int
48 removed func(uint64) bool
49 }
50
51 type runRef struct {
52 path string
53 count int64
54 }
55
56 // flushAt bounds the memory buffer between explicit flushes.
57 const flushAt = 8 << 20
58
59 func openSorted(path string, recLen, cmpLen int) (*sortedFile, error) {
60 os.Remove(path + ".tmp")
61 f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o644)
62 if err != nil {
63 return nil, err
64 }
65 info, err := f.Stat()
66 if err != nil {
67 f.Close()
68 return nil, err
69 }
70 size := info.Size()
71 count := size / int64(recLen)
72 if size%int64(recLen) != 0 {
73 // A torn tail cannot be part of a complete record; the index can
74 // always be rebuilt from the WAL.
75 if err := f.Truncate(count * int64(recLen)); err != nil {
76 f.Close()
77 return nil, err
78 }
79 }
80
81 s := &sortedFile{
82 path: path,
83 recLen: recLen,
84 cmpLen: cmpLen,
85 f: f,
86 count: count,
87 sorted: true,
88 del: make(map[string]struct{}),
89 }
90 if err := s.discoverRuns(); err != nil {
91 f.Close()
92 return nil, err
93 }
94
95 thresh := size / 10
96 if thresh < 1<<20 {
97 thresh = 1 << 20
98 }
99 s.mergeThresh = thresh
100 return s, nil
101 }
102
103 // discoverRuns finds the immutable run files, oldest first.
104 func (s *sortedFile) discoverRuns() error {
105 names, err := filepath.Glob(s.path + ".buf*")
106 if err != nil {
107 return err
108 }
109 sort.Strings(names)
110 for _, name := range names {
111 info, err := os.Stat(name)
112 if err != nil {
113 continue
114 }
115 size := info.Size()
116 count := size / int64(s.recLen)
117 if size == 0 {
118 os.Remove(name)
119 continue
120 }
121 if size%int64(s.recLen) != 0 {
122 // A partially written run is the tail of a crash; drop the
123 // incomplete record so the run stays searchable.
124 if err := os.Truncate(name, count*int64(s.recLen)); err != nil {
125 return err
126 }
127 }
128 s.runs = append(s.runs, runRef{path: name, count: count})
129 if suffix := strings.TrimPrefix(name, s.path+".buf."); suffix != name {
130 var n int
131 if _, err := fmt.Sscanf(suffix, "%d", &n); err == nil && n >= s.runSeq {
132 s.runSeq = n + 1
133 }
134 } else if s.runSeq == 0 {
135 s.runSeq = 1
136 }
137 }
138 s.runData = make([][]byte, len(s.runs))
139 return nil
140 }
141
142 // Put appends a record. The record is copied.
143 func (s *sortedFile) Put(rec []byte) error {
144 if len(rec) != s.recLen {
145 panic("store: bad record length")
146 }
147 // A later Put cancels an earlier Delete for the same key.
148 if len(s.del) > 0 {
149 delete(s.del, string(rec[:s.cmpLen]))
150 }
151 s.buf = append(s.buf, rec...)
152 s.bufN++
153 s.sorted = false
154 if len(s.buf) >= flushAt {
155 return s.QuickFlush(false)
156 }
157 return nil
158 }
159
160 // Delete marks a key (its first cmpLen bytes) as removed.
161 func (s *sortedFile) Delete(key []byte) {
162 k := make([]byte, s.cmpLen)
163 copy(k, key)
164 s.del[string(k)] = struct{}{}
165 }
166
167 func (s *sortedFile) isDeleted(rec []byte) bool {
168 if len(s.del) > 0 {
169 if _, ok := s.del[string(rec[:s.cmpLen])]; ok {
170 return true
171 }
172 }
173 if s.removed != nil && s.serOff >= 0 && len(rec) >= s.serOff+SerialLen {
174 return s.removed(GetSerial(rec[s.serOff : s.serOff+SerialLen]))
175 }
176 return false
177 }
178
179 func (s *sortedFile) Dirty() bool { return s.bufN > 0 || len(s.del) > 0 || len(s.runs) > 0 }
180
181 func (s *sortedFile) Count() int64 { return s.count + int64(s.bufN) + s.runCount() }
182
183 func (s *sortedFile) runCount() int64 {
184 var n int64
185 for _, r := range s.runs {
186 n += r.count
187 }
188 return n
189 }
190
191 func (s *sortedFile) runBytes() int64 {
192 var n int64
193 for _, r := range s.runs {
194 n += r.count * int64(s.recLen)
195 }
196 return n
197 }
198
199 type recSorter struct {
200 data []byte
201 recLen int
202 cmpLen int
203 tmp []byte
204 }
205
206 func (r *recSorter) Len() int { return len(r.data) / r.recLen }
207 func (r *recSorter) Less(i, j int) bool {
208 a := r.data[i*r.recLen : i*r.recLen+r.cmpLen]
209 b := r.data[j*r.recLen : j*r.recLen+r.cmpLen]
210 return bytes.Compare(a, b) < 0
211 }
212 func (r *recSorter) Swap(i, j int) {
213 a := r.data[i*r.recLen : (i+1)*r.recLen]
214 b := r.data[j*r.recLen : (j+1)*r.recLen]
215 copy(r.tmp, a)
216 copy(a, b)
217 copy(b, r.tmp)
218 }
219
220 func (s *sortedFile) ensureSorted() {
221 if s.sorted || s.bufN <= 1 {
222 s.sorted = true
223 return
224 }
225 sort.Sort(&recSorter{s.buf, s.recLen, s.cmpLen, make([]byte, s.recLen)})
226 s.sorted = true
227 }
228
229 // QuickFlush writes the memory buffer as a new immutable run. sync controls
230 // fsync, which is what makes the checkpoint safe to advance.
231 func (s *sortedFile) QuickFlush(sync bool) error {
232 if s.bufN == 0 {
233 return nil
234 }
235 s.ensureSorted()
236 name := fmt.Sprintf("%s.buf.%06d", s.path, s.runSeq)
237 s.runSeq++
238 tmp := name + ".tmp"
239 f, err := os.OpenFile(tmp, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0o644)
240 if err != nil {
241 return err
242 }
243 if _, err := f.Write(s.buf[:s.bufN*s.recLen]); err != nil {
244 f.Close()
245 return err
246 }
247 if sync {
248 if err := f.Sync(); err != nil {
249 f.Close()
250 return err
251 }
252 }
253 if err := f.Close(); err != nil {
254 return err
255 }
256 if err := os.Rename(tmp, name); err != nil {
257 return err
258 }
259 count := int64(s.bufN)
260 s.runs = append(s.runs, runRef{path: name, count: count})
261 if s.bulk {
262 // Bulk imports create many runs and defer every merge to Flush.
263 // Keeping their bytes resident would make memory grow with the whole
264 // import; loadRuns reads them back when a scan actually needs them.
265 s.runData = append(s.runData, nil)
266 } else {
267 data := make([]byte, len(s.buf[:s.bufN*s.recLen]))
268 copy(data, s.buf[:s.bufN*s.recLen])
269 s.runData = append(s.runData, data)
270 }
271 s.buf = s.buf[:0]
272 s.bufN = 0
273 s.sorted = true
274
275 if s.bulk {
276 // Bulk imports write one sorted run per buffer and defer the single
277 // merge to Flush. Merging every run as the main file grows turns a
278 // linear import into repeated rewrites of the whole index.
279 return nil
280 }
281 if s.runBytes() >= s.mergeThresh {
282 return s.Flush()
283 }
284 return nil
285 }
286
287 // loadRuns reads every run into memory once. Runs are bounded by the merge
288 // threshold, so this is bounded memory, and the cache avoids re-reading the
289 // same run for every point lookup in a query.
290 func (s *sortedFile) loadRuns() error {
291 if len(s.runData) == len(s.runs) {
292 all := true
293 for _, d := range s.runData {
294 if d == nil {
295 all = false
296 break
297 }
298 }
299 if all {
300 return nil
301 }
302 }
303 if len(s.runData) != len(s.runs) {
304 s.runData = make([][]byte, len(s.runs))
305 }
306 for i, r := range s.runs {
307 if s.runData[i] != nil {
308 continue
309 }
310 data := make([]byte, r.count*int64(s.recLen))
311 f, err := os.Open(r.path)
312 if err != nil {
313 return err
314 }
315 _, err = f.ReadAt(data, 0)
316 f.Close()
317 if err != nil && err != io.EOF {
318 return err
319 }
320 s.runData[i] = data
321 }
322 return nil
323 }
324
325 func (s *sortedFile) cmpBound(rec, bound []byte) int {
326 n := s.cmpLen
327 if len(bound) < n {
328 n = len(bound)
329 }
330 return bytes.Compare(rec[:n], bound[:n])
331 }
332
333 func (s *sortedFile) readAt(dst []byte, idx int64) error {
334 _, err := s.f.ReadAt(dst, idx*int64(s.recLen))
335 return err
336 }
337
338 // lowerBound returns the first main-file record >= key.
339 func (s *sortedFile) lowerBound(key []byte) int64 {
340 lo, hi := int64(0), s.count
341 rec := make([]byte, s.recLen)
342 for lo < hi {
343 mid := lo + (hi-lo)/2
344 if err := s.readAt(rec, mid); err != nil {
345 return s.count
346 }
347 if s.cmpBound(rec, key) < 0 {
348 lo = mid + 1
349 } else {
350 hi = mid
351 }
352 }
353 return lo
354 }
355
356 // Get returns the first record whose cmpLen prefix equals key.
357 func (s *sortedFile) Get(key []byte) ([]byte, bool, error) {
358 // Memory is newest, then runs newest-to-oldest, then the main file.
359 if s.bufN > 0 {
360 s.ensureSorted()
361 lo, hi := 0, s.bufN
362 for lo < hi {
363 mid := lo + (hi-lo)/2
364 rec := s.buf[mid*s.recLen:]
365 if s.cmpBound(rec, key) < 0 {
366 lo = mid + 1
367 } else {
368 hi = mid
369 }
370 }
371 if lo < s.bufN {
372 rec := s.buf[lo*s.recLen : (lo+1)*s.recLen]
373 if s.cmpBound(rec, key) == 0 && !s.isDeleted(rec) {
374 out := make([]byte, s.recLen)
375 copy(out, rec)
376 return out, true, nil
377 }
378 }
379 }
380 if err := s.loadRuns(); err != nil {
381 return nil, false, err
382 }
383 for i := len(s.runs) - 1; i >= 0; i-- {
384 data := s.runData[i]
385 n := s.runs[i].count
386 lo, hi := int64(0), n
387 for lo < hi {
388 mid := lo + (hi-lo)/2
389 rec := data[mid*int64(s.recLen):]
390 if s.cmpBound(rec, key) < 0 {
391 lo = mid + 1
392 } else {
393 hi = mid
394 }
395 }
396 if lo < n {
397 rec := data[lo*int64(s.recLen):]
398 if s.cmpBound(rec, key) == 0 && !s.isDeleted(rec) {
399 out := make([]byte, s.recLen)
400 copy(out, rec[:s.recLen])
401 return out, true, nil
402 }
403 }
404 }
405 idx := s.lowerBound(key)
406 if idx < s.count {
407 rec := make([]byte, s.recLen)
408 if err := s.readAt(rec, idx); err != nil {
409 return nil, false, err
410 }
411 if s.cmpBound(rec, key) == 0 {
412 if s.isDeleted(rec) {
413 return nil, false, nil
414 }
415 return rec, true, nil
416 }
417 }
418 return nil, false, nil
419 }
420
421 // Scan emits every record in [start, end) in key order. A nil end means the
422 // end of the index. fn returning false stops the scan.
423 func (s *sortedFile) Scan(start, end []byte, fn func(rec []byte) bool) error {
424 if s.bufN > 0 {
425 s.ensureSorted()
426 }
427 if err := s.loadRuns(); err != nil {
428 return err
429 }
430
431 var recs [][]byte
432 inRange := func(rec []byte) bool {
433 if start != nil && s.cmpBound(rec, start) < 0 {
434 return false
435 }
436 if end != nil && s.cmpBound(rec, end) >= 0 {
437 return false
438 }
439 return true
440 }
441 add := func(rec []byte) {
442 if !inRange(rec) || s.isDeleted(rec) {
443 return
444 }
445 out := make([]byte, s.recLen)
446 copy(out, rec[:s.recLen])
447 recs = append(recs, out)
448 }
449
450 from := int64(0)
451 if start != nil {
452 from = s.lowerBound(start)
453 }
454 chunkRecs := 256
455 chunk := make([]byte, s.recLen*chunkRecs)
456 pastEnd := func(rec []byte) bool { return end != nil && s.cmpBound(rec, end) >= 0 }
457
458 // The main file is already sorted. With no runs or buffered records there
459 // is nothing to merge, so stream records straight to fn. The general path
460 // below collects and sorts the entire range; for a broad scan that can be
461 // millions of records even when the caller stops after a few results.
462 if len(s.runs) == 0 && s.bufN == 0 {
463 for i := from; i < s.count; {
464 n := int64(chunkRecs)
465 if s.count-i < n {
466 n = s.count - i
467 }
468 if _, err := s.f.ReadAt(chunk[:int(n)*s.recLen], i*int64(s.recLen)); err != nil && err != io.EOF {
469 return err
470 }
471 for j := int64(0); j < n; j++ {
472 rec := chunk[j*int64(s.recLen) : (j+1)*int64(s.recLen)]
473 if pastEnd(rec) {
474 return nil
475 }
476 if s.isDeleted(rec) {
477 continue
478 }
479 if !fn(rec) {
480 return nil
481 }
482 }
483 i += n
484 }
485 return nil
486 }
487
488 for i := from; i < s.count; {
489 n := int64(chunkRecs)
490 if s.count-i < n {
491 n = s.count - i
492 }
493 if _, err := s.f.ReadAt(chunk[:int(n)*s.recLen], i*int64(s.recLen)); err != nil && err != io.EOF {
494 return err
495 }
496 stop := false
497 for j := int64(0); j < n; j++ {
498 rec := chunk[j*int64(s.recLen) : (j+1)*int64(s.recLen)]
499 if pastEnd(rec) {
500 stop = true
501 break
502 }
503 add(rec)
504 }
505 if stop {
506 break
507 }
508 i += n
509 }
510 for i := range s.runs {
511 data := s.runData[i]
512 for j := int64(0); j < s.runs[i].count; j++ {
513 rec := data[j*int64(s.recLen):]
514 if pastEnd(rec) {
515 break
516 }
517 add(rec)
518 }
519 }
520 for i := 0; i < s.bufN; i++ {
521 rec := s.buf[i*s.recLen : (i+1)*s.recLen]
522 if pastEnd(rec) {
523 break
524 }
525 add(rec)
526 }
527
528 if len(recs) == 0 {
529 return nil
530 }
531 sort.Slice(recs, func(i, j int) bool {
532 c := bytes.Compare(recs[i][:s.cmpLen], recs[j][:s.cmpLen])
533 if c != 0 {
534 return c < 0
535 }
536 return bytes.Compare(recs[i], recs[j]) < 0
537 })
538 var prev []byte
539 for _, rec := range recs {
540 if prev != nil && bytes.Equal(prev, rec) {
541 continue
542 }
543 prev = rec
544 if !fn(rec) {
545 return nil
546 }
547 }
548 return nil
549 }
550
551 // Flush merges the main file and all runs into a new main file, dropping
552 // deleted and duplicate records.
553 func (s *sortedFile) Flush() error {
554 if s.bufN > 0 {
555 if err := s.QuickFlush(false); err != nil {
556 return err
557 }
558 }
559 if len(s.runs) == 0 && len(s.del) == 0 {
560 return nil
561 }
562 if err := s.loadRuns(); err != nil {
563 return err
564 }
565
566 // Combine the runs into one sorted buffer. Their total size is bounded by
567 // mergeThresh, and the main file is streamed rather than loaded.
568 total := 0
569 for _, d := range s.runData {
570 total += len(d)
571 }
572 merged := make([]byte, 0, total)
573 for _, d := range s.runData {
574 merged = append(merged, d...)
575 }
576 if len(merged) > 0 {
577 sort.Sort(&recSorter{merged, s.recLen, s.cmpLen, make([]byte, s.recLen)})
578 }
579
580 tmpPath := s.path + ".tmp"
581 tmp, err := os.OpenFile(tmpPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0o644)
582 if err != nil {
583 return err
584 }
585 w := bufio.NewWriterSize(tmp, 1<<16)
586
587 const chunkRecs = 1024
588 mainBuf := make([]byte, chunkRecs*s.recLen)
589 mainPos := int64(0)
590 mainRec := func() ([]byte, error) {
591 if mainPos >= s.count {
592 return nil, nil
593 }
594 idx := mainPos % chunkRecs
595 if idx == 0 {
596 n := int64(chunkRecs)
597 if s.count-mainPos < n {
598 n = s.count - mainPos
599 }
600 if _, err := s.f.ReadAt(mainBuf[:int(n)*s.recLen], mainPos*int64(s.recLen)); err != nil && err != io.EOF {
601 return nil, err
602 }
603 }
604 rec := mainBuf[idx*int64(s.recLen) : (idx+1)*int64(s.recLen)]
605 mainPos++
606 return rec, nil
607 }
608 mergePos := int64(0)
609 mergeRec := func() []byte {
610 if mergePos >= int64(len(merged)) {
611 return nil
612 }
613 rec := merged[mergePos : mergePos+int64(s.recLen)]
614 mergePos += int64(s.recLen)
615 return rec
616 }
617
618 prev := make([]byte, 0, s.recLen)
619 emit := func(rec []byte) error {
620 if s.isDeleted(rec) {
621 return nil
622 }
623 if len(prev) == s.recLen && bytes.Equal(prev, rec) {
624 return nil
625 }
626 prev = append(prev[:0], rec...)
627 _, err := w.Write(rec)
628 return err
629 }
630
631 mr, err := mainRec()
632 if err != nil {
633 tmp.Close()
634 return err
635 }
636 ur := mergeRec()
637 for mr != nil || ur != nil {
638 var rec []byte
639 switch {
640 case mr == nil:
641 rec = ur
642 ur = mergeRec()
643 case ur == nil:
644 rec = mr
645 mr, err = mainRec()
646 if err != nil {
647 tmp.Close()
648 return err
649 }
650 default:
651 c := bytes.Compare(mr[:s.cmpLen], ur[:s.cmpLen])
652 switch {
653 case c < 0:
654 rec = mr
655 mr, err = mainRec()
656 if err != nil {
657 tmp.Close()
658 return err
659 }
660 case c > 0:
661 rec = ur
662 ur = mergeRec()
663 default:
664 if bytes.Compare(mr, ur) <= 0 {
665 rec = mr
666 mr, err = mainRec()
667 if err != nil {
668 tmp.Close()
669 return err
670 }
671 } else {
672 rec = ur
673 ur = mergeRec()
674 }
675 }
676 }
677 if err := emit(rec); err != nil {
678 tmp.Close()
679 return err
680 }
681 }
682
683 if err := w.Flush(); err != nil {
684 tmp.Close()
685 return err
686 }
687 if err := tmp.Sync(); err != nil {
688 tmp.Close()
689 return err
690 }
691 if err := tmp.Close(); err != nil {
692 return err
693 }
694 if err := s.f.Close(); err != nil {
695 return err
696 }
697 if err := os.Rename(tmpPath, s.path); err != nil {
698 return err
699 }
700 f, err := os.OpenFile(s.path, os.O_RDWR, 0o644)
701 if err != nil {
702 return err
703 }
704 s.f = f
705 info, err := f.Stat()
706 if err != nil {
707 return err
708 }
709 s.count = info.Size() / int64(s.recLen)
710
711 for _, r := range s.runs {
712 os.Remove(r.path)
713 }
714 s.runs = nil
715 s.runData = nil
716 s.del = make(map[string]struct{})
717 return nil
718 }
719
720 func (s *sortedFile) Close() error {
721 if err := s.Flush(); err != nil {
722 return err
723 }
724 if err := s.f.Sync(); err != nil {
725 return err
726 }
727 if err := s.f.Close(); err != nil {
728 return err
729 }
730 for _, r := range s.runs {
731 os.Remove(r.path)
732 }
733 return nil
734 }
735
736 // Clear removes every record from the main file and every run.
737 func (s *sortedFile) Clear() error {
738 s.buf = s.buf[:0]
739 s.bufN = 0
740 s.sorted = true
741 s.del = make(map[string]struct{})
742 for _, r := range s.runs {
743 os.Remove(r.path)
744 }
745 s.runs = nil
746 s.runData = nil
747 if err := s.f.Truncate(0); err != nil {
748 return err
749 }
750 s.count = 0
751 return nil
752 }
753