Optimize memory usage by avoiding intermediate buffer in message serialization (...

[btrtrc.git] / peer-conn-msg-writer.go

diff --git a/peer-conn-msg-writer.go b/peer-conn-msg-writer.go
index 1bacc59d188c59ceb726abf4911939092ac9f574..4f17e5edec4886560cf56e7b2974723f6cb82ee0 100644 (file)
--- a/peer-conn-msg-writer.go
+++ b/peer-conn-msg-writer.go
@@ -117,10 +117,21 @@ func (cn *peerConnMsgWriter) run(keepAliveTimeout time.Duration) {
        }
 }
 
+func (cn *peerConnMsgWriter) writeToBuffer(msg pp.Message) (err error) {
+       originalLen := cn.writeBuffer.Len()
+       defer func() {
+               if err != nil {
+                       // Since an error occurred during buffer write, revert buffer to its original state before the write.
+                       cn.writeBuffer.Truncate(originalLen)
+               }
+       }()
+       return msg.WriteTo(cn.writeBuffer)
+}
+
 func (cn *peerConnMsgWriter) write(msg pp.Message) bool {
        cn.mu.Lock()
        defer cn.mu.Unlock()
-       cn.writeBuffer.Write(msg.MustMarshalBinary())
+       cn.writeToBuffer(msg)
        cn.writeCond.Broadcast()
        return !cn.writeBufferFull()
 }