enhance: Add dynamic cgo pool for proxy CGO call (#34768)

Related to #34705

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2024-07-19 15:11:41 +08:00 committed by GitHub
parent 2ac7164c39
commit e4e18cb8c3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

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"
)
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))
check := bool(C.CheckVecIndexWithDataType(cIndexName, cType))
return check
result = bool(C.CheckVecIndexWithDataType(cIndexName, cType))
return nil, nil
}).Await()
return result
}