release reserved memory on close

libp2p · Feb 28, 2022 · 239b975 · 239b975
1 parent abd728e
commit 239b975
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/multiplex.go b/multiplex.go
@@ -92,7 +92,8 @@ type Multiplex struct {
 	channels map[streamID]*Stream
 	chLock   sync.Mutex
 
-	bufIn, bufOut chan struct{}
+	bufIn, bufOut  chan struct{}
+	reservedMemory int
 }
 
 // NewMultiplex creates a new multiplexer session.
@@ -128,6 +129,7 @@ func NewMultiplex(con net.Conn, initiator bool, memoryManager MemoryManager) (*M
 		if err = mp.memoryManager.ReserveMemory(2*MaxMessageSize, prio); err != nil {
 			break
 		}
+		mp.reservedMemory += 2 * MaxMessageSize
 		bufs++
 	}
 
@@ -186,6 +188,7 @@ func (mp *Multiplex) closeNoWait() {
 	select {
 	case <-mp.shutdown:
 	default:
+		mp.memoryManager.ReleaseMemory(mp.reservedMemory)
 		mp.con.Close()
 		close(mp.shutdown)
 	}