github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/morebufio/readerat.go (about)

     1  package morebufio
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/Schaudge/grailbase/errors"
     9  	"github.com/Schaudge/grailbase/ioctx"
    10  )
    11  
    12  type readerAt struct {
    13  	r ioctx.ReaderAt
    14  	// mu guards updates to all subsequent fields. It's not held during reads.
    15  	mu sync.Mutex
    16  	// seekerInUse is true while waiting for operations (like reads) on seeker.
    17  	seekerInUse bool
    18  	seeker      *readSeeker
    19  }
    20  
    21  // NewReaderAt constructs a buffered ReaderAt. While ReaderAt allows arbitrary concurrent reads,
    22  // this implementation has only one buffer (for simplicity), so reads will only be usefully buffered
    23  // if reads are serial and generally contiguous (just like a plain ioctx.Reader). Concurrent reads
    24  // will be "passed through" to the underlying ReaderAt, just without buffering.
    25  func NewReaderAtSize(r ioctx.ReaderAt, size int) ioctx.ReaderAt {
    26  	return &readerAt{
    27  		r:      r,
    28  		seeker: NewReadSeekerSize(&readerAtSeeker{r: r}, size),
    29  	}
    30  }
    31  
    32  // ReadAt implements ioctx.ReaderAt.
    33  func (r *readerAt) ReadAt(ctx context.Context, dst []byte, off int64) (int, error) {
    34  	acquired := r.tryAcquireSeeker()
    35  	if !acquired {
    36  		return r.r.ReadAt(ctx, dst, off)
    37  	}
    38  	defer r.releaseSeeker()
    39  	if _, err := r.seeker.Seek(ctx, off, io.SeekStart); err != nil {
    40  		return 0, errors.E(err, "seeking for ReadAt")
    41  	}
    42  	return r.seeker.Read(ctx, dst)
    43  }
    44  
    45  func (r *readerAt) tryAcquireSeeker() bool {
    46  	r.mu.Lock()
    47  	defer r.mu.Unlock()
    48  	if r.seekerInUse {
    49  		return false
    50  	}
    51  	r.seekerInUse = true
    52  	return true
    53  }
    54  
    55  func (r *readerAt) releaseSeeker() {
    56  	r.mu.Lock()
    57  	defer r.mu.Unlock()
    58  	if !r.seekerInUse {
    59  		panic("release of unacquired seeker")
    60  	}
    61  	r.seekerInUse = false
    62  }
    63  
    64  // readerAtSeeker is a simple ioctx.ReadSeeker that only supports seeking by io.SeekCurrent,
    65  // which is all readSeeker requires.
    66  type readerAtSeeker struct {
    67  	r   ioctx.ReaderAt
    68  	pos int64
    69  }
    70  
    71  func (r *readerAtSeeker) Read(ctx context.Context, p []byte) (int, error) {
    72  	n, err := r.r.ReadAt(ctx, p, r.pos)
    73  	r.pos += int64(n)
    74  	return n, err
    75  }
    76  
    77  func (r *readerAtSeeker) Seek(ctx context.Context, request int64, whence int) (int64, error) {
    78  	if whence == io.SeekCurrent {
    79  		r.pos += request
    80  		return r.pos, nil
    81  	}
    82  	// Pretend the end position is zero. readSeeker requests this at initialization but we
    83  	// won't use it after that.
    84  	r.pos = 0
    85  	return r.pos, nil
    86  }