fix: [2.4] replicate message exception when the ttMsgEnable config is…

… changed dynamically (#38440) - issue: #38177 - pr: #38178 Signed-off-by: SimFG <bang.fu@zilliz.com>
milvus-io · Dec 14, 2024 · 4c896c6 · 4c896c6
1 parent 1da4ac4
commit 4c896c6
Show file tree

Hide file tree

Showing 11 changed files with 44 additions and 38 deletions.
diff --git a/internal/datanode/data_sync_service.go b/internal/datanode/data_sync_service.go
@@ -283,7 +283,7 @@ func getServiceWithChannel(initCtx context.Context, node *DataNode, info *datapb
 		metrics.DataNodeNumProducers.WithLabelValues(fmt.Sprint(config.serverID)).Inc()
 		log.Info("datanode AsProducer", zap.String("TimeTickChannelName", Params.CommonCfg.DataCoordTimeTick.GetValue()))
 
-		m.EnableProduce(true)
+		m.ForceEnableProduce(true)
 
 		updater = newMqStatsUpdater(config, m)
 	}

diff --git a/internal/datanode/flow_graph_dmstream_input_node_test.go b/internal/datanode/flow_graph_dmstream_input_node_test.go
@@ -103,7 +103,7 @@ func (mtm *mockTtMsgStream) CheckTopicValid(channel string) error {
 	return nil
 }
 
-func (mtm *mockTtMsgStream) EnableProduce(can bool) {
+func (mtm *mockTtMsgStream) ForceEnableProduce(can bool) {
 }
 
 func TestNewDmInputNode(t *testing.T) {

diff --git a/internal/proxy/mock_msgstream_test.go b/internal/proxy/mock_msgstream_test.go
@@ -10,10 +10,10 @@ import (
 
 type mockMsgStream struct {
 	msgstream.MsgStream
-	asProducer    func([]string)
-	setRepack     func(repackFunc msgstream.RepackFunc)
-	close         func()
-	enableProduce func(bool)
+	asProducer         func([]string)
+	setRepack          func(repackFunc msgstream.RepackFunc)
+	close              func()
+	forceEnableProduce func(bool)
 }
 
 func (m *mockMsgStream) AsProducer(producers []string) {
@@ -34,9 +34,9 @@ func (m *mockMsgStream) Close() {
 	}
 }
 
-func (m *mockMsgStream) EnableProduce(enabled bool) {
-	if m.enableProduce != nil {
-		m.enableProduce(enabled)
+func (m *mockMsgStream) ForceEnableProduce(enabled bool) {
+	if m.forceEnableProduce != nil {
+		m.forceEnableProduce(enabled)
 	}
 }
 

diff --git a/internal/proxy/mock_test.go b/internal/proxy/mock_test.go
@@ -311,7 +311,7 @@ func (ms *simpleMockMsgStream) CheckTopicValid(topic string) error {
 	return nil
 }
 
-func (ms *simpleMockMsgStream) EnableProduce(enabled bool) {
+func (ms *simpleMockMsgStream) ForceEnableProduce(enabled bool) {
 }
 
 func newSimpleMockMsgStream() *simpleMockMsgStream {

diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
@@ -272,7 +272,7 @@ func (node *Proxy) Init() error {
 			zap.Error(err))
 		return err
 	}
-	node.replicateMsgStream.EnableProduce(true)
+	node.replicateMsgStream.ForceEnableProduce(true)
 	node.replicateMsgStream.AsProducer([]string{replicateMsgChannel})
 
 	node.sched, err = newTaskScheduler(node.ctx, node.tsoAllocator, node.factory)

diff --git a/internal/proxy/replicate_stream_manager.go b/internal/proxy/replicate_stream_manager.go
@@ -43,7 +43,7 @@ func (m *ReplicateStreamManager) newMsgStreamResource(channel string) resource.N
 		}
 		msgStream.SetRepackFunc(replicatePackFunc)
 		msgStream.AsProducer([]string{channel})
-		msgStream.EnableProduce(true)
+		msgStream.ForceEnableProduce(true)
 
 		res := resource.NewSimpleResource(msgStream, ReplicateMsgStreamTyp, channel, ReplicateMsgStreamExpireTime, func() {
 			msgStream.Close()

diff --git a/internal/proxy/replicate_stream_manager_test.go b/internal/proxy/replicate_stream_manager_test.go
@@ -34,7 +34,7 @@ func TestReplicateManager(t *testing.T) {
 		mockMsgStream.asProducer = func(producers []string) {
 			i++
 		}
-		mockMsgStream.enableProduce = func(b bool) {
+		mockMsgStream.forceEnableProduce = func(b bool) {
 			i++
 		}
 		mockMsgStream.close = func() {

diff --git a/pkg/mq/msgstream/mock_msgstream.go b/pkg/mq/msgstream/mock_msgstream.go
diff --git a/pkg/mq/msgstream/mq_msgstream.go b/pkg/mq/msgstream/mq_msgstream.go
@@ -58,18 +58,19 @@ type mqMsgStream struct {
 	consumers        map[string]mqwrapper.Consumer
 	consumerChannels []string
 
-	repackFunc    RepackFunc
-	unmarshal     UnmarshalDispatcher
-	receiveBuf    chan *MsgPack
-	closeRWMutex  *sync.RWMutex
-	streamCancel  func()
-	bufSize       int64
-	producerLock  *sync.RWMutex
-	consumerLock  *sync.Mutex
-	closed        int32
-	onceChan      sync.Once
-	enableProduce atomic.Value
-	configEvent   config.EventHandler
+	repackFunc         RepackFunc
+	unmarshal          UnmarshalDispatcher
+	receiveBuf         chan *MsgPack
+	closeRWMutex       *sync.RWMutex
+	streamCancel       func()
+	bufSize            int64
+	producerLock       *sync.RWMutex
+	consumerLock       *sync.Mutex
+	closed             int32
+	onceChan           sync.Once
+	ttMsgEnable        atomic.Value
+	forceEnableProduce atomic.Value
+	configEvent        config.EventHandler
 }
 
 // NewMqMsgStream is used to generate a new mqMsgStream object
@@ -104,14 +105,15 @@ func NewMqMsgStream(ctx context.Context,
 		closed:       0,
 	}
 	ctxLog := log.Ctx(ctx)
-	stream.enableProduce.Store(paramtable.Get().CommonCfg.TTMsgEnabled.GetAsBool())
+	stream.forceEnableProduce.Store(false)
+	stream.ttMsgEnable.Store(paramtable.Get().CommonCfg.TTMsgEnabled.GetAsBool())
 	stream.configEvent = config.NewHandler("enable send tt msg "+fmt.Sprint(streamCounter.Inc()), func(event *config.Event) {
 		value, err := strconv.ParseBool(event.Value)
 		if err != nil {
 			ctxLog.Warn("Failed to parse bool value", zap.String("v", event.Value), zap.Error(err))
 			return
 		}
-		stream.enableProduce.Store(value)
+		stream.ttMsgEnable.Store(value)
 		ctxLog.Info("Msg Stream state updated", zap.Bool("can_produce", stream.isEnabledProduce()))
 	})
 	paramtable.Get().Watch(paramtable.Get().CommonCfg.TTMsgEnabled.Key, stream.configEvent)
@@ -265,12 +267,12 @@ func (ms *mqMsgStream) GetProduceChannels() []string {
 	return ms.producerChannels
 }
 
-func (ms *mqMsgStream) EnableProduce(can bool) {
-	ms.enableProduce.Store(can)
+func (ms *mqMsgStream) ForceEnableProduce(can bool) {
+	ms.forceEnableProduce.Store(can)
 }
 
 func (ms *mqMsgStream) isEnabledProduce() bool {
-	return ms.enableProduce.Load().(bool)
+	return ms.forceEnableProduce.Load().(bool) || ms.ttMsgEnable.Load().(bool)
 }
 
 func (ms *mqMsgStream) Produce(msgPack *MsgPack) error {

diff --git a/pkg/mq/msgstream/mq_msgstream_test.go b/pkg/mq/msgstream/mq_msgstream_test.go
@@ -114,6 +114,8 @@ func TestStream_ConfigEvent(t *testing.T) {
 }
 
 func TestStream_PulsarMsgStream_Insert(t *testing.T) {
+	Params.Save(Params.CommonCfg.TTMsgEnabled.Key, "false")
+	defer Params.Remove(Params.CommonCfg.TTMsgEnabled.Key)
 	pulsarAddress := getPulsarAddress()
 	c1, c2 := funcutil.RandomString(8), funcutil.RandomString(8)
 	producerChannels := []string{c1, c2}
@@ -129,12 +131,12 @@ func TestStream_PulsarMsgStream_Insert(t *testing.T) {
 	outputStream := getPulsarOutputStream(ctx, pulsarAddress, consumerChannels, consumerSubName)
 
 	{
-		inputStream.EnableProduce(false)
+		inputStream.ForceEnableProduce(false)
 		err := inputStream.Produce(&msgPack)
 		require.Error(t, err)
 	}
 
-	inputStream.EnableProduce(true)
+	inputStream.ForceEnableProduce(true)
 	err := inputStream.Produce(&msgPack)
 	require.NoErrorf(t, err, fmt.Sprintf("produce error = %v", err))
 
@@ -187,6 +189,8 @@ func TestStream_PulsarMsgStream_TimeTick(t *testing.T) {
 }
 
 func TestStream_PulsarMsgStream_BroadCast(t *testing.T) {
+	Params.Save(Params.CommonCfg.TTMsgEnabled.Key, "false")
+	defer Params.Remove(Params.CommonCfg.TTMsgEnabled.Key)
 	pulsarAddress := getPulsarAddress()
 	c1, c2 := funcutil.RandomString(8), funcutil.RandomString(8)
 	producerChannels := []string{c1, c2}
@@ -202,12 +206,12 @@ func TestStream_PulsarMsgStream_BroadCast(t *testing.T) {
 	outputStream := getPulsarOutputStream(ctx, pulsarAddress, consumerChannels, consumerSubName)
 
 	{
-		inputStream.EnableProduce(false)
+		inputStream.ForceEnableProduce(false)
 		_, err := inputStream.Broadcast(&msgPack)
 		require.Error(t, err)
 	}
 
-	inputStream.EnableProduce(true)
+	inputStream.ForceEnableProduce(true)
 	_, err := inputStream.Broadcast(&msgPack)
 	require.NoErrorf(t, err, fmt.Sprintf("broadcast error = %v", err))
 

diff --git a/pkg/mq/msgstream/msgstream.go b/pkg/mq/msgstream/msgstream.go
@@ -70,7 +70,7 @@ type MsgStream interface {
 	GetLatestMsgID(channel string) (MessageID, error)
 	CheckTopicValid(channel string) error
 
-	EnableProduce(can bool)
+	ForceEnableProduce(can bool)
 }
 
 type Factory interface {