enhance: [2.4] Add dynamic cgo pool for proxy CGO call (#34768) (#34842)

Cherry-pick from master
pr: #34768
Related to #34705

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
pull/34869/head
congqixia 2024-07-21 18:05:41 +08:00 committed by GitHub
parent 8834649650
commit d1e0d4ea4a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 44 additions and 6 deletions

View File

@ -24,15 +24,53 @@ package proxy
import "C"
import (
"runtime"
"sync"
"unsafe"
"go.uber.org/atomic"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/conc"
"github.com/milvus-io/milvus/pkg/util/hardware"
)
func CheckVecIndexWithDataTypeExist(name string, dType schemapb.DataType) bool {
cIndexName := C.CString(name)
cType := uint32(dType)
defer C.free(unsafe.Pointer(cIndexName))
check := bool(C.CheckVecIndexWithDataType(cIndexName, cType))
return check
var (
dp atomic.Pointer[conc.Pool[any]]
dynOnce sync.Once
)
func initDynamicPool() {
dynOnce.Do(func() {
pool := conc.NewPool[any](
hardware.GetCPUNum(),
conc.WithPreAlloc(false),
conc.WithDisablePurge(false),
conc.WithPreHandler(runtime.LockOSThread), // lock os thread for cgo thread disposal
)
dp.Store(pool)
log.Info("init dynamicPool done", zap.Int("size", hardware.GetCPUNum()))
})
}
// GetDynamicPool returns the singleton pool for dynamic cgo operations.
func GetDynamicPool() *conc.Pool[any] {
initDynamicPool()
return dp.Load()
}
func CheckVecIndexWithDataTypeExist(name string, dType schemapb.DataType) bool {
var result bool
GetDynamicPool().Submit(func() (any, error) {
cIndexName := C.CString(name)
cType := uint32(dType)
defer C.free(unsafe.Pointer(cIndexName))
result = bool(C.CheckVecIndexWithDataType(cIndexName, cType))
return nil, nil
}).Await()
return result
}