tstor/pkg/ioutils/cachereader.go

package ioutils

import (
	"context"
	"errors"
	"io"
	"sync"

	"github.com/royalcat/ctxio"
)

type FileReader interface {
	ctxio.ReaderAt
	ctxio.Reader
	ctxio.Closer
}

type CacheReader struct {
	m sync.Mutex

	fo int64
	fr *FileBuffer
	to int64
	tr ctxio.Reader
}

var _ FileReader = (*CacheReader)(nil)

func NewCacheReader(r ctxio.Reader) (FileReader, error) {
	fr := NewFileBuffer(nil)
	tr := ctxio.TeeReader(r, fr)
	return &CacheReader{fr: fr, tr: tr}, nil
}

func (dtr *CacheReader) ReadAt(ctx context.Context, p []byte, off int64) (int, error) {
	dtr.m.Lock()
	defer dtr.m.Unlock()
	tb := off + int64(len(p))

	if tb > dtr.fo {
		w, err := ctxio.CopyN(ctx, ctxio.Discard, dtr.tr, tb-dtr.fo)
		dtr.to += w
		if err != nil && err != io.EOF {
			return 0, err
		}
	}

	n, err := dtr.fr.ReadAt(ctx, p, off)
	dtr.fo += int64(n)
	return n, err
}

func (dtr *CacheReader) Read(ctx context.Context, p []byte) (n int, err error) {
	dtr.m.Lock()
	defer dtr.m.Unlock()
	// use directly tee reader here
	n, err = dtr.tr.Read(ctx, p)
	dtr.to += int64(n)
	return
}

func (dtr *CacheReader) Close(ctx context.Context) error {
	frcloser := dtr.fr.Close(ctx)

	var closeerr error
	if rc, ok := dtr.tr.(ctxio.ReadCloser); ok {
		closeerr = rc.Close(ctx)
	}

	return errors.Join(frcloser, closeerr)
}