ring: add multiple rewrite WriterTo to ring buffer API

This commit is contained in:
Dan Kortschak 2018-06-03 22:31:44 +09:30
parent a37923439c
commit 4e1ae308ff
2 changed files with 107 additions and 20 deletions

View File

@ -30,7 +30,6 @@ LICENSE
package ring package ring
import ( import (
"bytes"
"errors" "errors"
"io" "io"
"time" "time"
@ -48,8 +47,8 @@ var (
// The buffer has a writable head and a readable tail with a queue from the head // The buffer has a writable head and a readable tail with a queue from the head
// to the tail. Concurrent read a write operations are safe. // to the tail. Concurrent read a write operations are safe.
type Buffer struct { type Buffer struct {
head, tail *bytes.Buffer head, tail *Chunk
full, empty chan *bytes.Buffer full, empty chan *Chunk
timeout time.Duration timeout time.Duration
} }
@ -61,12 +60,12 @@ func NewBuffer(len, size int, timeout time.Duration) *Buffer {
return nil return nil
} }
b := Buffer{ b := Buffer{
full: make(chan *bytes.Buffer, len), full: make(chan *Chunk, len),
empty: make(chan *bytes.Buffer, len), empty: make(chan *Chunk, len),
timeout: timeout, timeout: timeout,
} }
for i := 0; i < len; i++ { for i := 0; i < len; i++ {
b.empty <- bytes.NewBuffer(make([]byte, 0, size)) b.empty <- newChunk(make([]byte, 0, size))
} }
return &b return &b
} }
@ -92,7 +91,7 @@ func (b *Buffer) Write(p []byte) (int, error) {
case <-timer.C: case <-timer.C:
select { select {
case b.head = <-b.full: case b.head = <-b.full:
b.head.Reset() b.head.reset()
dropped = true dropped = true
default: default:
// This should never happen. // This should never happen.
@ -102,16 +101,16 @@ func (b *Buffer) Write(p []byte) (int, error) {
timer.Stop() timer.Stop()
} }
} }
if len(p) > b.head.Cap() { if len(p) > b.head.cap() {
return 0, ErrTooLong return 0, ErrTooLong
} }
if len(p) > b.head.Cap()-b.head.Len() { if len(p) > b.head.cap()-b.head.len() {
b.full <- b.head b.full <- b.head
b.head = nil b.head = nil
return b.Write(p) return b.Write(p)
} }
n, err := b.head.Write(p) n, err := b.head.write(p)
if b.head.Cap()-b.head.Len() == 0 { if b.head.cap()-b.head.len() == 0 {
b.full <- b.head b.full <- b.head
b.head = nil b.head = nil
} }
@ -150,21 +149,22 @@ func (b *Buffer) Close() error {
// //
// Next is safe to use concurrently with write operations, but may not be used concurrently with // Next is safe to use concurrently with write operations, but may not be used concurrently with
// another Read call or Next call. A goroutine calling Next must not call Flush or Close. // another Read call or Next call. A goroutine calling Next must not call Flush or Close.
func (b *Buffer) Next(timeout time.Duration) error { func (b *Buffer) Next(timeout time.Duration) (*Chunk, error) {
if b.tail == nil { if b.tail == nil {
timer := time.NewTimer(timeout) timer := time.NewTimer(timeout)
var ok bool var ok bool
select { select {
case <-timer.C: case <-timer.C:
return ErrTimeout return nil, ErrTimeout
case b.tail, ok = <-b.full: case b.tail, ok = <-b.full:
timer.Stop() timer.Stop()
if !ok { if !ok {
return io.EOF return nil, io.EOF
} }
} }
} }
return nil b.tail.owner = b
return b.tail, nil
} }
// Read reads bytes from the current tail of the ring buffer into p and returns the number of // Read reads bytes from the current tail of the ring buffer into p and returns the number of
@ -176,11 +176,98 @@ func (b *Buffer) Read(p []byte) (int, error) {
if b.tail == nil { if b.tail == nil {
return 0, io.EOF return 0, io.EOF
} }
n, err := b.tail.Read(p) n, err := b.tail.read(p)
if b.tail.Len() == 0 { if b.tail.len() == 0 {
b.tail.Reset() b.tail.reset()
b.empty <- b.tail b.empty <- b.tail
b.tail = nil b.tail = nil
} }
return n, err return n, err
} }
// Chunk is a simplified version of byte buffer without the capacity to grow beyond the
// buffer's original cap, and a modified WriteTo method that allows multiple calls without
// consuming the buffered data.
type Chunk struct {
buf []byte
off int
owner *Buffer
}
func newChunk(buf []byte) *Chunk {
return &Chunk{buf: buf[:0]}
}
func (b *Chunk) len() int {
return len(b.buf) - b.off
}
func (b *Chunk) cap() int {
return cap(b.buf)
}
func (b *Chunk) reset() {
b.buf = b.buf[:0]
b.off = 0
}
func (b *Chunk) write(p []byte) (n int, err error) {
if len(p) > cap(b.buf)-len(b.buf) {
err = ErrTooLong
}
l := len(b.buf)
m := l + len(p)
if m > cap(b.buf) {
m = cap(b.buf)
}
b.buf = b.buf[:m]
n = copy(b.buf[l:], p)
return n, err
}
func (b *Chunk) read(p []byte) (n int, err error) {
if b.len() <= 0 {
if len(p) == 0 {
return 0, nil
}
return 0, io.EOF
}
n = copy(p, b.buf[b.off:])
b.off += n
return n, nil
}
// WriteTo writes data to w until there's no more data to write or when an error occurs.
// The return value n is the number of bytes written. Any error encountered during the
// write is also returned. Repeated called to WriteTo will write the same data until
// the Chunk's Close method is called.
//
// WriteTo will panic if the Chunk has not been obtained through a call to Buffer.Next or
// has been closed. WriteTo must be used in the same goroutine as the call to Next.
func (b *Chunk) WriteTo(w io.Writer) (n int, err error) {
if b.owner == nil || b.owner.tail != b {
panic("ring: invalid use of ring buffer chunk")
}
n, err = w.Write(b.buf)
if n > len(b.buf) {
panic("ring: invalid byte count")
}
if n != len(b.buf) {
return n, io.ErrShortWrite
}
return n, nil
}
// Close closes the Chunk, reseting its data and releasing it back to the Buffer. A Chunk
// may not be used after it has been closed. Close must be used in the same goroutine as
// the call to Next.
func (b *Chunk) Close() error {
if b.owner == nil || b.owner.tail != b {
panic("ring: invalid use of ring buffer chunk")
}
b.reset()
b.owner = nil
b.owner.tail = nil
b.owner.empty <- b
return nil
}

View File

@ -139,7 +139,7 @@ func TestRoundTrip(t *testing.T) {
var timeouts int var timeouts int
elements: elements:
for { for {
err := b.Next(test.nextTimeout) _, err := b.Next(test.nextTimeout)
switch err { switch err {
case nil: case nil:
timeouts = 0 timeouts = 0
@ -236,7 +236,7 @@ func BenchmarkRoundTrip(b *testing.B) {
var timeouts int var timeouts int
elements: elements:
for { for {
err := rb.Next(timeout) _, err := rb.Next(timeout)
switch err { switch err {
case nil: case nil:
timeouts = 0 timeouts = 0