/
pool.go
168 lines (146 loc) · 3.6 KB
/
pool.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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
package zlog
import (
"fmt"
"io"
"reflect"
"sync"
)
// Pooled buffers, modeled on the way the slog.JSONHandler does it.
// BufPool is the global pool of buffers.
var bufPool = sync.Pool{
New: func() any {
n := make([]byte, 0, 1024)
if len(n) != 0 {
panic("WTF")
}
return (*buffer)(&n)
},
}
// Buffer is a byte buffer implemented over a slice.
//
// Implementing it this way makes all the helper functions methods instead of
// just free functions.
type buffer []byte
// NewBuffer returns a buffer from the global pool, allocating if necessary.
func newBuffer() *buffer {
return bufPool.Get().(*buffer)
}
// Custom methods first:
// Release returns modestly sized buffers back to the [bufPool] and leaks large
// buffers.
//
// As a convenience, this may be called on a nil receiver.
func (b *buffer) Release() {
const maxSz = 16 << 10
if b == nil {
return
}
if cap(*b) <= maxSz {
*b = (*b)[:0]
bufPool.Put(b)
}
}
// Tail reports the last-written byte, like an backwards [io.ByteReader].
func (b *buffer) Tail() byte {
return (*b)[len(*b)-1]
}
// Unwrite slices off the last-written byte.
func (b *buffer) Unwrite() {
*b = (*b)[:len(*b)-1]
}
// ReplaceTail replaces the last-written byte.
//
// This method is more efficient than [Unwrite] + [WriteByte].
func (b *buffer) ReplaceTail(c byte) {
(*b)[len(*b)-1] = c
}
// Clone returns a new buffer with the contents of the receiver.
//
// The receiver does not have Release called on it. As a convenience, this may
// be called on a nil receiver.
func (b *buffer) Clone() (out *buffer) {
out = newBuffer()
if b == nil {
return out
}
if cap(*b) > cap(*out) {
// Leak small buffer, this new larger will end up in the pool
// instead.
*out = make([]byte, len(*b), cap(*b))
}
if len(*b) > len(*out) {
*out = (*out)[:len(*b)]
}
copy(*out, *b)
return out
}
// Boring methods:
var (
_ io.Writer = (*buffer)(nil)
_ io.WriterTo = (*buffer)(nil)
_ io.ByteWriter = (*buffer)(nil)
_ io.StringWriter = (*buffer)(nil)
)
// WriteString implements [io.StringWriter].
func (b *buffer) WriteString(s string) (int, error) {
*b = append(*b, s...)
return len(s), nil
}
// WriteByte implements [io.ByteWriter].
func (b *buffer) WriteByte(c byte) error {
*b = append(*b, c)
return nil
}
// Write implements [io.Writer].
func (b *buffer) Write(in []byte) (int, error) {
*b = append(*b, in...)
return len(in), nil
}
// WriteTo implements [io.WriterTo].
func (b *buffer) WriteTo(w io.Writer) (int64, error) {
n, err := w.Write(*b)
return int64(n), err
}
// The global statePools.
var pools = map[reflect.Type]interface{}{
reflect.TypeOf(stateJSON{}): &statePool[*stateJSON]{
New: func() *stateJSON { return new(stateJSON) },
},
reflect.TypeOf(stateJournal{}): &statePool[*stateJournal]{
New: func() *stateJournal { return new(stateJournal) },
},
}
// GetPool returns the type-specific [statePool].
//
// GetPool panics if the type parameter is not a pointer type.
func getPool[T state]() *statePool[T] {
var t T
typ := reflect.TypeOf(t).Elem()
var v interface{}
var ok bool
v, ok = pools[typ]
if !ok {
panic(fmt.Sprintf("programmer error: called with unexpected type: %T", t))
}
return v.(*statePool[T])
}
// StatePool is just a typed wrapper around a [sync.Pool].
type statePool[V state] struct {
sync.Pool
// New is a constructor for the concrete type.
New func() V
}
// Get returns a new stateObj.
func (p *statePool[V]) Get(g []string, b *buffer) (v V) {
if x := p.Pool.Get(); x != nil {
v = x.(V)
} else {
v = p.New()
}
v.Reset(g, b)
return v
}
// Put returns the stateObj to the pool.
func (p *statePool[V]) Put(v V) {
p.Pool.Put(v)
}