diff --git a/net.go b/net.go index ffa5653..e4833a2 100644 --- a/net.go +++ b/net.go @@ -8,6 +8,7 @@ package net import ( "errors" + "io" "time" ) @@ -318,3 +319,77 @@ var errClosed = errNetClosing{} // in another error, and should normally be tested using // errors.Is(err, net.ErrClosed). var ErrClosed error = errClosed + +// buffersWriter is the interface implemented by Conns that support a +// "writev"-like batch write optimization. +// writeBuffers should fully consume and write all chunks from the +// provided Buffers, else it should report a non-nil error. +type buffersWriter interface { + writeBuffers(*Buffers) (int64, error) +} + +// Buffers contains zero or more runs of bytes to write. +// +// On certain machines, for certain types of connections, this is +// optimized into an OS-specific batch write operation (such as +// "writev"). +type Buffers [][]byte + +var ( + _ io.WriterTo = (*Buffers)(nil) + _ io.Reader = (*Buffers)(nil) +) + +// WriteTo writes contents of the buffers to w. +// +// WriteTo implements [io.WriterTo] for [Buffers]. +// +// WriteTo modifies the slice v as well as v[i] for 0 <= i < len(v), +// but does not modify v[i][j] for any i, j. +func (v *Buffers) WriteTo(w io.Writer) (n int64, err error) { + if wv, ok := w.(buffersWriter); ok { + return wv.writeBuffers(v) + } + for _, b := range *v { + nb, err := w.Write(b) + n += int64(nb) + if err != nil { + v.consume(n) + return n, err + } + } + v.consume(n) + return n, nil +} + +// Read from the buffers. +// +// Read implements [io.Reader] for [Buffers]. +// +// Read modifies the slice v as well as v[i] for 0 <= i < len(v), +// but does not modify v[i][j] for any i, j. +func (v *Buffers) Read(p []byte) (n int, err error) { + for len(p) > 0 && len(*v) > 0 { + n0 := copy(p, (*v)[0]) + v.consume(int64(n0)) + p = p[n0:] + n += n0 + } + if len(*v) == 0 { + err = io.EOF + } + return +} + +func (v *Buffers) consume(n int64) { + for len(*v) > 0 { + ln0 := int64(len((*v)[0])) + if ln0 > n { + (*v)[0] = (*v)[0][n:] + return + } + n -= ln0 + (*v)[0] = nil + *v = (*v)[1:] + } +}