mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-11-30 02:48:45 +08:00
125 lines
3.3 KiB
Go
125 lines
3.3 KiB
Go
|
package cgo
|
||
|
|
||
|
import (
|
||
|
"math"
|
||
|
"reflect"
|
||
|
"sync"
|
||
|
|
||
|
"go.uber.org/atomic"
|
||
|
|
||
|
"github.com/milvus-io/milvus/pkg/metrics"
|
||
|
"github.com/milvus-io/milvus/pkg/util/hardware"
|
||
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||
|
)
|
||
|
|
||
|
const (
|
||
|
registerIndex = 0
|
||
|
maxSelectCase = 65535
|
||
|
defaultRegisterBuf = 1
|
||
|
)
|
||
|
|
||
|
var (
|
||
|
futureManager *activeFutureManager
|
||
|
initOnce sync.Once
|
||
|
)
|
||
|
|
||
|
// InitCGO initializes the cgo caller and future manager.
|
||
|
// Please call this function before using any cgo utilities.
|
||
|
func InitCGO() {
|
||
|
initOnce.Do(func() {
|
||
|
nodeID := paramtable.GetStringNodeID()
|
||
|
chSize := int64(math.Ceil(float64(hardware.GetCPUNum()) * paramtable.Get().QueryNodeCfg.CGOPoolSizeRatio.GetAsFloat()))
|
||
|
if chSize <= 0 {
|
||
|
chSize = 1
|
||
|
}
|
||
|
caller = &cgoCaller{
|
||
|
// TODO: temporary solution, need to find a better way to set the pool size.
|
||
|
ch: make(chan struct{}, chSize),
|
||
|
nodeID: nodeID,
|
||
|
}
|
||
|
futureManager = newActiveFutureManager(nodeID)
|
||
|
futureManager.Run()
|
||
|
})
|
||
|
}
|
||
|
|
||
|
type futureManagerStat struct {
|
||
|
ActiveCount int64
|
||
|
}
|
||
|
|
||
|
func newActiveFutureManager(nodeID string) *activeFutureManager {
|
||
|
manager := &activeFutureManager{
|
||
|
activeCount: atomic.NewInt64(0),
|
||
|
activeFutures: make([]basicFuture, 0),
|
||
|
cases: make([]reflect.SelectCase, 1),
|
||
|
register: make(chan basicFuture, defaultRegisterBuf),
|
||
|
nodeID: nodeID,
|
||
|
}
|
||
|
manager.cases[0] = reflect.SelectCase{
|
||
|
Dir: reflect.SelectRecv,
|
||
|
Chan: reflect.ValueOf(manager.register),
|
||
|
}
|
||
|
return manager
|
||
|
}
|
||
|
|
||
|
// activeFutureManager manages the active futures.
|
||
|
// it will transfer the cancel signal into cgo.
|
||
|
type activeFutureManager struct {
|
||
|
activeCount *atomic.Int64
|
||
|
activeFutures []basicFuture
|
||
|
cases []reflect.SelectCase
|
||
|
register chan basicFuture
|
||
|
nodeID string
|
||
|
}
|
||
|
|
||
|
// Run starts the active future manager.
|
||
|
func (m *activeFutureManager) Run() {
|
||
|
go func() {
|
||
|
for {
|
||
|
m.doSelect()
|
||
|
}
|
||
|
}()
|
||
|
}
|
||
|
|
||
|
// Register registers a future when it's created into the manager.
|
||
|
func (m *activeFutureManager) Register(c basicFuture) {
|
||
|
m.register <- c
|
||
|
}
|
||
|
|
||
|
// Stat returns the stat of the manager, only for testing now.
|
||
|
func (m *activeFutureManager) Stat() futureManagerStat {
|
||
|
return futureManagerStat{
|
||
|
ActiveCount: m.activeCount.Load(),
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// doSelect selects the active futures and cancel the finished ones.
|
||
|
func (m *activeFutureManager) doSelect() {
|
||
|
index, newCancelableObject, _ := reflect.Select(m.getSelectableCases())
|
||
|
if index == registerIndex {
|
||
|
newCancelable := newCancelableObject.Interface().(basicFuture)
|
||
|
m.cases = append(m.cases, reflect.SelectCase{
|
||
|
Dir: reflect.SelectRecv,
|
||
|
Chan: reflect.ValueOf(newCancelable.Context().Done()),
|
||
|
})
|
||
|
m.activeFutures = append(m.activeFutures, newCancelable)
|
||
|
} else {
|
||
|
m.cases = append(m.cases[:index], m.cases[index+1:]...)
|
||
|
offset := index - 1
|
||
|
// cancel the future and move it into gc manager.
|
||
|
m.activeFutures[offset].cancel(m.activeFutures[offset].Context().Err())
|
||
|
m.activeFutures = append(m.activeFutures[:offset], m.activeFutures[offset+1:]...)
|
||
|
}
|
||
|
activeTotal := len(m.activeFutures)
|
||
|
m.activeCount.Store(int64(activeTotal))
|
||
|
metrics.ActiveFutureTotal.WithLabelValues(
|
||
|
m.nodeID,
|
||
|
).Set(float64(activeTotal))
|
||
|
}
|
||
|
|
||
|
func (m *activeFutureManager) getSelectableCases() []reflect.SelectCase {
|
||
|
if len(m.cases) <= maxSelectCase {
|
||
|
return m.cases
|
||
|
}
|
||
|
return m.cases[0:maxSelectCase]
|
||
|
}
|