From a54d6cb36108ad695292f109e21adfc6f82466f5 Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Sun, 12 Feb 2023 19:51:18 +0800
Subject: [PATCH] control of shaper heap size

---
 session.go | 17 ++++++++++++++++-
 stream.go  |  4 ++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/session.go b/session.go
index 4659743..28fdd04 100644
--- a/session.go
+++ b/session.go
@@ -13,6 +13,7 @@ import (
 
 const (
 	defaultAcceptBacklog = 1024
+	maxShaperSize        = 1024
 )
 
 var (
@@ -420,8 +421,10 @@ func (s *Session) shaperLoop() {
 	var reqs shaperHeap
 	var next writeRequest
 	var chWrite chan writeRequest
+	var chShaper chan writeRequest
 
 	for {
+		// chWrite is not available until it has packet to send
 		if len(reqs) > 0 {
 			chWrite = s.writes
 			next = heap.Pop(&reqs).(writeRequest)
@@ -429,10 +432,22 @@ func (s *Session) shaperLoop() {
 			chWrite = nil
 		}
 
+		// control heap size, chShaper is not available until packets are less than maximum allowed
+		if len(reqs) >= maxShaperSize {
+			chShaper = nil
+		} else {
+			chShaper = s.shaper
+		}
+
+		// assertion on non nil
+		if chShaper == nil && chWrite == nil {
+			panic("both channel are nil")
+		}
+
 		select {
 		case <-s.die:
 			return
-		case r := <-s.shaper:
+		case r := <-chShaper:
 			if chWrite != nil { // next is valid, reshape
 				heap.Push(&reqs, next)
 			}
diff --git a/stream.go b/stream.go
index a95fb56..94e858e 100644
--- a/stream.go
+++ b/stream.go
@@ -325,7 +325,7 @@ func (s *Stream) Write(b []byte) (n int, err error) {
 		}
 		frame.data = bts[:sz]
 		bts = bts[sz:]
-		n, err := s.sess.writeFrameInternal(frame, deadline, s.numWritten)
+		n, err := s.sess.writeFrameInternal(frame, deadline, 0)
 		s.numWritten++
 		sent += n
 		if err != nil {
@@ -393,7 +393,7 @@ func (s *Stream) writeV2(b []byte) (n int, err error) {
 				}
 				frame.data = bts[:sz]
 				bts = bts[sz:]
-				n, err := s.sess.writeFrameInternal(frame, deadline, atomic.LoadUint32(&s.numWritten))
+				n, err := s.sess.writeFrameInternal(frame, deadline, 0)
 				atomic.AddUint32(&s.numWritten, uint32(sz))
 				sent += n
 				if err != nil {