github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/ruler/storage/wal/util.go (about)

     1  // This directory was copied and adapted from https://github.com/grafana/agent/tree/main/pkg/metrics.
     2  // We cannot vendor the agent in since the agent vendors loki in, which would cause a cyclic dependency.
     3  // NOTE: many changes have been made to the original code for our use-case.
     4  package wal
     5  
     6  import (
     7  	"path/filepath"
     8  	"sync"
     9  
    10  	"github.com/prometheus/prometheus/tsdb/record"
    11  	"github.com/prometheus/prometheus/tsdb/wal"
    12  )
    13  
    14  type walReplayer struct {
    15  	w wal.WriteTo
    16  }
    17  
    18  func (r walReplayer) Replay(dir string) error {
    19  	w, err := wal.Open(nil, dir)
    20  	if err != nil {
    21  		return err
    22  	}
    23  
    24  	dir, startFrom, err := wal.LastCheckpoint(w.Dir())
    25  	if err != nil && err != record.ErrNotFound {
    26  		return err
    27  	}
    28  
    29  	if err == nil {
    30  		sr, err := wal.NewSegmentsReader(dir)
    31  		if err != nil {
    32  			return err
    33  		}
    34  
    35  		err = r.replayWAL(wal.NewReader(sr))
    36  		if closeErr := sr.Close(); closeErr != nil && err == nil {
    37  			err = closeErr
    38  		}
    39  		if err != nil {
    40  			return err
    41  		}
    42  
    43  		startFrom++
    44  	}
    45  
    46  	_, last, err := wal.Segments(w.Dir())
    47  	if err != nil {
    48  		return err
    49  	}
    50  
    51  	for i := startFrom; i <= last; i++ {
    52  		s, err := wal.OpenReadSegment(wal.SegmentName(w.Dir(), i))
    53  		if err != nil {
    54  			return err
    55  		}
    56  
    57  		sr := wal.NewSegmentBufReader(s)
    58  		err = r.replayWAL(wal.NewReader(sr))
    59  		if closeErr := sr.Close(); closeErr != nil && err == nil {
    60  			err = closeErr
    61  		}
    62  		if err != nil {
    63  			return err
    64  		}
    65  	}
    66  
    67  	return nil
    68  }
    69  
    70  func (r walReplayer) replayWAL(reader *wal.Reader) error {
    71  	var dec record.Decoder
    72  
    73  	for reader.Next() {
    74  		rec := reader.Record()
    75  		switch dec.Type(rec) {
    76  		case record.Series:
    77  			series, err := dec.Series(rec, nil)
    78  			if err != nil {
    79  				return err
    80  			}
    81  			r.w.StoreSeries(series, 0)
    82  		case record.Samples:
    83  			samples, err := dec.Samples(rec, nil)
    84  			if err != nil {
    85  				return err
    86  			}
    87  			r.w.Append(samples)
    88  		case record.Exemplars:
    89  			exemplars, err := dec.Exemplars(rec, nil)
    90  			if err != nil {
    91  				return err
    92  			}
    93  			r.w.AppendExemplars(exemplars)
    94  		}
    95  	}
    96  
    97  	return nil
    98  }
    99  
   100  type walDataCollector struct {
   101  	mut       sync.Mutex
   102  	samples   []record.RefSample
   103  	series    []record.RefSeries
   104  	exemplars []record.RefExemplar
   105  }
   106  
   107  func (c *walDataCollector) AppendExemplars(exemplars []record.RefExemplar) bool {
   108  	c.mut.Lock()
   109  	defer c.mut.Unlock()
   110  
   111  	c.exemplars = append(c.exemplars, exemplars...)
   112  	return true
   113  }
   114  
   115  func (c *walDataCollector) Append(samples []record.RefSample) bool {
   116  	c.mut.Lock()
   117  	defer c.mut.Unlock()
   118  
   119  	c.samples = append(c.samples, samples...)
   120  	return true
   121  }
   122  
   123  func (c *walDataCollector) StoreSeries(series []record.RefSeries, _ int) {
   124  	c.mut.Lock()
   125  	defer c.mut.Unlock()
   126  
   127  	c.series = append(c.series, series...)
   128  }
   129  
   130  func (c *walDataCollector) UpdateSeriesSegment(series []record.RefSeries, index int) {}
   131  
   132  func (c *walDataCollector) SeriesReset(_ int) {}
   133  
   134  // SubDirectory returns the subdirectory within a Storage directory used for
   135  // the Prometheus WAL.
   136  func SubDirectory(base string) string {
   137  	return filepath.Join(base, "wal")
   138  }