kafka: make the verify lightweight (#5617)

cmd/kafka-consumer/consumer.go

@@ -42,28 +42,50 @@ func getPartitionNum(o *option) (int32, error) {
 	}
 	defer admin.Close()
 
+	topics := strings.Split(o.topic, ",")
+	maxPartitionNum := int32(0)
 	timeout := 3000
-	for i := 0; i <= 30; i++ {
-		resp, err := admin.GetMetadata(&o.topic, false, timeout)
-		if err != nil {
-			if err.(kafka.Error).Code() == kafka.ErrTransport {
-				log.Info("retry get partition number", zap.Int("retryTime", i), zap.Int("timeout", timeout))
-				timeout += 100
-				continue
+	for _, topic := range topics {
+		topic = strings.TrimSpace(topic)
+		if topic == "" {
+			continue
+		}
+		found := false
+		for i := 0; i <= 30; i++ {
+			resp, err := admin.GetMetadata(&topic, false, timeout)
+			if err != nil {
+				var kafkaErr kafka.Error
+				if errors.As(err, &kafkaErr) && kafkaErr.Code() == kafka.ErrTransport {
+					log.Info("retry get partition number", zap.String("topic", topic), zap.Int("retryTime", i), zap.Int("timeout", timeout))
+					timeout += 100
+					continue
+				}
+				return 0, errors.Trace(err)
 			}
-			return 0, errors.Trace(err)
+
+			topicDetail, ok := resp.Topics[topic]
+			if ok && topicDetail.Error.Code() == kafka.ErrNoError {
+				numPartitions := int32(len(topicDetail.Partitions))
+				log.Info("get partition number of topic",
+					zap.String("topic", topic),
+					zap.Int32("partitionNum", numPartitions))
+				if numPartitions > maxPartitionNum {
+					maxPartitionNum = numPartitions
+				}
+				found = true
+				break
+			}
+			log.Info("retry get partition number", zap.String("topic", topic))
+			time.Sleep(1 * time.Second)
 		}
-		if topicDetail, ok := resp.Topics[o.topic]; ok {
-			numPartitions := int32(len(topicDetail.Partitions))
-			log.Info("get partition number of topic",
-				zap.String("topic", o.topic),
-				zap.Int32("partitionNum", numPartitions))
-			return numPartitions, nil
+		if !found {
+			return 0, errors.Errorf("get partition number(%s) timeout", topic)
 		}
-		log.Info("retry get partition number", zap.String("topic", o.topic))
-		time.Sleep(1 * time.Second)
 	}
-	return 0, errors.Errorf("get partition number(%s) timeout", o.topic)
+	if maxPartitionNum == 0 {
+		return 0, errors.Errorf("get partition number(%s) timeout", o.topic)
+	}
+	return maxPartitionNum, nil
 }
 
 type consumer struct {

cmd/kafka-consumer/option.go

@@ -122,11 +122,11 @@ func (o *option) Adjust(upstreamURIStr string, configFile string) {
 		}
 		o.partitionNum = int32(c)
 	}
-	partitionNum, err := getPartitionNum(o)
-	if err != nil {
-		log.Panic("cannot get the partition number", zap.String("topic", o.topic), zap.Error(err))
-	}
 	if o.partitionNum == 0 {
+		partitionNum, err := getPartitionNum(o)
+		if err != nil {
+			log.Panic("cannot get the partition number", zap.String("topic", o.topic), zap.Error(err))
+		}
 		o.partitionNum = partitionNum
 	}
 

downstreamadapter/sink/eventrouter/topic/expression.go

@@ -68,15 +68,14 @@ func (e Expression) validate() error {
 		return nil
 	}
 
-	return errors.ErrKafkaInvalidTopicExpression.GenWithStackByArgs(e)
+	return errors.ErrKafkaInvalidConfig.GenWithStack("invalid topic expression: %s", e)
 }
 
 // ValidateForAvro checks whether topic pattern is {schema}_{table}, the only allowed
 func (e Expression) validateForAvro() error {
 	if ok := avroTopicNameRE.MatchString(string(e)); !ok {
-		return errors.ErrKafkaInvalidTopicExpression.GenWithStackByArgs(e,
-			"topic rule for Avro must contain {schema} and {table}",
-		)
+		return errors.ErrKafkaInvalidConfig.GenWithStack(
+			"invalid topic expression %s: topic rule for Avro must contain {schema} and {table}", e)
 	}
 
 	return nil

downstreamadapter/sink/eventrouter/topic/expression_test.go

@@ -265,11 +265,11 @@ func TestInvalidExpression(t *testing.T) {
 	topicExpr := Expression(invalidExpr)
 
 	err := topicExpr.validate()
-	require.ErrorIs(t, err, errors.ErrKafkaInvalidTopicExpression)
+	require.ErrorIs(t, err, errors.ErrKafkaInvalidConfig)
 	require.ErrorContains(t, err, invalidExpr)
 
 	err = topicExpr.validateForAvro()
-	require.ErrorIs(t, err, errors.ErrKafkaInvalidTopicExpression)
+	require.ErrorIs(t, err, errors.ErrKafkaInvalidConfig)
 	require.ErrorContains(t, err, "Avro")
 	require.ErrorContains(t, err, invalidExpr)
 }

downstreamadapter/sink/kafka/helper.go

@@ -21,23 +21,24 @@ import (
 	"github.com/pingcap/ticdc/downstreamadapter/sink/eventrouter"
 	"github.com/pingcap/ticdc/downstreamadapter/sink/helper"
 	"github.com/pingcap/ticdc/downstreamadapter/sink/topicmanager"
-	commonType "github.com/pingcap/ticdc/pkg/common"
+	"github.com/pingcap/ticdc/pkg/common"
 	"github.com/pingcap/ticdc/pkg/config"
-	"github.com/pingcap/ticdc/pkg/errors"
 	"github.com/pingcap/ticdc/pkg/sink/codec"
-	"github.com/pingcap/ticdc/pkg/sink/codec/common"
+	codecCommon "github.com/pingcap/ticdc/pkg/sink/codec/common"
 	"github.com/pingcap/ticdc/pkg/sink/kafka"
+	"github.com/pingcap/ticdc/pkg/sink/kafka/claimcheck"
 	"github.com/pingcap/tidb/br/pkg/utils"
 )
 
 type components struct {
 	encoderGroup   codec.EncoderGroup
-	encoder        common.EventEncoder
+	encoder        codecCommon.EventEncoder
 	columnSelector *columnselector.ColumnSelectors
 	eventRouter    *eventrouter.EventRouter
 	topicManager   topicmanager.TopicManager
 	adminClient    kafka.ClusterAdminClient
 	factory        kafka.Factory
+	claimCheck     *claimcheck.ClaimCheck
 }
 
 func (c components) close() {
@@ -47,93 +48,94 @@ func (c components) close() {
 	if c.topicManager != nil {
 		c.topicManager.Close()
 	}
+	if c.claimCheck != nil {
+		c.claimCheck.Close()
+	}
 }
 
-func newKafkaSinkComponentWithFactory(ctx context.Context,
-	changefeedID commonType.ChangeFeedID,
+func newKafkaSinkComponent(
+	ctx context.Context,
+	changefeedID common.ChangeFeedID,
 	sinkURI *url.URL,
 	sinkConfig *config.SinkConfig,
-	factoryCreator kafka.FactoryCreator,
 ) (components, config.Protocol, error) {
-	kafkaComponent := components{}
+	var (
+		comp components
+		err  error
+	)
+	// must release resources when error occurs.
+	defer func() {
+		if err != nil {
+			comp.close()
+		}
+	}()
 	protocol, err := helper.GetProtocol(utils.GetOrZero(sinkConfig.Protocol))
 	if err != nil {
-		return kafkaComponent, config.ProtocolUnknown, errors.Trace(err)
+		return comp, config.ProtocolUnknown, err
 	}
 
 	topic, err := helper.GetTopic(sinkURI)
 	if err != nil {
-		return kafkaComponent, protocol, errors.Trace(err)
+		return comp, protocol, err
 	}
 
 	options := kafka.NewOptions()
 	if err = options.Apply(changefeedID, sinkURI, sinkConfig); err != nil {
-		return kafkaComponent, protocol, errors.WrapError(errors.ErrKafkaInvalidConfig, err)
+		return comp, protocol, err
 	}
 	options.Topic = topic
 
-	kafkaComponent.factory, err = factoryCreator(ctx, options, changefeedID)
+	comp.factory, err = kafka.NewSaramaFactory(ctx, options, changefeedID)
 	if err != nil {
-		return kafkaComponent, protocol, errors.WrapError(errors.ErrKafkaNewProducer, err)
+		return comp, protocol, err
 	}
 
-	kafkaComponent.eventRouter, err = eventrouter.NewEventRouter(
-		sinkConfig, topic, false, protocol == config.ProtocolAvro)
+	isAvroLike := protocol == config.ProtocolAvro
+	comp.eventRouter, err = eventrouter.NewEventRouter(
+		sinkConfig, topic, false, isAvroLike)
 	if err != nil {
-		return kafkaComponent, protocol, errors.Trace(err)
+		return comp, protocol, err
 	}
 
-	kafkaComponent.columnSelector, err = columnselector.New(sinkConfig)
+	comp.columnSelector, err = columnselector.New(sinkConfig)
 	if err != nil {
-		return kafkaComponent, protocol, errors.Trace(err)
+		return comp, protocol, err
 	}
 
 	encoderConfig, err := helper.GetEncoderConfig(changefeedID, sinkURI, protocol, sinkConfig, options.MaxMessageBytes)
 	if err != nil {
-		return kafkaComponent, protocol, errors.Trace(err)
+		return comp, protocol, err
 	}
 
-	kafkaComponent.encoderGroup, err = codec.NewEncoderGroup(ctx, sinkConfig, encoderConfig, changefeedID)
+	comp.claimCheck, err = claimcheck.New(ctx, encoderConfig.LargeMessageHandle, changefeedID)
 	if err != nil {
-		return kafkaComponent, protocol, errors.Trace(err)
+		return comp, protocol, err
 	}
 
-	kafkaComponent.encoder, err = codec.NewEventEncoder(ctx, encoderConfig)
+	comp.encoderGroup, err = codec.NewEncoderGroup(ctx, sinkConfig, encoderConfig, comp.claimCheck, changefeedID)
 	if err != nil {
-		return kafkaComponent, protocol, errors.Trace(err)
+		return comp, protocol, err
 	}
 
-	kafkaComponent.adminClient, err = kafkaComponent.factory.AdminClient(ctx)
+	comp.encoder, err = codec.NewEventEncoder(ctx, encoderConfig, comp.claimCheck)
 	if err != nil {
-		return kafkaComponent, protocol, errors.WrapError(errors.ErrKafkaNewProducer, err)
+		return comp, protocol, err
 	}
 
-	// We must close adminClient when this func return cause by an error
-	// otherwise the adminClient will never be closed and lead to a goroutine leak.
-	defer func() {
-		if err != nil && kafkaComponent.adminClient != nil {
-			kafkaComponent.adminClient.Close()
-		}
-	}()
+	comp.adminClient, err = comp.factory.AdminClient(ctx)
+	if err != nil {
+		return comp, protocol, err
+	}
 
-	kafkaComponent.topicManager, err = topicmanager.GetTopicManagerAndTryCreateTopic(
+	comp.topicManager, err = topicmanager.GetTopicManagerAndTryCreateTopic(
 		ctx,
 		changefeedID,
 		topic,
 		options.DeriveTopicConfig(),
-		kafkaComponent.adminClient,
+		comp.adminClient,
 	)
 	if err != nil {
-		return kafkaComponent, protocol, errors.Trace(err)
+		return comp, protocol, err
 	}
-	return kafkaComponent, protocol, nil
-}
-
-func newKafkaSinkComponent(
-	ctx context.Context,
-	changefeedID commonType.ChangeFeedID,
-	sinkURI *url.URL,
-	sinkConfig *config.SinkConfig,
-) (components, config.Protocol, error) {
-	return newKafkaSinkComponentWithFactory(ctx, changefeedID, sinkURI, sinkConfig, kafka.NewSaramaFactory)
+	return comp, protocol, nil
 }