-
Notifications
You must be signed in to change notification settings - Fork 96
/
bufpool.go
103 lines (90 loc) · 2.74 KB
/
bufpool.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package moxio
import (
"bufio"
"errors"
"fmt"
"io"
"golang.org/x/exp/slog"
"github.com/mjl-/mox/mlog"
)
// todo: instead of a bufpool, should maybe just make an alternative to bufio.Reader with a big enough buffer that we can fully use to read a line.
var ErrLineTooLong = errors.New("line from remote too long") // Returned by Bufpool.Readline.
// Bufpool caches byte slices for reuse during parsing of line-terminated commands.
type Bufpool struct {
c chan []byte
size int
}
// NewBufpool makes a new pool, initially empty, but holding at most "max" buffers of "size" bytes each.
func NewBufpool(max, size int) *Bufpool {
return &Bufpool{
c: make(chan []byte, max),
size: size,
}
}
// get returns a buffer from the pool if available, otherwise allocates a new buffer.
// The buffer should be returned with a call to put.
func (b *Bufpool) get() []byte {
var buf []byte
// Attempt to get buffer from pool. Otherwise create new buffer.
select {
case buf = <-b.c:
default:
}
if buf == nil {
buf = make([]byte, b.size)
}
return buf
}
// put puts a "buf" back in the pool. Put clears the first "n" bytes, which should
// be all the bytes that have been read in the buffer. If the pool is full, the
// buffer is discarded, and will be cleaned up by the garbage collector.
// The caller should no longer reference "buf" after a call to put.
func (b *Bufpool) put(log mlog.Log, buf []byte, n int) {
if len(buf) != b.size {
log.Error("buffer with bad size returned, ignoring", slog.Int("badsize", len(buf)), slog.Int("expsize", b.size))
return
}
for i := 0; i < n; i++ {
buf[i] = 0
}
select {
case b.c <- buf:
default:
}
}
// Readline reads a \n- or \r\n-terminated line. Line is returned without \n or \r\n.
// If the line was too long, ErrLineTooLong is returned.
// If an EOF is encountered before a \n, io.ErrUnexpectedEOF is returned.
func (b *Bufpool) Readline(log mlog.Log, r *bufio.Reader) (line string, rerr error) {
var nread int
buf := b.get()
defer func() {
b.put(log, buf, nread)
}()
// Read until newline. If we reach the end of the buffer first, we write back an
// error and abort the connection because our protocols cannot be recovered. We
// don't want to consume data until we finally see a newline, which may be never.
for {
if nread >= len(buf) {
return "", fmt.Errorf("%w: no newline after all %d bytes", ErrLineTooLong, nread)
}
c, err := r.ReadByte()
if err == io.EOF {
return "", io.ErrUnexpectedEOF
} else if err != nil {
return "", fmt.Errorf("reading line from remote: %w", err)
}
if c == '\n' {
var s string
if nread > 0 && buf[nread-1] == '\r' {
s = string(buf[:nread-1])
} else {
s = string(buf[:nread])
}
nread++
return s, nil
}
buf[nread] = c
nread++
}
}