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

     1  package morebufio
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"github.com/Schaudge/grailbase/ioctx"
     8  )
     9  
    10  type readSeeker struct {
    11  	r ioctx.ReadSeeker
    12  	// buf is the buffer, resized as necessary after reading from r.
    13  	buf []byte
    14  	// off is the caller's current offset into the buffer. buf[off:] is unread.
    15  	off int
    16  
    17  	// filePos is the caller's current position r's stream. This can be different from r's position,
    18  	// for example when there's unread data in buf. Equals -1 when uninitialized.
    19  	filePos int64
    20  	// fileEnd is the offset of the end of r, uused for efficiently seeking within r. Equals -1 when
    21  	// uninitialized.
    22  	fileEnd int64
    23  }
    24  
    25  var _ ioctx.ReadSeeker = (*readSeeker)(nil)
    26  
    27  // minBufferSize equals bufio.minBufferSize.
    28  const minBufferSize = 16
    29  
    30  // NewReadSeekerSize returns a buffered io.ReadSeeker whose buffer has at least the specified size.
    31  // If r is already a readSeeker with sufficient size, returns r.
    32  func NewReadSeekerSize(r ioctx.ReadSeeker, size int) *readSeeker {
    33  	if b, ok := r.(*readSeeker); ok && len(b.buf) >= size {
    34  		return b
    35  	}
    36  	if size < minBufferSize {
    37  		size = minBufferSize
    38  	}
    39  	return &readSeeker{r, make([]byte, 0, size), 0, -1, -1}
    40  }
    41  
    42  // Read implements ioctx.Reader.
    43  func (b *readSeeker) Read(ctx context.Context, p []byte) (int, error) {
    44  	if len(p) == 0 {
    45  		return 0, nil
    46  	}
    47  	if err := b.initFilePos(ctx); err != nil {
    48  		return 0, err
    49  	}
    50  	var err error
    51  	if b.off == len(b.buf) {
    52  		b.buf = b.buf[:cap(b.buf)]
    53  		var n int
    54  		n, err = b.r.Read(ctx, b.buf)
    55  		b.buf, b.off = b.buf[:n], 0
    56  	}
    57  	n := copy(p, b.buf[b.off:])
    58  	b.off += n
    59  	if b.off < len(b.buf) && err == io.EOF {
    60  		// We've reached EOF from filling the buffer but the caller hasn't reached the end yet.
    61  		// Clear EOF for now; we'll find it again after the caller reaches the end of the buffer.
    62  		err = nil
    63  	}
    64  	b.filePos += int64(n)
    65  	return n, err
    66  }
    67  
    68  // Seek implements ioctx.Seeker.
    69  func (b *readSeeker) Seek(ctx context.Context, request int64, whence int) (int64, error) {
    70  	if err := b.initFilePos(ctx); err != nil {
    71  		return 0, err
    72  	}
    73  	var diff int64
    74  	switch whence {
    75  	case io.SeekStart:
    76  		diff = request - b.filePos
    77  	case io.SeekCurrent:
    78  		diff = request
    79  	case io.SeekEnd:
    80  		diff = b.fileEnd + request - b.filePos
    81  	default:
    82  		panic(whence)
    83  	}
    84  	if -int64(b.off) <= diff && diff <= int64(len(b.buf)-b.off) {
    85  		// Seek within buffer without changing file position.
    86  		b.off += int(diff)
    87  		b.filePos += diff
    88  		return b.filePos, nil
    89  	}
    90  	// Discard the buffer and seek the underlying reader.
    91  	diff -= int64(len(b.buf) - b.off)
    92  	b.buf, b.off = b.buf[:0], 0
    93  	var err error
    94  	b.filePos, err = b.r.Seek(ctx, diff, io.SeekCurrent)
    95  	return b.filePos, err
    96  }
    97  
    98  // initFilePos idempotently initializes filePos and fileEnd.
    99  func (b *readSeeker) initFilePos(ctx context.Context) error {
   100  	if b.filePos >= 0 && b.fileEnd >= 0 {
   101  		return nil
   102  	}
   103  	var err error
   104  	b.filePos, err = b.r.Seek(ctx, 0, io.SeekCurrent)
   105  	if err != nil {
   106  		return err
   107  	}
   108  	b.fileEnd, err = b.r.Seek(ctx, 0, io.SeekEnd)
   109  	if err != nil {
   110  		return err
   111  	}
   112  	_, err = b.r.Seek(ctx, b.filePos, io.SeekStart)
   113  	return err
   114  }