bufio.Copy

v1.go

package egonelbre

import (
	"errors"
	"io"
)

const (
	maxChunk = 16 << 10
)

func Copy(dst io.Writer, src io.Reader, buffer int) (written int64, err error) {
	var rerr, werr error

	buf := make([]byte, buffer, buffer)
	n := (buffer + maxChunk - 1) / maxChunk

	hot := make(chan []byte, 1)
	cold := make(chan []byte, n)
	work := make(chan []byte, n)

	rq := make(chan struct{})

	for i := 0; i < n; i += 1 {
		s, e := i*maxChunk, i*maxChunk+maxChunk
		if e > len(buf) {
			e = len(buf)
		}
		cold <- buf[s:e:e]
	}

	go func() {
		defer close(rq)
		defer close(work)

		var (
			nr   int
			next []byte
			ok   bool
		)

		for {
			select {
			case next, ok = <-hot:
			default:
				select {
				case next, ok = <-hot:
				case next, ok = <-cold:
				}
			}

			if !ok {
				break
			}

			nr, rerr = src.Read(next[:cap(next)])
			work <- next[:nr]
			if rerr != nil {
				break
			}
		}
	}()

	var nr int
	for next := range work {
		nr, werr = dst.Write(next)
		written += int64(nr)
		if werr != nil {
			break
		} else if nr != len(next) {
			werr = errors.New("Short write.")
			break
		}
		select {
		case hot <- next:
		case v := <-hot:
			hot <- next
			cold <- v
		}
	}
	close(hot)
	close(cold)

	// Wait until reader also closes
	<-rq

	switch {
	case rerr == nil || rerr == io.EOF:
		return written, werr
	// what if both errors happened?
	default:
		return written, rerr
	}
}

v2.go

package egonelbre

import (
	"errors"
	"io"
)

const (
	maxChunk = 16 << 10
)

func Copy(dst io.Writer, src io.Reader, buffer int) (written int64, err error) {
	var rerr, werr error

	buf := make([]byte, buffer, buffer)
	n := (buffer + maxChunk - 1) / maxChunk

	free := make(chan []byte, n)
	work := make(chan []byte, n)

	rq := make(chan struct{})

	for i := 0; i < n; i += 1 {
		s, e := i*maxChunk, i*maxChunk+maxChunk
		if e > len(buf) {
			e = len(buf)
		}
		free <- buf[s:e:e]
	}

	go func() {
		defer close(rq)
		defer close(work)

		var nr int
		for {
			next, ok := <-free
			if !ok {
				break
			}

			nr, rerr = src.Read(next[:cap(next)])
			work <- next[:nr]
			if rerr != nil {
				break
			}
		}
	}()

	var nr int
	for next := range work {
		nr, werr = dst.Write(next)
		written += int64(nr)
		if werr != nil {
			break
		} else if nr != len(next) {
			werr = errors.New("Short write.")
			break
		}
		free <- next
	}
	close(free)

	// Wait until reader also closes
	<-rq

	switch {
	case rerr == nil || rerr == io.EOF:
		return written, werr
	// what if both errors happened?
	default:
		return written, rerr
	}
}