loki/pkg/storage/bloom/v1/reader.go

package v1

import (
	"bytes"
	"io"
	"os"
	"path/filepath"

	"github.com/pkg/errors"

	"github.com/grafana/dskit/multierror"

	iter "github.com/grafana/loki/v3/pkg/iter/v2"
)

type BlockReader interface {
	Index() (io.ReadSeeker, error)
	Blooms() (io.ReadSeeker, error)
	TarEntries() (iter.Iterator[TarEntry], error)
	Cleanup() error
}

// In memory reader
type ByteReader struct {
	index, blooms *bytes.Buffer
}

func NewByteReader(index, blooms *bytes.Buffer) *ByteReader {
	return &ByteReader{index: index, blooms: blooms}
}

func (r *ByteReader) Index() (io.ReadSeeker, error) {
	return bytes.NewReader(r.index.Bytes()), nil
}

func (r *ByteReader) Blooms() (io.ReadSeeker, error) {
	return bytes.NewReader(r.blooms.Bytes()), nil
}

func (r *ByteReader) TarEntries() (iter.Iterator[TarEntry], error) {
	indexLn := r.index.Len()
	index, err := r.Index()
	if err != nil {
		return nil, err
	}
	bloomLn := r.blooms.Len()
	blooms, err := r.Blooms()
	if err != nil {
		return nil, err
	}
	entries := []TarEntry{
		{
			Name: SeriesFileName,
			Size: int64(indexLn),
			Body: index,
		},
		{
			Name: BloomFileName,
			Size: int64(bloomLn),
			Body: blooms,
		},
	}

	return iter.NewSliceIter(entries), err
}

func (r *ByteReader) Cleanup() error {
	r.index.Reset()
	r.blooms.Reset()
	return nil
}

// File reader
type DirectoryBlockReader struct {
	dir           string
	blooms, index *os.File

	initialized bool
}

func NewDirectoryBlockReader(dir string) *DirectoryBlockReader {
	return &DirectoryBlockReader{
		dir:         dir,
		initialized: false,
	}
}

func (r *DirectoryBlockReader) Init() error {
	if !r.initialized {
		var err error
		r.index, err = os.Open(filepath.Join(r.dir, SeriesFileName))
		if err != nil {
			return errors.Wrap(err, "opening series file")
		}

		r.blooms, err = os.Open(filepath.Join(r.dir, BloomFileName))
		if err != nil {
			return errors.Wrap(err, "opening bloom file")
		}

		r.initialized = true
	}
	return nil
}

func (r *DirectoryBlockReader) Index() (io.ReadSeeker, error) {
	if !r.initialized {
		if err := r.Init(); err != nil {
			return nil, err
		}
	}
	return r.index, nil
}

func (r *DirectoryBlockReader) Blooms() (io.ReadSeeker, error) {
	if !r.initialized {
		if err := r.Init(); err != nil {
			return nil, err
		}
	}
	return r.blooms, nil
}

func (r *DirectoryBlockReader) TarEntries() (iter.Iterator[TarEntry], error) {
	var err error
	if !r.initialized {
		if err = r.Init(); err != nil {
			return nil, err
		}
	}

	_, err = r.index.Seek(0, io.SeekStart)
	if err != nil {
		return nil, errors.Wrap(err, "error seeking series file")
	}

	idxInfo, err := r.index.Stat()
	if err != nil {
		return nil, errors.Wrap(err, "error stat'ing series file")
	}

	_, err = r.blooms.Seek(0, io.SeekStart)
	if err != nil {
		return nil, errors.Wrap(err, "error seeking bloom file")
	}

	bloomInfo, err := r.blooms.Stat()
	if err != nil {
		return nil, errors.Wrap(err, "error stat'ing bloom file")
	}

	entries := []TarEntry{
		{
			Name: SeriesFileName,
			Size: idxInfo.Size(),
			Body: r.index,
		},
		{
			Name: BloomFileName,
			Size: bloomInfo.Size(),
			Body: r.blooms,
		},
	}

	return iter.NewSliceIter(entries), nil
}

func (r *DirectoryBlockReader) Cleanup() error {
	r.initialized = false
	err := multierror.New()
	err.Add(os.Remove(r.index.Name()))
	err.Add(os.Remove(r.blooms.Name()))
	err.Add(os.RemoveAll(r.dir))
	return err.Err()
}