-
Notifications
You must be signed in to change notification settings - Fork 493
/
prealloc.go
50 lines (42 loc) · 1.17 KB
/
prealloc.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
package tempopb
import (
"github.com/prometheus/prometheus/pkg/pool"
)
var (
// buckets: [0.5KiB, 1KiB, 2KiB, 4KiB, 8KiB, 16KiB]
bytePool = pool.New(500, 16_000, 2, func(size int) interface{} { return make([]byte, 0, size) })
)
// PreallocBytes is a (repeated bytes slices) which preallocs slices on Unmarshal.
type PreallocBytes struct {
Slice []byte
}
// Unmarshal implements proto.Message.
func (r *PreallocBytes) Unmarshal(dAtA []byte) error {
r.Slice = bytePool.Get(len(dAtA)).([]byte)
r.Slice = r.Slice[:len(dAtA)]
copy(r.Slice, dAtA)
return nil
}
// MarshalTo implements proto.Marshaller.
// returned int is not used
func (r *PreallocBytes) MarshalTo(dAtA []byte) (int, error) {
copy(dAtA[:], r.Slice[:])
return len(r.Slice), nil
}
// Size implements proto.Sizer.
func (r *PreallocBytes) Size() (n int) {
if r == nil {
return 0
}
return len(r.Slice)
}
// ReuseTraceBytes puts the byte slice back into bytePool for reuse.
func ReuseTraceBytes(trace *TraceBytes) {
for _, t := range trace.Traces {
bytePool.Put(t[:0])
}
}
// SliceFromBytePool gets a slice from the byte pool
func SliceFromBytePool(size int) []byte {
return bytePool.Get(size).([]byte)[:size]
}