From fc70dbde873be0c09fa0f8f603772d2d95d59a8c Mon Sep 17 00:00:00 2001
From: Ling Jin <7138436+3AceShowHand@users.noreply.github.com>
Date: Tue, 4 Jan 2022 14:14:35 +0800
Subject: [PATCH] metrics(cdc): fix mq sink write row count metrics. (#4192)

close pingcap/tiflow#3431
---
 cdc/sink/codec/avro.go            | 1 +
 cdc/sink/codec/canal.go           | 4 +++-
 cdc/sink/codec/canal_flat.go      | 4 +++-
 cdc/sink/codec/canal_flat_test.go | 2 ++
 cdc/sink/codec/canal_test.go      | 1 +
 cdc/sink/codec/json_test.go       | 1 +
 cdc/sink/codec/maxwell.go         | 1 +
 cdc/sink/codec/maxwell_test.go    | 1 +
 8 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/cdc/sink/codec/avro.go b/cdc/sink/codec/avro.go
index a0a962acb12..13fbcf8fa02 100644
--- a/cdc/sink/codec/avro.go
+++ b/cdc/sink/codec/avro.go
@@ -123,6 +123,7 @@ func (a *AvroEventBatchEncoder) AppendRowChangedEvent(e *model.RowChangedEvent)
 	}
 
 	mqMessage.Key = evlp
+	mqMessage.IncRowsCount()
 	a.resultBuf = append(a.resultBuf, mqMessage)
 
 	return EncoderNeedAsyncWrite, nil
diff --git a/cdc/sink/codec/canal.go b/cdc/sink/codec/canal.go
index 8e3ed61ecfb..20fdfa363ae 100644
--- a/cdc/sink/codec/canal.go
+++ b/cdc/sink/codec/canal.go
@@ -458,7 +458,8 @@ func (d *CanalEventBatchEncoder) EncodeDDLEvent(e *model.DDLEvent) (*MQMessage,
 
 // Build implements the EventBatchEncoder interface
 func (d *CanalEventBatchEncoder) Build() []*MQMessage {
-	if len(d.messages.Messages) == 0 {
+	rowCount := len(d.messages.Messages)
+	if rowCount == 0 {
 		return nil
 	}
 
@@ -472,6 +473,7 @@ func (d *CanalEventBatchEncoder) Build() []*MQMessage {
 		log.Panic("Error when serializing Canal packet", zap.Error(err))
 	}
 	ret := NewMQMessage(config.ProtocolCanal, nil, value, 0, model.MqMessageTypeRow, nil, nil)
+	ret.SetRowsCount(rowCount)
 	d.messages.Reset()
 	d.resetPacket()
 	return []*MQMessage{ret}
diff --git a/cdc/sink/codec/canal_flat.go b/cdc/sink/codec/canal_flat.go
index b369215c709..98bcff8fa55 100644
--- a/cdc/sink/codec/canal_flat.go
+++ b/cdc/sink/codec/canal_flat.go
@@ -366,7 +366,9 @@ func (c *CanalFlatEventBatchEncoder) Build() []*MQMessage {
 			log.Panic("CanalFlatEventBatchEncoder", zap.Error(err))
 			return nil
 		}
-		ret[i] = NewMQMessage(config.ProtocolCanalJSON, nil, value, msg.getTikvTs(), model.MqMessageTypeRow, msg.getSchema(), msg.getTable())
+		m := NewMQMessage(config.ProtocolCanalJSON, nil, value, msg.getTikvTs(), model.MqMessageTypeRow, msg.getSchema(), msg.getTable())
+		m.IncRowsCount()
+		ret[i] = m
 	}
 	c.resolvedBuf = c.resolvedBuf[0:0]
 	return ret
diff --git a/cdc/sink/codec/canal_flat_test.go b/cdc/sink/codec/canal_flat_test.go
index 25fa8e7aac4..fb968f2b2ac 100644
--- a/cdc/sink/codec/canal_flat_test.go
+++ b/cdc/sink/codec/canal_flat_test.go
@@ -318,6 +318,8 @@ func (s *canalFlatSuite) TestBatching(c *check.C) {
 			c.Assert(msgs, check.HasLen, int(resolvedTs-lastResolved))
 
 			for j := range msgs {
+				c.Assert(msgs[j].GetRowsCount(), check.Equals, 1)
+
 				var msg canalFlatMessage
 				err := json.Unmarshal(msgs[j].Value, &msg)
 				c.Assert(err, check.IsNil)
diff --git a/cdc/sink/codec/canal_test.go b/cdc/sink/codec/canal_test.go
index 8b6354a67dd..c25b3cb9800 100644
--- a/cdc/sink/codec/canal_test.go
+++ b/cdc/sink/codec/canal_test.go
@@ -99,6 +99,7 @@ func (s *canalBatchSuite) TestCanalEventBatchEncoder(c *check.C) {
 		c.Assert(res, check.HasLen, 1)
 		c.Assert(res[0].Key, check.IsNil)
 		c.Assert(len(res[0].Value), check.Equals, size)
+		c.Assert(res[0].GetRowsCount(), check.Equals, len(cs))
 
 		packet := &canal.Packet{}
 		err := proto.Unmarshal(res[0].Value, packet)
diff --git a/cdc/sink/codec/json_test.go b/cdc/sink/codec/json_test.go
index d7f1499ec30..387990d44a3 100644
--- a/cdc/sink/codec/json_test.go
+++ b/cdc/sink/codec/json_test.go
@@ -137,6 +137,7 @@ func (s *batchSuite) testBatchCodec(c *check.C, newEncoder func() EventBatchEnco
 		if len(cs) > 0 {
 			res := encoder.Build()
 			c.Assert(res, check.HasLen, 1)
+			c.Assert(res[0].GetRowsCount(), check.Equals, len(cs))
 			decoder, err := newDecoder(res[0].Key, res[0].Value)
 			c.Assert(err, check.IsNil)
 			checkRowDecoder(decoder, cs)
diff --git a/cdc/sink/codec/maxwell.go b/cdc/sink/codec/maxwell.go
index a85d5ce3678..42f54aafc10 100644
--- a/cdc/sink/codec/maxwell.go
+++ b/cdc/sink/codec/maxwell.go
@@ -298,6 +298,7 @@ func (d *MaxwellEventBatchEncoder) Build() []*MQMessage {
 	}
 
 	ret := NewMQMessage(config.ProtocolMaxwell, d.keyBuf.Bytes(), d.valueBuf.Bytes(), 0, model.MqMessageTypeRow, nil, nil)
+	ret.SetRowsCount(d.batchSize)
 	d.Reset()
 	return []*MQMessage{ret}
 }
diff --git a/cdc/sink/codec/maxwell_test.go b/cdc/sink/codec/maxwell_test.go
index 5e6fea4a58e..1f98bedac44 100644
--- a/cdc/sink/codec/maxwell_test.go
+++ b/cdc/sink/codec/maxwell_test.go
@@ -54,6 +54,7 @@ func (s *maxwellbatchSuite) testmaxwellBatchCodec(c *check.C, newEncoder func()
 			continue
 		}
 		c.Assert(messages, check.HasLen, 1)
+		c.Assert(messages[0].GetRowsCount(), check.Equals, len(cs))
 		c.Assert(len(messages[0].Key)+len(messages[0].Value), check.Equals, size)
 	}