kafka: introduce franz-go as the kafka client

downstreamadapter/sink/kafka/helper.go

@@ -51,6 +51,9 @@ func (c components) close() {
 	if c.claimCheck != nil {
 		c.claimCheck.Close()
 	}
+	if c.factory != nil {
+		c.factory.Close()
+	}
 }
 
 func newKafkaSinkComponent(
@@ -85,7 +88,7 @@ func newKafkaSinkComponent(
 	}
 	options.Topic = topic
 
-	comp.factory, err = kafka.NewSaramaFactory(ctx, options, changefeedID)
+	comp.factory, err = kafka.NewFactory(ctx, options, changefeedID)
 	if err != nil {
 		return comp, protocol, err
 	}
@@ -109,6 +112,9 @@ func newKafkaSinkComponent(
 	if err != nil {
 		return comp, protocol, err
 	}
+	if options.Client == kafka.KafkaClientFranz {
+		encoderConfig.WithKafkaRecordBatchSize()
+	}
 
 	comp.claimCheck, err = claimcheck.New(ctx, encoderConfig.LargeMessageHandle, changefeedID)
 	if err != nil {

downstreamadapter/sink/kafka/sink.go

@@ -72,10 +72,6 @@ func (s *sink) SinkType() common.SinkType {
 	return common.KafkaSinkType
 }
 
-var createKafkaFactory = func(createSaramaFactory func() (kafka.Factory, error)) (kafka.Factory, error) {
-	return createSaramaFactory()
-}
-
 func Verify(ctx context.Context, changefeedID common.ChangeFeedID, uri *url.URL, sinkConfig *config.SinkConfig) error {
 	protocol, err := helper.GetProtocol(util.GetOrZero(sinkConfig.Protocol))
 	if err != nil {
@@ -100,6 +96,9 @@ func Verify(ctx context.Context, changefeedID common.ChangeFeedID, uri *url.URL,
 	if err != nil {
 		return err
 	}
+	if options.Client == kafka.KafkaClientFranz {
+		encoderConfig.WithKafkaRecordBatchSize()
+	}
 
 	claimCheck, err := claimcheck.New(ctx, encoderConfig.LargeMessageHandle, changefeedID)
 	if err != nil {
@@ -116,12 +115,11 @@ func Verify(ctx context.Context, changefeedID common.ChangeFeedID, uri *url.URL,
 		return err
 	}
 
-	factory, err := createKafkaFactory(func() (kafka.Factory, error) {
-		return kafka.NewSaramaFactory(ctx, options, changefeedID)
-	})
+	factory, err := kafka.NewFactory(ctx, options, changefeedID)
 	if err != nil {
 		return err
 	}
+	defer factory.Close()
 
 	adminClient, err := factory.AdminClient(ctx)
 	if err != nil {
@@ -427,7 +425,6 @@ func (s *sink) sendMessages(ctx context.Context) error {
 			for _, message := range future.Messages {
 				start := time.Now()
 				if err = s.statistics.RecordBatchExecution(func() (int, int64, error) {
-					message.SetPartitionKey(future.Key.PartitionKey)
 					if err = s.dmlProducer.AsyncSend(
 						ctx,
 						future.Key.Topic,
@@ -471,11 +468,11 @@ func (s *sink) sendDDLEvent(event *commonEvent.DDLEvent) error {
 		ddlType := e.GetDDLType().String()
 		if s.partitionRule == helper.PartitionAll {
 			err = s.statistics.RecordDDLExecution(func() (string, error) {
-				return ddlType, s.ddlProducer.SendMessages(topic, partitionNum, message)
+				return ddlType, s.ddlProducer.SendMessages(s.ctx, topic, partitionNum, message)
 			})
 		} else {
 			err = s.statistics.RecordDDLExecution(func() (string, error) {
-				return ddlType, s.ddlProducer.SendMessage(topic, 0, message)
+				return ddlType, s.ddlProducer.SendMessage(s.ctx, topic, 0, message)
 			})
 		}
 		if err != nil {
@@ -541,7 +538,7 @@ func (s *sink) sendCheckpoint(ctx context.Context) error {
 				if err != nil {
 					return err
 				}
-				err = s.ddlProducer.SendMessages(topic, partitionNum, msg)
+				err = s.ddlProducer.SendMessages(ctx, topic, partitionNum, msg)
 				if err != nil {
 					return err
 				}
@@ -552,7 +549,7 @@ func (s *sink) sendCheckpoint(ctx context.Context) error {
 					if err != nil {
 						return err
 					}
-					err = s.ddlProducer.SendMessages(topic, partitionNum, msg)
+					err = s.ddlProducer.SendMessages(ctx, topic, partitionNum, msg)
 					if err != nil {
 						return err
 					}

downstreamadapter/sink/kafka/sink_test.go

@@ -35,6 +35,7 @@ import (
 	"github.com/pingcap/ticdc/pkg/sink/kafka"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/stretchr/testify/require"
+	"github.com/twmb/franz-go/pkg/kfake"
 	"go.uber.org/atomic"
 )
 
@@ -91,6 +92,9 @@ func TestSinkWorkersReturnContextError(t *testing.T) {
 }
 
 func TestVerifyInvalidConfig(t *testing.T) {
+	cluster := kfake.MustCluster(kfake.NumBrokers(1), kfake.SeedTopics(1, kafkaSinkTestTopic))
+	defer cluster.Close()
+
 	schemaRegistry := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
 		http.Error(w, "invalid response", http.StatusInternalServerError)
 	}))
@@ -101,28 +105,10 @@ func TestVerifyInvalidConfig(t *testing.T) {
 		Protocol:       &avroProtocol,
 		SchemaRegistry: &schemaRegistry.URL,
 	}
-	sinkURI, err := url.Parse("kafka://127.0.0.1:9092/" + kafkaSinkTestTopic +
+	sinkURI, err := url.Parse("kafka://" + cluster.ListenAddrs()[0] + "/" + kafkaSinkTestTopic +
 		"?required-acks=1&kafka-version=2.4.0")
 	require.NoError(t, err)
 
-	ctrl := gomock.NewController(t)
-	adminClient := kafka.NewMockAdminClient(ctrl)
-	factory := kafka.NewMockFactory(ctrl)
-	gomock.InOrder(
-		factory.EXPECT().AdminClient(gomock.Any()).Return(adminClient, nil),
-		adminClient.EXPECT().GetTopicsMeta([]string{kafkaSinkTestTopic}, false).Return(
-			map[string]kafka.TopicDetail{kafkaSinkTestTopic: {Name: kafkaSinkTestTopic}}, nil),
-		adminClient.EXPECT().Close(),
-	)
-
-	originalCreateKafkaFactory := createKafkaFactory
-	createKafkaFactory = func(_ func() (kafka.Factory, error)) (kafka.Factory, error) {
-		return factory, nil
-	}
-	t.Cleanup(func() {
-		createKafkaFactory = originalCreateKafkaFactory
-	})
-
 	changefeedID := common.NewChangefeedID4Test("test", "verify-invalid-config")
 	err = Verify(context.Background(), changefeedID, sinkURI, sinkConfig)
 	require.ErrorContains(t, err, "ErrAvroSchemaAPIError")
@@ -208,7 +194,7 @@ func TestKafkaSinkBasicFunctionality(t *testing.T) {
 			}
 			return nil
 		}).Times(2)
-	syncProducer.EXPECT().SendMessages(gomock.Any(), int32(1), gomock.Any()).Return(nil)
+	syncProducer.EXPECT().SendMessages(gomock.Any(), gomock.Any(), int32(1), gomock.Any()).Return(nil)
 	defer cancel()
 	go kafkaSink.Run(ctx)
 
@@ -248,6 +234,7 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 		gomock.InOrder(
 			adminClient.EXPECT().Close(),
 			topicManager.EXPECT().Close(),
+			factory.EXPECT().Close(),
 		)
 
 		kafkaSink, err := newWithComponents(
@@ -276,6 +263,7 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 			asyncProducer.EXPECT().Close(),
 			adminClient.EXPECT().Close(),
 			topicManager.EXPECT().Close(),
+			factory.EXPECT().Close(),
 		)
 
 		kafkaSink, err := newWithComponents(
@@ -307,6 +295,7 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 			asyncProducer.EXPECT().Close().Do(func() { closeCount.Add(1) }),
 			adminClient.EXPECT().Close().Do(func() { closeCount.Add(1) }),
 			topicManager.EXPECT().Close().Do(func() { closeCount.Add(1) }),
+			factory.EXPECT().Close().Do(func() { closeCount.Add(1) }),
 		)
 
 		kafkaSink, err := newWithComponents(
@@ -322,7 +311,7 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 		require.True(t, kafkaSink.IsNormal())
 
 		kafkaSink.Close()
-		require.Equal(t, int64(4), closeCount.Load())
+		require.Equal(t, int64(5), closeCount.Load())
 		require.False(t, kafkaSink.IsNormal())
 		kafkaSink.AddDMLEvent(&commonEvent.DMLEvent{})
 		require.Zero(t, kafkaSink.eventChan.Len())
@@ -443,8 +432,8 @@ func TestKafkaSinkDDL(t *testing.T) {
 		kafkaSink, topicManager, _, syncProducer := newKafkaSinkForTest(
 			t, t.Context(), config.ProtocolOpen, &config.SinkConfig{})
 		topicManager.EXPECT().GetPartitionNum(gomock.Any(), kafkaSinkTestTopic).Return(int32(4), nil)
-		syncProducer.EXPECT().SendMessages(kafkaSinkTestTopic, int32(4), gomock.Any()).
-			DoAndReturn(func(_ string, _ int32, message *codecCommon.Message) error {
+		syncProducer.EXPECT().SendMessages(gomock.Any(), kafkaSinkTestTopic, int32(4), gomock.Any()).
+			DoAndReturn(func(_ context.Context, _ string, _ int32, message *codecCommon.Message) error {
 				require.NotEmpty(t, message.Key)
 				require.NotEmpty(t, message.Value)
 				return nil
@@ -457,8 +446,8 @@ func TestKafkaSinkDDL(t *testing.T) {
 		kafkaSink, topicManager, _, syncProducer := newKafkaSinkForTest(
 			t, t.Context(), config.ProtocolCanalJSON, &config.SinkConfig{})
 		topicManager.EXPECT().GetPartitionNum(gomock.Any(), kafkaSinkTestTopic).Return(int32(4), nil)
-		syncProducer.EXPECT().SendMessage(kafkaSinkTestTopic, int32(0), gomock.Any()).
-			DoAndReturn(func(_ string, _ int32, message *codecCommon.Message) error {
+		syncProducer.EXPECT().SendMessage(gomock.Any(), kafkaSinkTestTopic, int32(0), gomock.Any()).
+			DoAndReturn(func(_ context.Context, _ string, _ int32, message *codecCommon.Message) error {
 				require.NotEmpty(t, message.Value)
 				return nil
 			})
@@ -480,7 +469,7 @@ func TestKafkaSinkDDL(t *testing.T) {
 			t, t.Context(), config.ProtocolOpen, &config.SinkConfig{})
 		cause := errors.ErrKafkaSendMessage.GenWithStackByArgs()
 		topicManager.EXPECT().GetPartitionNum(gomock.Any(), kafkaSinkTestTopic).Return(int32(2), nil)
-		syncProducer.EXPECT().SendMessages(kafkaSinkTestTopic, int32(2), gomock.Any()).Return(cause)
+		syncProducer.EXPECT().SendMessages(gomock.Any(), kafkaSinkTestTopic, int32(2), gomock.Any()).Return(cause)
 
 		require.Equal(t, cause, kafkaSink.WriteBlockEvent(ddlEvent))
 		require.False(t, kafkaSink.IsNormal())
@@ -508,8 +497,8 @@ func TestKafkaSinkCheckpoint(t *testing.T) {
 		kafkaSink, topicManager, _, syncProducer := newKafkaSinkForTest(
 			t, t.Context(), config.ProtocolOpen, &config.SinkConfig{})
 		topicManager.EXPECT().GetPartitionNum(gomock.Any(), kafkaSinkTestTopic).Return(int32(3), nil)
-		syncProducer.EXPECT().SendMessages(kafkaSinkTestTopic, int32(3), gomock.Any()).
-			DoAndReturn(func(_ string, _ int32, message *codecCommon.Message) error {
+		syncProducer.EXPECT().SendMessages(gomock.Any(), kafkaSinkTestTopic, int32(3), gomock.Any()).
+			DoAndReturn(func(_ context.Context, _ string, _ int32, message *codecCommon.Message) error {
 				require.NotEmpty(t, message.Key)
 				return nil
 			})
@@ -535,7 +524,7 @@ func TestKafkaSinkCheckpoint(t *testing.T) {
 		partitionCounts := map[string]int32{"topic-a": 2, "topic-b": 3, kafkaSinkTestTopic: 4}
 		for topic, partitionCount := range partitionCounts {
 			topicManager.EXPECT().GetPartitionNum(gomock.Any(), topic).Return(partitionCount, nil)
-			syncProducer.EXPECT().SendMessages(topic, partitionCount, gomock.Any()).Return(nil)
+			syncProducer.EXPECT().SendMessages(gomock.Any(), topic, partitionCount, gomock.Any()).Return(nil)
 		}
 		kafkaSink.checkpointChan <- 100
 		close(kafkaSink.checkpointChan)
@@ -567,7 +556,7 @@ func TestKafkaSinkCheckpoint(t *testing.T) {
 		// return the error and stop, so exactly one GetPartitionNum and one
 		// SendMessages call are expected regardless of the topic order.
 		topicManager.EXPECT().GetPartitionNum(gomock.Any(), gomock.Any()).Return(int32(2), nil)
-		syncProducer.EXPECT().SendMessages(gomock.Any(), int32(2), gomock.Any()).Return(cause)
+		syncProducer.EXPECT().SendMessages(gomock.Any(), gomock.Any(), int32(2), gomock.Any()).Return(cause)
 		kafkaSink.checkpointChan <- 100
 
 		require.Equal(t, cause, kafkaSink.sendCheckpoint(t.Context()))
@@ -623,6 +612,7 @@ func newKafkaSinkForTest(
 	factory.EXPECT().AsyncProducer(gomock.Any()).Return(asyncProducer, nil)
 	factory.EXPECT().SyncProducer(gomock.Any()).Return(syncProducer, nil)
 	factory.EXPECT().MetricsCollector(nil).Return(noopMetricsCollector{})
+	factory.EXPECT().Close().AnyTimes()
 
 	kafkaSink, err := newWithComponents(ctx, changefeedID, common.DefaultKeyspaceID, protocol, components{
 		encoderGroup:   encoderGroup,

downstreamadapter/sink/topicmanager/kafka_topic_manager.go

@@ -125,7 +125,7 @@ func (m *kafkaTopicManager) backgroundRefreshMeta(ctx context.Context) {
 		case <-ticker.C:
 			// We ignore the error here, because the error may be caused by the
 			// network problem, and we can try to get the metadata next time.
-			topicPartitionNums, _ := m.fetchAllTopicsPartitionsNum()
+			topicPartitionNums, _ := m.fetchAllTopicsPartitionsNum(ctx)
 			for topic, partitionNum := range topicPartitionNums {
 				m.tryUpdatePartitionsAndLogging(topic, partitionNum)
 			}
@@ -157,15 +157,15 @@ func (m *kafkaTopicManager) tryUpdatePartitionsAndLogging(topic string, partitio
 // The error returned by this method could be a transient error that is fixable by the underlying logic.
 // When handling this error, please be cautious.
 // If you simply throw the error to the caller, it may impact the robustness of your program.
-func (m *kafkaTopicManager) fetchAllTopicsPartitionsNum() (map[string]int32, error) {
+func (m *kafkaTopicManager) fetchAllTopicsPartitionsNum(ctx context.Context) (map[string]int32, error) {
 	var topics []string
 	m.topics.Range(func(key, _ any) bool {
 		topics = append(topics, key.(string))
 		return true
 	})
 
 	start := time.Now()
-	numPartitions, err := m.admin.GetTopicsPartitionsNum(topics)
+	numPartitions, err := m.admin.GetTopicsPartitionsNum(ctx, topics)
 	if err != nil {
 		log.Warn(
 			"kafka topic metadata refresh failed",
@@ -202,7 +202,7 @@ func (m *kafkaTopicManager) waitUntilTopicVisible(
 	err := retry.Do(ctx, func() error {
 		// ignoreTopicError is set to false since we just create the topic,
 		// make sure the topic is visible.
-		meta, err := m.admin.GetTopicsMeta(topics, false)
+		meta, err := m.admin.GetTopicsMeta(ctx, topics, false)
 		if err != nil {
 			return err
 		}
@@ -231,20 +231,17 @@ func (m *kafkaTopicManager) waitUntilTopicVisible(
 
 // createTopic creates a topic with the given name
 // and returns the number of partitions.
-func (m *kafkaTopicManager) createTopic(
-	_ context.Context,
-	topicName string,
-) (int32, error) {
+func (m *kafkaTopicManager) createTopic(ctx context.Context, topicName string) (int32, error) {
 	if !m.cfg.AutoCreate {
 		return 0, errors.ErrKafkaInvalidConfig.GenWithStack("`auto-create-topic` is false, and %s not found", topicName)
 	}
 
-	if err := m.cfg.ValidateReplicationFactor(m.admin); err != nil {
+	if err := m.cfg.ValidateReplicationFactor(ctx, m.admin); err != nil {
 		return 0, err
 	}
 
 	start := time.Now()
-	err := m.admin.CreateTopic(&kafka.TopicDetail{
+	err := m.admin.CreateTopic(ctx, &kafka.TopicDetail{
 		Name:              topicName,
 		NumPartitions:     m.cfg.PartitionNum,
 		ReplicationFactor: m.cfg.ReplicationFactor,
@@ -274,7 +271,7 @@ func (m *kafkaTopicManager) CreateTopicAndWaitUntilVisible(
 	ctx context.Context, topicName string,
 ) (int32, error) {
 	// If the topic is not in the cache, try to get its metadata.
-	topicDetails, err := m.admin.GetTopicsMeta([]string{topicName}, false)
+	topicDetails, err := m.admin.GetTopicsMeta(ctx, []string{topicName}, false)
 	if err == nil {
 		if numPartition, ok := m.tryStoreTopicMeta(topicName, topicDetails); ok {
 			return numPartition, nil