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
|
package refclose
import (
"runtime/pprof"
"sync"
)
var profile = pprof.NewProfile("refs")
type RefPool struct {
mu sync.Mutex
rs map[interface{}]*resource
}
type Closer func()
func (me *RefPool) inc(key interface{}) {
me.mu.Lock()
defer me.mu.Unlock()
r := me.rs[key]
if r == nil {
r = new(resource)
if me.rs == nil {
me.rs = make(map[interface{}]*resource)
}
me.rs[key] = r
}
r.numRefs++
}
func (me *RefPool) dec(key interface{}) {
me.mu.Lock()
defer me.mu.Unlock()
r := me.rs[key]
r.numRefs--
if r.numRefs > 0 {
return
}
if r.numRefs < 0 {
panic(r.numRefs)
}
r.closer()
delete(me.rs, key)
}
type resource struct {
closer Closer
numRefs int
}
func (me *RefPool) NewRef(key interface{}) (ret *Ref) {
me.inc(key)
ret = &Ref{
pool: me,
key: key,
}
profile.Add(ret, 0)
return
}
type Ref struct {
mu sync.Mutex
pool *RefPool
key interface{}
closed bool
}
func (me *Ref) SetCloser(closer Closer) {
me.pool.mu.Lock()
defer me.pool.mu.Unlock()
me.pool.rs[me.key].closer = closer
}
func (me *Ref) panicIfClosed() {
if me.closed {
panic("ref is closed")
}
}
func (me *Ref) Release() {
me.mu.Lock()
defer me.mu.Unlock()
me.panicIfClosed()
profile.Remove(me)
me.pool.dec(me.key)
}
func (me *Ref) Key() interface{} {
me.mu.Lock()
defer me.mu.Unlock()
me.panicIfClosed()
return me.key
}
|