milvus/internal/util/cgo/manager_active.go

125 lines
3.3 KiB
Go
Raw Normal View History

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]
}