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  }