Update ProtocolMessage and ProtocolWrapper moving. (#1612)

sogou · Aug 20, 2024 · fcabd2d · fcabd2d
1 parent 0c11e8c
commit fcabd2d
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git 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
@@ -114,7 +114,7 @@ class ProtocolMessage : public CommMessageOut, public CommMessageIn
 		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 @@ class ProtocolMessage : public CommMessageOut, public CommMessageIn
 			delete this->attachment;
 			this->attachment = message.attachment;
 			message.attachment = NULL;
-			this->wrapper = message.wrapper;
 		}
 
 		return *this;
@@ -153,22 +152,28 @@ class ProtocolWrapper : public ProtocolMessage
 		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 @@ class ProtocolWrapper : public ProtocolMessage
 		if (&wrapper != this)
 		{
 			*(ProtocolMessage *)this = std::move(wrapper);
-			wrapper.message->wrapper = this;
-			this->message = wrapper.message;
+			this->set_message(wrapper.message);
 			wrapper.message = NULL;
 		}