tstor/src/vfs/archive_cache.go

package vfs

import (
	"context"
	"errors"

	lru "github.com/hashicorp/golang-lru/v2"
	"github.com/royalcat/ctxio"
)

// TODO переделать кеш в демон

const blockSize int64 = 1024 * 16 // 16KB
const defaultBlockCount = 32768   // 512MB of total usage

type archiveFileIndex struct {
	archive  string
	filename string
}

type blockIndex struct {
	index archiveFileIndex
	off   int64
}

var blockCache *lru.Cache[blockIndex, []byte]

func ChangeBufferSize(blockCount int) {
	blockCache.Resize(blockCount)
}

func init() {
	var err error
	blockCache, err = lru.New[blockIndex, []byte](defaultBlockCount)
	if err != nil {
		panic(err)
	}
}

func newRandomReaderFromLinear(index archiveFileIndex, size int64, readerFactory archiveFileReaderFactory) *randomReaderFromLinear {
	return &randomReaderFromLinear{
		index:         index,
		size:          size,
		readerFactory: readerFactory,
	}
}

type randomReaderFromLinear struct {
	index         archiveFileIndex
	readerFactory archiveFileReaderFactory
	reader        ctxio.ReadCloser
	readen        int64
	size          int64
	closed        bool
}

var _ ctxio.ReaderAt = (*randomReaderFromLinear)(nil)
var _ ctxio.Closer = (*randomReaderFromLinear)(nil)

// ReadAt implements ctxio.ReaderAt.
func (a *randomReaderFromLinear) ReadAt(ctx context.Context, p []byte, off int64) (n int, err error) {
	ctx, span := tracer.Start(ctx, "archive.RandomReader.ReadAt")
	defer span.End()

	if a.closed {
		return 0, errors.New("reader is closed")
	}

	if off >= a.size {
		return 0, ctxio.EOF
	}

	aligntOff := (off / blockSize) * blockSize

	block, ok := blockCache.Get(blockIndex{index: a.index, off: aligntOff})
	if ok {
		n = copy(p, block[off-aligntOff:])
		if len(block) < int(blockSize) {
			err = ctxio.EOF
		}

		return n, err
	}

	span.AddEvent("cache miss, reading from file")
	if err := a.readTo(ctx, aligntOff+blockSize); err != nil && err != ctxio.EOF {
		return 0, err
	}

	block, ok = blockCache.Get(blockIndex{index: a.index, off: aligntOff})
	if !ok {
		// WTF this theoretically shouldn't happen under normal scenarios
		return 0, errors.New("block not found or block cache under too much pressure, try to increase the cache size")
	}

	n = copy(p, block[off-aligntOff:])
	if len(block) < int(blockSize) {
		err = ctxio.EOF
	}
	return n, err
}

func (a *randomReaderFromLinear) readTo(ctx context.Context, targetOffset int64) (err error) {
	if a.reader == nil || a.readen > targetOffset {
		a.reader, err = a.readerFactory(context.TODO())
		if err != nil {
			return err
		}
		a.readen = 0
	}

	for off := a.readen; off < targetOffset; off += blockSize {
		// TODO sync.Pool ?
		buf := make([]byte, blockSize)
		n, err := a.reader.Read(ctx, buf)
		if err != nil && err != ctxio.EOF {
			return err
		}
		a.readen += int64(n)
		if int64(n) < blockSize {
			buf = buf[:n]
		}

		blockCache.Add(blockIndex{index: a.index, off: off}, buf)
	}

	return nil
}

// Close implements ctxio.Closer.
func (a *randomReaderFromLinear) Close(ctx context.Context) error {
	if a.closed {
		return nil
	}
	a.closed = true

	var errs []error

	if a.reader != nil {
		errs = append(errs, a.reader.Close(ctx))
	}

	for _, block := range blockCache.Keys() {
		if block.index == a.index {
			blockCache.Remove(block)
		}
	}

	return errors.Join(errs...)
}