From 3951e858058d4024e1005b6637a121cb2ed456ab Mon Sep 17 00:00:00 2001 From: FanOne <294350394@qq.com> Date: Tue, 11 Jun 2024 23:47:19 +0800 Subject: [PATCH] feat:avoid too many goroutines by ants pool Signed-off-by: FanOne <294350394@qq.com> --- pkg/util/distance/calc_distance.go | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/pkg/util/distance/calc_distance.go b/pkg/util/distance/calc_distance.go index 3a3e4558d796c..776b42cc3ce70 100644 --- a/pkg/util/distance/calc_distance.go +++ b/pkg/util/distance/calc_distance.go @@ -6,6 +6,7 @@ import ( "sync" "github.com/cockroachdb/errors" + "github.com/panjf2000/ants/v2" ) /** @@ -130,15 +131,27 @@ func CalcFloatDistance(dim int64, left, right []float32, metric string) ([]float distArray := make([]float32, leftNum*rightNum) - // Multi-threads to calculate distance. TODO: avoid too many go routines + // Multi-threads to calculate distance. var waitGroup sync.WaitGroup - CalcWorker := func(index int64) { - CalcFFBatch(dim, left, index, right, metricUpper, &distArray) + calcWorker := func(index interface{}) { + i := index.(int64) + CalcFFBatch(dim, left, i, right, metricUpper, &distArray) waitGroup.Done() } + // avoid too many goroutines by ants pool + poolSize := int(leftNum / 3) + pool, err := ants.NewPoolWithFunc(poolSize, calcWorker) + if err != nil { + return nil, err + } + defer pool.Release() + for i := int64(0); i < leftNum; i++ { waitGroup.Add(1) - go CalcWorker(i) + err = pool.Invoke(i) + if err != nil { + return nil, err + } } waitGroup.Wait()