Skip to content

Commit

Permalink
enhance: Add dynamic cgo pool for proxy CGO call (milvus-io#34768)
Browse files Browse the repository at this point in the history
Related to milvus-io#34705

Signed-off-by: Congqi Xia <[email protected]>
  • Loading branch information
congqixia committed Jul 19, 2024
1 parent ff4bd2c commit 2c08722
Showing 1 changed file with 43 additions and 5 deletions.
48 changes: 43 additions & 5 deletions internal/proxy/cgo_util.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
cIndexName := C.CString(name)
cType := uint32(dType)
defer C.free(unsafe.Pointer(cIndexName))
check := bool(C.CheckVecIndexWithDataType(cIndexName, cType))
return check
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
}

0 comments on commit 2c08722

Please sign in to comment.