tstor/pkg/ctxio/reader.go

package ctxio

import (
	"context"
	"io"
)

type FileReader interface {
	Reader
	ReaderAt
	Closer
}

type contextReader struct {
	ctx context.Context
	r   Reader
}

func (r *contextReader) Read(p []byte) (n int, err error) {
	if r.ctx.Err() != nil {
		return 0, r.ctx.Err()
	}

	return r.r.Read(r.ctx, p)
}

func IoReaderAt(ctx context.Context, r ReaderAt) io.ReaderAt {
	return &contextReaderAt{ctx: ctx, r: r}
}

type contextReaderAt struct {
	ctx context.Context
	r   ReaderAt
}

func (c *contextReaderAt) ReadAt(p []byte, off int64) (n int, err error) {
	if c.ctx.Err() != nil {
		return 0, c.ctx.Err()
	}

	return c.r.ReadAt(c.ctx, p, off)
}

func IoReader(ctx context.Context, r Reader) io.Reader {
	return &contextReader{ctx: ctx, r: r}
}

func WrapIoReader(r io.Reader) Reader {
	return &wrapReader{r: r}
}

type wrapReader struct {
	r io.Reader
}

var _ Reader = (*wrapReader)(nil)

// Read implements Reader.
func (c *wrapReader) Read(ctx context.Context, p []byte) (n int, err error) {
	if ctx.Err() != nil {
		return 0, ctx.Err()
	}
	return c.r.Read(p)
}

func WrapIoWriter(w io.Writer) Writer {
	return &wrapWriter{w: w}
}

type wrapWriter struct {
	w io.Writer
}

var _ Writer = (*wrapWriter)(nil)

// Write implements Writer.
func (c *wrapWriter) Write(ctx context.Context, p []byte) (n int, err error) {
	if ctx.Err() != nil {
		return 0, ctx.Err()
	}
	return c.w.Write(p)
}

func WrapIoReadCloser(r io.ReadCloser) ReadCloser {
	return &wrapReadCloser{r: r}
}

type wrapReadCloser struct {
	r io.ReadCloser
}

var _ Reader = (*wrapReadCloser)(nil)

// Read implements Reader.
func (c *wrapReadCloser) Read(ctx context.Context, p []byte) (n int, err error) {
	if ctx.Err() != nil {
		return 0, ctx.Err()
	}
	return c.r.Read(p)
}

// Close implements ReadCloser.
func (c *wrapReadCloser) Close(ctx context.Context) error {
	return c.r.Close()
}