cubefs

1
// Copyright 2022 The CubeFS Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12
// implied. See the License for the specific language governing
13
// permissions and limitations under the License.
14

15
package resourcepool
16

17
import (
18
	"errors"
19
	"sort"
20
)
21

22
// ErrNoSuitableSizeClass no suitable pool of size
23
var ErrNoSuitableSizeClass = errors.New("no suitable size class")
24

25
// zero bytes, high performance at the 16KB, see more in the benchmark:
26
// BenchmarkZero/4MB-16KB-4                   13338             88378 ns/op
27
// BenchmarkZero/8MB-16KB-4                    6670            183987 ns/op
28
// BenchmarkZero/16MB-16KB-4                   1926            590422 ns/op
29
const zeroLen = 1 << 14
30

31
var zero = make([]byte, zeroLen)
32

33
// MemPool reused buffer pool
34
type MemPool struct {
35
	pool     []Pool
36
	poolSize []int
37
}
38

39
// Status memory pools status
40
type Status []PoolStatus
41

42
// PoolStatus status of the pool
43
type PoolStatus struct {
44
	Size     int `json:"size"`
45
	Capacity int `json:"capacity"`
46
	Running  int `json:"running"`
47
	Idle     int `json:"idle"`
48
}
49

50
// NewMemPool returns a MemPool within chan pool
51
func NewMemPool(sizeClasses map[int]int) *MemPool {
52
	return NewMemPoolWith(sizeClasses, func(size, capacity int) Pool {
53
		return NewChanPool(func() []byte {
54
			return make([]byte, size)
55
		}, capacity)
56
	})
57
}
58

59
// NewMemPoolWith new MemPool with size-class and self-defined pool
60
func NewMemPoolWith(sizeClasses map[int]int, newPool func(size, capacity int) Pool) *MemPool {
61
	pool := make([]Pool, 0, len(sizeClasses))
62
	poolSize := make([]int, 0, len(sizeClasses))
63
	for sizeClass := range sizeClasses {
64
		if sizeClass > 0 {
65
			poolSize = append(poolSize, sizeClass)
66
		}
67
	}
68

69
	sort.Ints(poolSize)
70
	for _, sizeClass := range poolSize {
71
		pool = append(pool, newPool(sizeClass, sizeClasses[sizeClass]))
72
	}
73

74
	return &MemPool{
75
		pool:     pool,
76
		poolSize: poolSize,
77
	}
78
}
79

80
// Get return a suitable buffer
81
func (p *MemPool) Get(size int) ([]byte, error) {
82
	for idx, ps := range p.poolSize {
83
		if size <= ps {
84
			buf, err := p.pool[idx].Get()
85
			if err != nil {
86
				return nil, err
87
			}
88
			buff := buf.([]byte)
89
			return buff[:size], nil
90
		}
91
	}
92

93
	return nil, ErrNoSuitableSizeClass
94
}
95

96
// Alloc return a buffer, make a new if oversize
97
func (p *MemPool) Alloc(size int) ([]byte, error) {
98
	buf, err := p.Get(size)
99
	if err == ErrNoSuitableSizeClass {
100
		return make([]byte, size), nil
101
	}
102

103
	return buf, err
104
}
105

106
// Put adds x to the pool, appropriately resize
107
func (p *MemPool) Put(b []byte) error {
108
	sizeClass := cap(b)
109
	b = b[0:sizeClass]
110
	for ii := len(p.poolSize) - 1; ii >= 0; ii-- {
111
		if sizeClass >= p.poolSize[ii] {
112
			b = b[0:p.poolSize[ii]]
113
			p.pool[ii].Put(b)
114
			return nil
115
		}
116
	}
117

118
	return ErrNoSuitableSizeClass
119
}
120

121
// Zero clean up the buffer b to zero bytes
122
func (p *MemPool) Zero(b []byte) {
123
	Zero(b)
124
}
125

126
// Status returns status of memory pool
127
func (p *MemPool) Status() Status {
128
	st := make(Status, len(p.poolSize))
129
	for idx, size := range p.poolSize {
130
		pool := p.pool[idx]
131
		st[idx] = PoolStatus{
132
			Size:     size,
133
			Capacity: pool.Cap(),
134
			Running:  pool.Len(),
135
			Idle:     pool.Idle(),
136
		}
137
	}
138
	return st
139
}
140

141
// Zero clean up the buffer b to zero bytes
142
func Zero(b []byte) {
143
	for len(b) > 0 {
144
		n := copy(b, zero)
145
		b = b[n:]
146
	}
147
}
148