test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-12-01 19:39:21 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

enhance: Use pool to limit kafka cgo thread number (#37744)

Browse Source 
Related to #37665

Thread number went rocket high when there is lots of kafka consumers on
datanode. Since the internal implementation is CGO, using which directly
will make cgo thread leaked.

This PR add a worker pool for kafka API utilzing CGO calls to limit
thread number.

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2024-11-18 10:42:31 +08:00 committed by GitHub
parent 3f1614e9d9
commit 3743fc2b00
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 113 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

84
pkg/mq/msgstream/mqwrapper/kafka/kafka_consumer.go
View File

@ -99,7 +99,10 @@ func newKafkaConsumer(config *kafka.ConfigMap, bufSize int64, topic string, grou
func (kc *Consumer) createKafkaConsumer() error {
var err error
kc.c, err = kafka.NewConsumer(kc.config)
getPool().Submit(func() (any, error) {
kc.c, err = kafka.NewConsumer(kc.config)
return nil, err
}).Await()
if err != nil {
log.Error("create kafka consumer failed", zap.String("topic", kc.topic), zap.Error(err))
return err
@ -133,7 +136,12 @@ func (kc *Consumer) Chan() <-chan common.Message {
return
default:
readTimeout := paramtable.Get().KafkaCfg.ReadTimeout.GetAsDuration(time.Second)
e, err := kc.c.ReadMessage(readTimeout)
var e *kafka.Message
var err error
getPool().Submit(func() (any, error) {
e, err = kc.c.ReadMessage(readTimeout)
return nil, err
}).Await()
if err != nil {
// if we failed to read message in 30 Seconds, print out a warn message since there should always be a tt
log.Warn("consume msg failed", zap.String("topic", kc.topic), zap.String("groupID", kc.groupID), zap.Error(err))
@ -166,38 +174,40 @@ func (kc *Consumer) Seek(id common.MessageID, inclusive bool) error {
}
func (kc *Consumer) internalSeek(offset kafka.Offset, inclusive bool) error {
log.Info("kafka consumer seek start", zap.String("topic name", kc.topic),
zap.Any("Msg offset", offset), zap.Bool("inclusive", inclusive))
_, err := getPool().Submit(func() (any, error) {
log.Info("kafka consumer seek start", zap.String("topic name", kc.topic),
zap.Any("Msg offset", offset), zap.Bool("inclusive", inclusive))
start := time.Now()
err := kc.c.Assign([]kafka.TopicPartition{{Topic: &kc.topic, Partition: mqwrapper.DefaultPartitionIdx, Offset: offset}})
if err != nil {
log.Warn("kafka consumer assign failed ", zap.String("topic name", kc.topic), zap.Any("Msg offset", offset), zap.Error(err))
return nil, err
}
start := time.Now()
err := kc.c.Assign([]kafka.TopicPartition{{Topic: &kc.topic, Partition: mqwrapper.DefaultPartitionIdx, Offset: offset}})
if err != nil {
log.Warn("kafka consumer assign failed ", zap.String("topic name", kc.topic), zap.Any("Msg offset", offset), zap.Error(err))
return err
}
cost := time.Since(start).Milliseconds()
if cost > 200 {
log.Warn("kafka consumer assign take too long!", zap.String("topic name", kc.topic),
zap.Any("Msg offset", offset), zap.Bool("inclusive", inclusive), zap.Int64("time cost(ms)", cost))
}
cost := time.Since(start).Milliseconds()
if cost > 200 {
log.Warn("kafka consumer assign take too long!", zap.String("topic name", kc.topic),
// If seek timeout is not 0 the call twice will return error isStarted RD_KAFKA_RESP_ERR__STATE.
// if the timeout is 0 it will initiate the seek but return immediately without any error reporting
kc.skipMsg = !inclusive
if err := kc.c.Seek(kafka.TopicPartition{
Topic: &kc.topic,
Partition: mqwrapper.DefaultPartitionIdx,
Offset: offset,
}, timeout); err != nil {
return nil, err
}
cost = time.Since(start).Milliseconds()
log.Info("kafka consumer seek finished", zap.String("topic name", kc.topic),
zap.Any("Msg offset", offset), zap.Bool("inclusive", inclusive), zap.Int64("time cost(ms)", cost))
}
// If seek timeout is not 0 the call twice will return error isStarted RD_KAFKA_RESP_ERR__STATE.
// if the timeout is 0 it will initiate the seek but return immediately without any error reporting
kc.skipMsg = !inclusive
if err := kc.c.Seek(kafka.TopicPartition{
Topic: &kc.topic,
Partition: mqwrapper.DefaultPartitionIdx,
Offset: offset,
}, timeout); err != nil {
return err
}
cost = time.Since(start).Milliseconds()
log.Info("kafka consumer seek finished", zap.String("topic name", kc.topic),
zap.Any("Msg offset", offset), zap.Bool("inclusive", inclusive), zap.Int64("time cost(ms)", cost))
kc.hasAssign = true
return nil
kc.hasAssign = true
return nil, nil
}).Await()
return err
}
func (kc *Consumer) Ack(message common.Message) {
@ -207,7 +217,13 @@ func (kc *Consumer) Ack(message common.Message) {
}
func (kc *Consumer) GetLatestMsgID() (common.MessageID, error) {
low, high, err := kc.c.QueryWatermarkOffsets(kc.topic, mqwrapper.DefaultPartitionIdx, timeout)
var low, high int64
var err error
getPool().Submit(func() (any, error) {
low, high, err = kc.c.QueryWatermarkOffsets(kc.topic, mqwrapper.DefaultPartitionIdx, timeout)
return nil, err
}).Await()
if err != nil {
return nil, err
}
@ -244,7 +260,11 @@ func (kc *Consumer) CheckTopicValid(topic string) error {
func (kc *Consumer) closeInternal() {
log.Info("close consumer ", zap.String("topic", kc.topic), zap.String("groupID", kc.groupID))
start := time.Now()
err := kc.c.Close()
var err error
getPool().Submit(func() (any, error) {
err = kc.c.Close()
return nil, err
}).Await()
if err != nil {
log.Warn("failed to close ", zap.String("topic", kc.topic), zap.Error(err))
}

14
pkg/mq/msgstream/mqwrapper/kafka/kafka_producer.go
View File

@ -44,11 +44,15 @@ func (kp *kafkaProducer) Send(ctx context.Context, message *mqcommon.ProducerMes
header := kafka.Header{Key: key, Value: []byte(value)}
headers = append(headers, header)
}
err := kp.p.Produce(&kafka.Message{
TopicPartition: kafka.TopicPartition{Topic: &kp.topic, Partition: mqwrapper.DefaultPartitionIdx},
Value: message.Payload,
Headers: headers,
}, kp.deliveryChan)
var err error
getPool().Submit(func() (any, error) {
err = kp.p.Produce(&kafka.Message{
TopicPartition: kafka.TopicPartition{Topic: &kp.topic, Partition: mqwrapper.DefaultPartitionIdx},
Value: message.Payload,
Headers: headers,
}, kp.deliveryChan)
return nil, err
})
if err != nil {
metrics.MsgStreamOpCounter.WithLabelValues(metrics.SendMsgLabel, metrics.FailLabel).Inc()
return nil, err

52
pkg/mq/msgstream/mqwrapper/kafka/pool.go Normal file
View File

@ -0,0 +1,52 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package kafka
import (
"runtime"
"sync"
"go.uber.org/atomic"
"go.uber.org/zap"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/conc"
"github.com/milvus-io/milvus/pkg/util/hardware"
)
var (
kafkaCPool atomic.Pointer[conc.Pool[any]]
initOnce sync.Once
)
func initPool() {
pool := conc.NewPool[any](
hardware.GetCPUNum(),
conc.WithPreAlloc(false),
conc.WithDisablePurge(false),
conc.WithPreHandler(runtime.LockOSThread), // lock os thread for cgo thread disposal
)
kafkaCPool.Store(pool)
log.Info("init dynamicPool done", zap.Int("size", hardware.GetCPUNum()))
}
// GetSQPool returns the singleton pool instance for search/query operations.
func getPool() *conc.Pool[any] {
initPool()
return kafkaCPool.Load()
}