github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/tools/blocksconvert/builder/series_iterator.go (about)

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