kafka: improve stability when creating many topics with Kafka v4

downstreamadapter/sink/kafka/helper.go

@@ -42,12 +42,12 @@ type components struct {
 }
 
 func (c components) close() {
-	if c.adminClient != nil {
-		c.adminClient.Close()
-	}
 	if c.topicManager != nil {
 		c.topicManager.Close()
 	}
+	if c.factory != nil {
+		c.factory.Close()
+	}
 	if c.claimCheck != nil {
 		c.claimCheck.Close()
 	}

downstreamadapter/sink/kafka/sink.go

@@ -122,12 +122,12 @@ func Verify(ctx context.Context, changefeedID common.ChangeFeedID, uri *url.URL,
 	if err != nil {
 		return err
 	}
+	defer factory.Close()
 
 	adminClient, err := factory.AdminClient(ctx)
 	if err != nil {
 		return err
 	}
-	defer adminClient.Close()
 
 	err = topicmanager.EnsureTopic(ctx, changefeedID, topic, options.DeriveTopicConfig(), adminClient)
 	if err != nil {
@@ -168,13 +168,15 @@ func newWithComponents(
 		if err == nil {
 			return
 		}
+		// Release shared Kafka resources first to help unblock producer shutdown
+		// when Kafka is unhealthy.
+		comp.close()
 		if syncProducer != nil {
 			syncProducer.Close()
 		}
 		if asyncProducer != nil {
 			asyncProducer.Close()
 		}
-		comp.close()
 		statistics.Close()
 	}()
 
@@ -232,6 +234,9 @@ func (s *sink) IsNormal() bool {
 }
 
 func (s *sink) AddDMLEvent(event *commonEvent.DMLEvent) {
+	if !s.isNormal.Load() {
+		return
+	}
 	s.eventChan.Push(event)
 }
 
@@ -260,6 +265,7 @@ func (s *sink) WriteBlockEvent(event commonEvent.BlockEvent) error {
 }
 
 func (s *sink) close() {
+	s.isNormal.Store(false)
 	s.eventChan.Close()
 	s.rowChan.Close()
 }
@@ -315,6 +321,11 @@ func (s *sink) calculateKeyPartitions(ctx context.Context) error {
 			if err != nil {
 				return err
 			}
+			select {
+			case <-ctx.Done():
+				return context.Cause(ctx)
+			default:
+			}
 			s.rowChan.Push(events...)
 		}
 	}
@@ -567,9 +578,12 @@ func (s *sink) getAllTableNames(ts uint64) []*commonEvent.SchemaTableName {
 }
 
 func (s *sink) Close() {
+	s.close()
+	// Release shared Kafka resources before closing producers to help unblock
+	// their shutdown.
+	s.comp.close()
 	s.ddlProducer.Close()
 	s.dmlProducer.Close()
-	s.comp.close()
 	s.statistics.Close()
 }
 

downstreamadapter/sink/kafka/sink_test.go

@@ -112,7 +112,7 @@ func TestVerifyInvalidConfig(t *testing.T) {
 		factory.EXPECT().AdminClient(gomock.Any()).Return(adminClient, nil),
 		adminClient.EXPECT().GetTopicsMeta([]string{kafkaSinkTestTopic}, true).Return(
 			map[string]kafka.TopicDetail{kafkaSinkTestTopic: {Name: kafkaSinkTestTopic}}, nil),
-		adminClient.EXPECT().Close(),
+		factory.EXPECT().Close(),
 	)
 
 	originalCreateKafkaFactory := createKafkaFactory
@@ -245,8 +245,10 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 		cause := errors.ErrKafkaSendMessage.GenWithStackByArgs()
 
 		factory.EXPECT().AsyncProducer(gomock.Any()).Return(nil, cause)
-		adminClient.EXPECT().Close()
-		topicManager.EXPECT().Close()
+		gomock.InOrder(
+			topicManager.EXPECT().Close(),
+			factory.EXPECT().Close(),
+		)
 
 		kafkaSink, err := newWithComponents(
 			t.Context(),
@@ -270,9 +272,11 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 
 		factory.EXPECT().AsyncProducer(gomock.Any()).Return(asyncProducer, nil)
 		factory.EXPECT().SyncProducer(gomock.Any()).Return(nil, cause)
-		asyncProducer.EXPECT().Close()
-		adminClient.EXPECT().Close()
-		topicManager.EXPECT().Close()
+		gomock.InOrder(
+			topicManager.EXPECT().Close(),
+			factory.EXPECT().Close(),
+			asyncProducer.EXPECT().Close(),
+		)
 
 		kafkaSink, err := newWithComponents(
 			t.Context(),
@@ -298,10 +302,12 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 		factory.EXPECT().AsyncProducer(gomock.Any()).Return(asyncProducer, nil)
 		factory.EXPECT().SyncProducer(gomock.Any()).Return(syncProducer, nil)
 		factory.EXPECT().MetricsCollector(adminClient).Return(noopMetricsCollector{})
-		asyncProducer.EXPECT().Close().Do(func() { closeCount.Add(1) })
-		syncProducer.EXPECT().Close().Do(func() { closeCount.Add(1) })
-		adminClient.EXPECT().Close().Do(func() { closeCount.Add(1) })
-		topicManager.EXPECT().Close().Do(func() { closeCount.Add(1) })
+		gomock.InOrder(
+			topicManager.EXPECT().Close().Do(func() { closeCount.Add(1) }),
+			factory.EXPECT().Close().Do(func() { closeCount.Add(1) }),
+			syncProducer.EXPECT().Close().Do(func() { closeCount.Add(1) }),
+			asyncProducer.EXPECT().Close().Do(func() { closeCount.Add(1) }),
+		)
 
 		kafkaSink, err := newWithComponents(
 			t.Context(),
@@ -313,8 +319,20 @@ func TestKafkaSinkConstructionAndCleanup(t *testing.T) {
 
 		require.NoError(t, err)
 		require.Zero(t, closeCount.Load())
+		require.True(t, kafkaSink.IsNormal())
+
 		kafkaSink.Close()
 		require.Equal(t, int64(4), closeCount.Load())
+		require.False(t, kafkaSink.IsNormal())
+		kafkaSink.AddDMLEvent(&commonEvent.DMLEvent{})
+		require.Zero(t, kafkaSink.eventChan.Len())
+
+		_, ok, err := kafkaSink.eventChan.GetWithContext(t.Context())
+		require.NoError(t, err)
+		require.False(t, ok)
+		_, ok, err = kafkaSink.rowChan.GetWithContext(t.Context())
+		require.NoError(t, err)
+		require.False(t, ok)
 	})
 }
 
@@ -394,6 +412,23 @@ func TestKafkaSinkDML(t *testing.T) {
 
 		require.Equal(t, cause, kafkaSink.calculateKeyPartitions(t.Context()))
 	})
+
+	t.Run("canceled after topic lookup", func(t *testing.T) {
+		dmlEvent := eventHelper.DML2Event("test", "t", "insert into t values (4, 'four')")
+		ctx, cancel := context.WithCancelCause(t.Context())
+		kafkaSink, topicManager, _, _ := newKafkaSinkForTest(
+			t, ctx, config.ProtocolOpen, &config.SinkConfig{})
+		cause := errors.ErrKafkaSinkClosed.GenWithStackByArgs()
+		topicManager.EXPECT().GetPartitionNum(gomock.Any(), kafkaSinkTestTopic).
+			DoAndReturn(func(context.Context, string) (int32, error) {
+				cancel(cause)
+				return 1, nil
+			})
+		kafkaSink.AddDMLEvent(dmlEvent)
+
+		require.Equal(t, cause, kafkaSink.calculateKeyPartitions(ctx))
+		require.Zero(t, kafkaSink.rowChan.Len())
+	})
 }
 
 func TestKafkaSinkDDL(t *testing.T) {
@@ -589,6 +624,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

@@ -46,6 +46,7 @@ type kafkaTopicManager struct {
 	topics sync.Map
 	// cancel is used to cancel the background goroutine.
 	cancel context.CancelFunc
+	wg     sync.WaitGroup
 }
 
 // newKafkaTopicManager creates a topic manager without starting background work.
@@ -91,7 +92,9 @@ func GetTopicManagerAndTryCreateTopic(
 	}
 	ctx, cancel := context.WithCancel(ctx)
 	topicManager.cancel = cancel
-	go topicManager.backgroundRefreshMeta(ctx)
+	topicManager.wg.Go(func() {
+		topicManager.backgroundRefreshMeta(ctx)
+	})
 
 	return topicManager, nil
 }
@@ -214,6 +217,9 @@ func (m *kafkaTopicManager) waitUntilTopicVisible(
 	}, retry.WithBackoffBaseDelay(500),
 		retry.WithBackoffMaxDelay(1000),
 		retry.WithMaxTries(6),
+		retry.WithIsRetryableErr(func(err error) bool {
+			return !kafka.IsUnretryableKafkaError(err)
+		}),
 	)
 	if err != nil {
 		log.Warn("kafka topic metadata refresh failed",
@@ -260,8 +266,6 @@ func (m *kafkaTopicManager) createTopic(
 		return 0, err
 	}
 
-	m.tryUpdatePartitionsAndLogging(topicName, m.cfg.PartitionNum)
-
 	return m.cfg.PartitionNum, nil
 }
 
@@ -308,6 +312,7 @@ func (m *kafkaTopicManager) CreateTopicAndWaitUntilVisible(
 	if err != nil {
 		return 0, err
 	}
+	m.tryUpdatePartitionsAndLogging(topicName, partitionNum)
 
 	log.Info(
 		"kafka topic created",
@@ -348,7 +353,10 @@ func (m *kafkaTopicManager) useConfiguredPartitionNum(topicName string, cause er
 	return m.cfg.PartitionNum
 }
 
-// Close exits the background goroutine.
+// Close cancels the background goroutine and waits for it to exit.
 func (m *kafkaTopicManager) Close() {
-	m.cancel()
+	if m.cancel != nil {
+		m.cancel()
+	}
+	m.wg.Wait()
 }