tstor/pkg/ioutils/seeker.go

package ioutils

import (
	"context"
	"io"
	"sync"

	"github.com/royalcat/ctxio"
)

type ioSeekerWrapper struct {
	ctx context.Context

	mu   sync.Mutex
	pos  int64
	size int64

	r ctxio.ReaderAt
}

func WrapIoReadSeeker(ctx context.Context, r ctxio.ReaderAt, size int64) io.ReadSeeker {
	return &ioSeekerWrapper{
		ctx:  ctx,
		r:    r,
		size: size,
	}
}

func (r *ioSeekerWrapper) Seek(offset int64, whence int) (int64, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	switch whence {
	case io.SeekStart:
		r.pos = offset
	case io.SeekCurrent:
		r.pos = r.pos + offset
	case io.SeekEnd:
		r.pos = r.size + offset
	}

	return r.pos, nil
}

func (r *ioSeekerWrapper) Read(p []byte) (int, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	n, err := r.r.ReadAt(r.ctx, p, r.pos)
	r.pos += int64(n)

	return n, err
}

var _ io.ReadSeekCloser = (*ioSeekerCloserWrapper)(nil)

type ioSeekerCloserWrapper struct {
	ctx context.Context

	mu   sync.Mutex
	pos  int64
	size int64

	r FileReader
}

func IoReadSeekCloserWrapper(ctx context.Context, r FileReader, size int64) io.ReadSeekCloser {
	return &ioSeekerCloserWrapper{
		ctx:  ctx,
		r:    r,
		size: size,
	}
}

func (r *ioSeekerCloserWrapper) Seek(offset int64, whence int) (int64, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	switch whence {
	case io.SeekStart:
		r.pos = offset
	case io.SeekCurrent:
		r.pos = r.pos + offset
	case io.SeekEnd:
		r.pos = r.size + offset
	}

	return r.pos, nil
}

func (r *ioSeekerCloserWrapper) Read(p []byte) (int, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	n, err := r.r.ReadAt(r.ctx, p, r.pos)
	r.pos += int64(n)

	return n, err
}

// Close implements io.ReadSeekCloser.
func (r *ioSeekerCloserWrapper) Close() error {
	return r.r.Close(r.ctx)
}