Repository: prometheus/prometheus
File: tsdb/head_wal.go (Go)

Code before:

func (h *Head) loadWAL(r *wlog.Reader, syms *labels.SymbolTable, multiRef map[chunks.HeadSeriesRef]chunks.HeadSeriesRef, mmappedChunks, oooMmappedChunks map[chunks.HeadSeriesRef][]*mmappedChunk) (err error) {
	// Track number of samples that referenced a series we don't know about
	// for error reporting.
	var unknownRefs atomic.Uint64
	var unknownExemplarRefs atomic.Uint64
	var unknownHistogramRefs atomic.Uint64
	var unknownMetadataRefs atomic.Uint64
	// Track number of series records that had overlapping m-map chunks.
	var mmapOverlappingChunks atomic.Uint64

	// Start workers that each process samples for a partition of the series ID space.
	var (
		wg             sync.WaitGroup
		concurrency    = h.opts.WALReplayConcurrency
		processors     = make([]walSubsetProcessor, concurrency)
		exemplarsInput chan record.RefExemplar

		shards          = make([][]record.RefSample, concurrency)
		histogramShards = make([][]histogramRecord, concurrency)

		decoded                      = make(chan interface{}, 10)
		decodeErr, seriesCreationErr error
	)

	_, last, err := wlog.Segments(h.wal.Dir())
	if err != nil {
		return fmt.Errorf("failed to get last segment to set WAL expiry for duplicate series: %w", err)
	}

	defer func() {
		// For CorruptionErr ensure to terminate all workers before exiting.
		_, ok := err.(*wlog.CorruptionErr)
		if ok || seriesCreationErr != nil {
			for i := 0; i < concurrency; i++ {
				processors[i].closeAndDrain()
			}
			close(exemplarsInput)
			wg.Wait()
		}
	}()

	wg.Add(concurrency)
	for i := 0; i < concurrency; i++ {
		processors[i].setup()

		go func(wp *walSubsetProcessor) {
			unknown, unknownHistograms, overlapping := wp.processWALSamples(h, mmappedChunks, oooMmappedChunks)
			unknownRefs.Add(unknown)
			mmapOverlappingChunks.Add(overlapping)
			unknownHistogramRefs.Add(unknownHistograms)

Review comment (refactor):
nit: this is breaking the layers a bit - `r` is an abstract reader of one segment.
Maybe pass in `last` from the caller?

Code after:

func (h *Head) loadWAL(r *wlog.Reader, syms *labels.SymbolTable, multiRef map[chunks.HeadSeriesRef]chunks.HeadSeriesRef, mmappedChunks, oooMmappedChunks map[chunks.HeadSeriesRef][]*mmappedChunk, lastSegment int) (err error) {
	// Track number of samples that referenced a series we don't know about
	// for error reporting.
	var unknownRefs atomic.Uint64
	var unknownExemplarRefs atomic.Uint64
	var unknownHistogramRefs atomic.Uint64
	var unknownMetadataRefs atomic.Uint64
	// Track number of series records that had overlapping m-map chunks.
	var mmapOverlappingChunks atomic.Uint64

	// Start workers that each process samples for a partition of the series ID space.
	var (
		wg             sync.WaitGroup
		concurrency    = h.opts.WALReplayConcurrency
		processors     = make([]walSubsetProcessor, concurrency)
		exemplarsInput chan record.RefExemplar

		shards          = make([][]record.RefSample, concurrency)
		histogramShards = make([][]histogramRecord, concurrency)

		decoded                      = make(chan interface{}, 10)
		decodeErr, seriesCreationErr error
	)

	defer func() {
		// For CorruptionErr ensure to terminate all workers before exiting.
		_, ok := err.(*wlog.CorruptionErr)
		if ok || seriesCreationErr != nil {
			for i := 0; i < concurrency; i++ {
				processors[i].closeAndDrain()
			}
			close(exemplarsInput)
			wg.Wait()
		}
	}()

	wg.Add(concurrency)
	for i := 0; i < concurrency; i++ {
		processors[i].setup()

		go func(wp *walSubsetProcessor) {
			unknown, unknownHistograms, overlapping := wp.processWALSamples(h, mmappedChunks, oooMmappedChunks)
			unknownRefs.Add(unknown)
			mmapOverlappingChunks.Add(overlapping)
			unknownHistogramRefs.Add(unknownHistograms)
			wg.Done()
		}(&processors[i])
	}

	wg.Add(1)