go-ml.dev/pkg/base@v0.0.0-20200610162856-60c38abac71b/fu/pool.go (about)

     1  package fu
     2  
     3  import (
     4  	"io"
     5  	"runtime"
     6  )
     7  
     8  type AtomicPool_ struct {
     9  	mask *AtomicMask_
    10  	pool []interface{}
    11  }
    12  
    13  func AtomicPool(acquire func(no int) interface{}) *AtomicPool_ {
    14  	a := &AtomicPool_{}
    15  	cpunum := runtime.NumCPU()
    16  	a.mask = ExtendableAtomicMask(
    17  		func(no int) bool {
    18  			// synchronized call
    19  			if no < cpunum {
    20  				if v := acquire(no); v != nil {
    21  					a.pool = append(a.pool, v)
    22  					return true
    23  				}
    24  			}
    25  			return false
    26  		})
    27  	return a
    28  }
    29  
    30  func (a *AtomicPool_) Allocate() (interface{}, int) {
    31  	n := a.mask.Lock()
    32  	return a.pool[n], n
    33  }
    34  
    35  func (a *AtomicPool_) Release(n int) {
    36  	a.mask.Unlock(n)
    37  }
    38  
    39  func (a *AtomicPool_) Close() {
    40  	a.mask.FinCallForAll(func(no int) {
    41  		if ci, ok := a.pool[no].(io.Closer); ok {
    42  			ci.Close()
    43  			a.pool[no] = nil // panic on reuse
    44  		}
    45  	})
    46  }