go_goutils/iox/funcs_ctxio.go

package iox

import (
	`bytes`
	`context`
	`io`
	`math`
)

func (c *CtxIO) Copy(dst io.Writer, src io.Reader) (written int64, err error) {
	if c.l.chunkLen > math.MaxInt64 {
		err = ErrChunkTooBig
	}
	return CopyCtxBufN(c.ctx, dst, src, int64(c.l.chunkLen))
}

func (c *CtxIO) CopyBufN(dst io.Writer, src io.Reader, n int64) (written int64, err error) {
	if n <= 0 {
		err = ErrBufTooSmall
		return
	}
	return CopyCtxBufN(c.ctx, dst, src, n)
}

func (c *CtxIO) GetChunkLen() (size uint) {
	return c.l.GetChunkLen()
}

func (c *CtxIO) Read(p []byte) (n int, err error) {

	var nr int64

	if nr, err = c.ReadWithContext(c.ctx, p); err != nil {
		if nr > math.MaxInt {
			n = math.MaxInt
		} else {
			n = int(nr)
		}
		return
	}

	if nr > math.MaxInt {
		n = math.MaxInt
	} else {
		n = int(nr)
	}

	return
}

func (c *CtxIO) ReadWithContext(ctx context.Context, p []byte) (n int64, err error) {

	var nr int
	var off int
	var buf []byte

	if p == nil || len(p) == 0 {
		return
	}
	if c.buf.Len() == 0 {
		err = io.EOF
		return
	}

	if c.l.chunkLen > uint(len(p)) {
		// Would normally be a single chunk, so one-shot it.
		nr, err = c.buf.Read(p)
		n = int64(nr)
		return
	}

	// Chunk over it.
endRead:
	for {
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return
		default:
			/*
				off(set) is the index of the *next position* to write to.
				Therefore the last offset == len(p),
				therefore:

					* if off == len(p), "done" (return no error, do *not* read from buf)
					* if off + c.l.chunkLen > len(p), buf should be len(p) - off instead
			*/
			if off == len(p) {
				break endRead
			}
			if uint(off)+c.l.chunkLen > uint(len(p)) {
				buf = make([]byte, len(p)-off)
			} else {
				buf = make([]byte, c.l.chunkLen)
			}
			nr, err = c.buf.Read(buf)
			n += int64(nr)
			if nr > 0 {
				off += nr
				copy(p[off:], buf[:nr])
			}
			if err == io.EOF {
				break endRead
			} else if err != nil {
				return
			}
		}
	}

	return
}

func (c *CtxIO) SetChunkLen(size uint) (err error) {
	return c.l.SetChunkLen(size)
}

func (c *CtxIO) SetContext(ctx context.Context) (err error) {

	if ctx == nil {
		err = ErrNilCtx
		return
	}

	c.ctx = ctx

	return
}

func (c *CtxIO) Write(p []byte) (n int, err error) {

	var nw int64

	if c.l.chunkLen > math.MaxInt64 {
		err = ErrChunkTooBig
		return
	}
	if nw, err = c.WriteNWithContext(c.ctx, p, int64(c.l.chunkLen)); err != nil {
		if nw > math.MaxInt {
			n = math.MaxInt
		} else {
			n = int(nw)
		}
		return
	}

	if nw > math.MaxInt {
		n = math.MaxInt
	} else {
		n = int(nw)
	}

	return
}

func (c *CtxIO) WriteNWithContext(ctx context.Context, p []byte, n int64) (written int64, err error) {
	return CopyCtxBufN(ctx, &c.buf, bytes.NewReader(p), n)
}

func (c *CtxIO) WriteRune(r rune) (n int, err error) {

	// We don't even bother listening for the ctx.Done because it's a single rune.
	n, err = c.buf.WriteRune(r)

	return
}

func (c *CtxIO) WriteWithContext(ctx context.Context, p []byte) (n int64, err error) {
	if c.l.chunkLen > math.MaxInt64 {
		err = ErrChunkTooBig
		return
	}
	return CopyCtxBufN(ctx, &c.buf, bytes.NewReader(p), int64(c.l.chunkLen))
}