From af0eded2310ebb3009ab314535cf523069ac5890 Mon Sep 17 00:00:00 2001
From: petricadaipegsp <155911522+petricadaipegsp@users.noreply.github.com>
Date: Fri, 22 Nov 2024 01:01:16 +0100
Subject: [PATCH] Release buffers in edge cases (#374)

---
 go-libp2p-blossomsub/comm.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/go-libp2p-blossomsub/comm.go b/go-libp2p-blossomsub/comm.go
index b5d8c81..9bad2a5 100644
--- a/go-libp2p-blossomsub/comm.go
+++ b/go-libp2p-blossomsub/comm.go
@@ -77,6 +77,7 @@ func (p *PubSub) handleNewStream(s network.Stream) {
 			return
 		}
 		if len(msgbytes) == 0 {
+			r.ReleaseMsg(msgbytes)
 			continue
 		}
 
@@ -186,18 +187,18 @@ func (p *PubSub) handleSendingMessages(ctx context.Context, s network.Stream, ou
 				s.Reset()
 				log.Debugf("writing message to %s: %s", s.Conn().RemotePeer(), err)
 				s.Close()
+				pool.Put(buf)
 				return
 			}
 
 			_, err = s.Write(buf)
+			pool.Put(buf)
 			if err != nil {
 				s.Reset()
 				log.Debugf("writing message to %s: %s", s.Conn().RemotePeer(), err)
 				s.Close()
 				return
 			}
-
-			pool.Put(buf)
 		case <-ctx.Done():
 			s.Close()
 			return