github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/tools/blocksconvert/builder/symbols_iterator.go (about) 1 package builder 2 3 import ( 4 "encoding/gob" 5 "io" 6 "os" 7 8 "github.com/golang/snappy" 9 tsdb_errors "github.com/prometheus/prometheus/tsdb/errors" 10 ) 11 12 type symbolsIterator struct { 13 files []*symbolsFile 14 errs []error 15 16 // To avoid returning duplicates, we remember last returned symbol. 17 lastReturned *string 18 } 19 20 func newSymbolsIterator(files []*symbolsFile) *symbolsIterator { 21 si := &symbolsIterator{ 22 files: files, 23 } 24 si.buildHeap() 25 return si 26 } 27 28 func (sit *symbolsIterator) buildHeap() { 29 // All files on the heap must have at least one element, so that "heapify" can order them. 30 // Here we verify that, and remove files with no more elements. 31 for ix := 0; ix < len(sit.files); { 32 f := sit.files[ix] 33 next, err := f.hasNext() 34 35 if err != nil { 36 sit.errs = append(sit.errs, err) 37 return 38 } 39 40 if !next { 41 if err := f.close(); err != nil { 42 sit.errs = append(sit.errs, err) 43 } 44 sit.files = append(sit.files[:ix], sit.files[ix+1:]...) 45 continue 46 } 47 48 ix++ 49 } 50 51 // Build heap, start with leaf nodes, and work towards to root. See comment at heapify for more details. 52 for ix := len(sit.files) - 1; ix >= 0; ix-- { 53 heapifySymbols(sit.files, ix) 54 } 55 } 56 57 // Next advances iterator forward, and returns next element. If there is no next element, returns false. 58 func (sit *symbolsIterator) Next() (string, bool) { 59 again: 60 if len(sit.errs) > 0 { 61 return "", false 62 } 63 64 if len(sit.files) == 0 { 65 return "", false 66 } 67 68 result := sit.files[0].pop() 69 70 hasNext, err := sit.files[0].hasNext() 71 if err != nil { 72 sit.errs = append(sit.errs, err) 73 } 74 75 if !hasNext { 76 if err := sit.files[0].close(); err != nil { 77 sit.errs = append(sit.errs, err) 78 } 79 80 // Move last file to the front, and heapify from there. 81 sit.files[0] = sit.files[len(sit.files)-1] 82 sit.files = sit.files[:len(sit.files)-1] 83 } 84 85 heapifySymbols(sit.files, 0) 86 87 if sit.lastReturned == nil || *sit.lastReturned != result { 88 sit.lastReturned = &result 89 return result, true 90 } 91 92 // Duplicate symbol, try next one. 93 goto again 94 } 95 96 func (sit *symbolsIterator) Error() error { 97 return tsdb_errors.NewMulti(sit.errs...).Err() 98 } 99 100 func (sit *symbolsIterator) Close() error { 101 errs := tsdb_errors.NewMulti() 102 for _, f := range sit.files { 103 errs.Add(f.close()) 104 } 105 return errs.Err() 106 } 107 108 func heapifySymbols(files []*symbolsFile, ix int) { 109 heapify(len(files), ix, func(i, j int) bool { 110 return files[i].peek() < files[j].peek() 111 }, func(i, j int) { 112 files[i], files[j] = files[j], files[i] 113 }) 114 } 115 116 type symbolsFile struct { 117 f *os.File 118 dec *gob.Decoder 119 120 next bool 121 nextSymbol string 122 } 123 124 func newSymbolsFile(f *os.File) *symbolsFile { 125 sn := snappy.NewReader(f) 126 dec := gob.NewDecoder(sn) 127 128 return &symbolsFile{ 129 f: f, 130 dec: dec, 131 } 132 } 133 134 func (sf *symbolsFile) close() error { 135 return sf.f.Close() 136 } 137 138 func (sf *symbolsFile) hasNext() (bool, error) { 139 if sf.next { 140 return true, nil 141 } 142 143 var s string 144 err := sf.dec.Decode(&s) 145 if err != nil { 146 if err == io.EOF { 147 return false, nil 148 } 149 return false, err 150 } 151 152 sf.next = true 153 sf.nextSymbol = s 154 return true, nil 155 } 156 157 func (sf *symbolsFile) peek() string { 158 if !sf.next { 159 panic("no next symbol") 160 } 161 162 return sf.nextSymbol 163 } 164 165 func (sf *symbolsFile) pop() string { 166 if !sf.next { 167 panic("no next symbol") 168 } 169 170 sf.next = false 171 return sf.nextSymbol 172 }