From fcabd2d8c680ae2d0a84bc26ef9946fc790829a3 Mon Sep 17 00:00:00 2001
From: xiehan <52160700+Barenboim@users.noreply.github.com>
Date: Tue, 20 Aug 2024 19:52:57 +0800
Subject: [PATCH] Update ProtocolMessage and ProtocolWrapper moving. (#1612)

---
 src/protocol/PackageWrapper.cc |  4 ++--
 src/protocol/ProtocolMessage.h | 20 ++++++++++++--------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/protocol/PackageWrapper.cc b/src/protocol/PackageWrapper.cc
index ce1152aa..a8fae118 100644
--- a/src/protocol/PackageWrapper.cc
+++ b/src/protocol/PackageWrapper.cc
@@ -39,7 +39,7 @@ int PackageWrapper::encode(struct iovec vectors[], int max)
 		}
 
 		cnt += ret;
-		this->message = this->next_out(this->message);
+		this->set_message(this->next_out(this->message));
 		if (!this->message)
 			return cnt;
 
@@ -57,7 +57,7 @@ int PackageWrapper::append(const void *buf, size_t *size)
 
 	if (ret > 0)
 	{
-		this->message = this->next_in(this->message);
+		this->set_message(this->next_in(this->message));
 		if (this->message)
 		{
 			this->renew();
diff --git a/src/protocol/ProtocolMessage.h b/src/protocol/ProtocolMessage.h
index bb54e16d..1f845ec5 100644
--- a/src/protocol/ProtocolMessage.h
+++ b/src/protocol/ProtocolMessage.h
@@ -114,7 +114,7 @@ public:
 		this->size_limit = message.size_limit;
 		this->attachment = message.attachment;
 		message.attachment = NULL;
-		this->wrapper = message.wrapper;
+		this->wrapper = NULL;
 	}
 
 	ProtocolMessage& operator = (ProtocolMessage&& message)
@@ -125,7 +125,6 @@ public:
 			delete this->attachment;
 			this->attachment = message.attachment;
 			message.attachment = NULL;
-			this->wrapper = message.wrapper;
 		}
 
 		return *this;
@@ -153,22 +152,28 @@ protected:
 		return this->message->inner();
 	}
 
+protected:
+	void set_message(ProtocolMessage *message)
+	{
+		this->message = message;
+		if (message)
+			message->wrapper = this;
+	}
+
 protected:
 	ProtocolMessage *message;
 
 public:
 	ProtocolWrapper(ProtocolMessage *message)
 	{
-		message->wrapper = this;
-		this->message = message;
+		this->set_message(message);
 	}
 
 public:
 	ProtocolWrapper(ProtocolWrapper&& wrapper) :
 		ProtocolMessage(std::move(wrapper))
 	{
-		wrapper.message->wrapper = this;
-		this->message = wrapper.message;
+		this->set_message(wrapper.message);
 		wrapper.message = NULL;
 	}
 
@@ -177,8 +182,7 @@ public:
 		if (&wrapper != this)
 		{
 			*(ProtocolMessage *)this = std::move(wrapper);
-			wrapper.message->wrapper = this;
-			this->message = wrapper.message;
+			this->set_message(wrapper.message);
 			wrapper.message = NULL;
 		}