tstor/pkg/ioutils/disk.go
royalcat bd75492b02
Some checks failed
docker / build-docker (linux/386) (push) Failing after 18s
docker / build-docker (linux/amd64) (push) Failing after 17s
docker / build-docker (linux/arm/v7) (push) Failing after 17s
docker / build-docker (linux/arm64) (push) Failing after 17s
docker / build-docker (linux/arm64/v8) (push) Failing after 15s
refactor
2024-06-02 22:53:33 +03:00

72 lines
1.4 KiB
Go

package ioutils
import (
"context"
"io"
"os"
"sync"
"github.com/royalcat/ctxio"
)
type DiskCacheReader struct {
m sync.Mutex
fo int64
fr *os.File
to int64
tr ctxio.Reader
}
var _ ctxio.ReaderAt = (*DiskCacheReader)(nil)
var _ ctxio.Reader = (*DiskCacheReader)(nil)
var _ ctxio.Closer = (*DiskCacheReader)(nil)
func NewDiskCacheReader(r ctxio.Reader) (*DiskCacheReader, error) {
tempDir, err := os.MkdirTemp("/tmp", "tstor")
if err != nil {
return nil, err
}
fr, err := os.CreateTemp(tempDir, "dtb_tmp")
if err != nil {
return nil, err
}
tr := ctxio.TeeReader(r, ctxio.WrapIoWriter(fr))
return &DiskCacheReader{fr: fr, tr: tr}, nil
}
func (dtr *DiskCacheReader) 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(p, off)
dtr.fo += int64(n)
return n, err
}
func (dtr *DiskCacheReader) 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 *DiskCacheReader) Close(ctx context.Context) error {
if err := dtr.fr.Close(); err != nil {
return err
}
return os.Remove(dtr.fr.Name())
}