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
// sync.Pool cache will be released by runtime.GC()
18
// see sync/pool.go: runtime_registerPoolCleanup(poolCleanup)
19

20
import (
21
	"errors"
22
	"sync"
23
	"sync/atomic"
24
)
25

26
// ErrPoolLimit pool elements exceed its capacity
27
var ErrPoolLimit = errors.New("resource pool limit")
28

29
// Pool resource pool support for sync.pool and capacity limit
30
// release resource if no used anymore
31
// no limit if capacity is negative
32
type Pool interface {
33
	// Get return nil and error if exceed pool's capacity
34
	Get() (interface{}, error)
35
	Put(x interface{})
36
	Cap() int
37
	Len() int
38
	// Idle return cached idle objects in pool.
39
	Idle() int
40
}
41

42
// sync pool Idle return -1 if no limit
43
type pool struct {
44
	sp       sync.Pool
45
	capacity int32
46
	current  int32
47
}
48

49
// NewPool return Pool with capacity, no limit if capacity is negative
50
func NewPool(newFunc func() interface{}, capacity int) Pool {
51
	return &pool{
52
		sp:       sync.Pool{New: newFunc},
53
		capacity: int32(capacity),
54
		current:  int32(0),
55
	}
56
}
57

58
func (p *pool) Get() (interface{}, error) {
59
	current := atomic.AddInt32(&p.current, 1)
60
	if p.capacity >= 0 && current > p.capacity {
61
		atomic.AddInt32(&p.current, -1)
62
		return nil, ErrPoolLimit
63
	}
64
	return p.sp.Get(), nil
65
}
66

67
func (p *pool) Put(x interface{}) {
68
	p.sp.Put(x)
69
	atomic.AddInt32(&p.current, -1)
70
}
71

72
func (p *pool) Cap() int {
73
	return int(p.capacity)
74
}
75

76
func (p *pool) Len() int {
77
	return int(atomic.LoadInt32(&p.current))
78
}
79

80
func (p *pool) Idle() int {
81
	if p.capacity < 0 {
82
		return -1
83
	}
84
	return p.Cap() - p.Len()
85
}
86