added shaperQueue for transmission queue management

Author: xtaci

session.go

@@ -23,7 +23,6 @@
 package smux
 
 import (
-	"container/heap"
 	"encoding/binary"
 	"errors"
 	"io"
@@ -480,22 +479,22 @@ func (s *Session) keepalive() {
 // shaperLoop implements a priority queue for write requests,
 // some control messages are prioritized over data messages
 func (s *Session) shaperLoop() {
-	var reqs shaperHeap
+	reqs := NewShaperQueue()
 	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 {
+		if !reqs.IsEmpty() {
 			chWrite = s.writes
-			next = heap.Pop(&reqs).(writeRequest)
+			next, _ = reqs.Pop()
 		} else {
 			chWrite = nil
 		}
 
 		// control heap size, chShaper is not available until packets are less than maximum allowed
-		if len(reqs) >= maxShaperSize {
+		if reqs.Len() >= maxShaperSize {
 			chShaper = nil
 		} else {
 			chShaper = s.shaper
@@ -510,10 +509,10 @@ func (s *Session) shaperLoop() {
 		case <-s.die:
 			return
 		case r := <-chShaper:
-			if chWrite != nil { // next is valid, reshape
-				heap.Push(&reqs, next)
+			if chWrite != nil { // re-enqueue the request if there is a pending write
+				reqs.Push(next)
 			}
-			heap.Push(&reqs, r)
+			reqs.Push(r)
 		case chWrite <- next:
 		}
 	}

shaper.go

@@ -22,6 +22,12 @@
 
 package smux
 
+import (
+	"container/heap"
+	"sync"
+	"time"
+)
+
 // _itimediff returns the time difference between two uint32 values.
 // The result is a signed 32-bit integer representing the difference between 'later' and 'earlier'.
 func _itimediff(later, earlier uint32) int32 {
@@ -54,3 +60,93 @@ func (h *shaperHeap) Pop() interface{} {
 	*h = old[0 : n-1]
 	return x
 }
+
+const (
+	streamExpireDuration = 1 * time.Minute
+)
+
+type shaperQueue struct {
+	streams    map[uint32]*shaperHeap
+	lastVisits map[uint32]time.Time
+	allSids    []uint32
+	nextIdx    uint32
+	count      uint32
+	mu         sync.Mutex
+}
+
+func NewShaperQueue() *shaperQueue {
+	return &shaperQueue{
+		streams:    make(map[uint32]*shaperHeap),
+		lastVisits: make(map[uint32]time.Time),
+	}
+}
+
+func (sq *shaperQueue) Push(req writeRequest) {
+	sq.mu.Lock()
+	defer sq.mu.Unlock()
+	sid := req.frame.sid
+	if _, ok := sq.streams[sid]; !ok {
+		sq.streams[sid] = new(shaperHeap)
+		sq.allSids = append(sq.allSids, sid)
+	}
+	h := sq.streams[sid]
+	heap.Push(h, req)
+	sq.lastVisits[sid] = time.Now()
+	sq.count++
+}
+
+// Pop uses Round Robin to pop writeRequests from the shaperQueue.
+func (sq *shaperQueue) Pop() (req writeRequest, ok bool) {
+	sq.mu.Lock()
+	defer sq.mu.Unlock()
+
+	if len(sq.allSids) == 0 {
+		return writeRequest{}, false
+	}
+
+	start := sq.nextIdx % uint32(len(sq.allSids))
+
+	// loop through all streams in a round-robin manner
+	for i := 0; i < len(sq.allSids); i++ {
+		idx := (int(start) + i) % len(sq.allSids)
+		sid := sq.allSids[idx]
+		h := sq.streams[sid]
+		if h == nil || h.Len() == 0 {
+			continue
+		}
+
+		// pop from the heap
+		req := heap.Pop(h).(writeRequest)
+		sq.count--
+
+		// If the heap is empty after popping, remove it from the map
+		if h.Len() == 0 && sq.lastVisits[sid].Add(streamExpireDuration).Before(time.Now()) {
+			delete(sq.streams, sid)
+			delete(sq.lastVisits, sid)
+			// copy the rest of allSids to overwrite the removed sid
+			sq.allSids = append(sq.allSids[:idx], sq.allSids[idx+1:]...)
+		}
+
+		// update nextSid for round-robin
+		if len(sq.allSids) == 0 {
+			sq.nextIdx = 0
+		} else {
+			sq.nextIdx = uint32((idx + 1) % len(sq.allSids))
+		}
+		return req, true
+	}
+
+	return writeRequest{}, false
+}
+
+func (sq *shaperQueue) IsEmpty() bool {
+	sq.mu.Lock()
+	defer sq.mu.Unlock()
+	return sq.count == 0
+}
+
+func (sq *shaperQueue) Len() int {
+	sq.mu.Lock()
+	defer sq.mu.Unlock()
+	return int(sq.count)
+}

shaper_test.go

@@ -24,7 +24,11 @@ package smux
 
 import (
 	"container/heap"
+	"fmt"
+	"math/rand"
+	"sync"
 	"testing"
+	"time"
 )
 
 func TestShaper(t *testing.T) {
@@ -70,3 +74,189 @@ func TestShaper2(t *testing.T) {
 		t.Log("sid:", w.frame.sid, "seq:", w.seq)
 	}
 }
+
+func TestShaperQueueFairness(t *testing.T) {
+	rand.Seed(time.Now().UnixNano())
+
+	sq := NewShaperQueue()
+
+	const streams = 10
+	const testDuration = 10 * time.Second
+
+	var wg sync.WaitGroup
+	sendCount := make([]uint64, streams)
+
+	stop := make(chan struct{})
+
+	// Producers: each stream pushes packets
+	for sid := 0; sid < streams; sid++ {
+		sid := sid
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			seq := uint32(0)
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				sq.Push(writeRequest{
+					frame: Frame{sid: uint32(sid)},
+					seq:   seq,
+				})
+				seq++
+				time.Sleep(time.Duration(rand.Intn(300)) * time.Microsecond)
+			}
+		}()
+	}
+
+	// Consumer: slow network, 1 pop every 10ms
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		ticker := time.NewTicker(10 * time.Millisecond)
+		for {
+			select {
+			case <-stop:
+				return
+			case <-ticker.C:
+				req, ok := sq.Pop()
+				if ok {
+					sendCount[req.frame.sid]++
+				}
+			}
+		}
+	}()
+
+	// ---- NEW: periodic live report ----
+	go func() {
+		ticker := time.NewTicker(500 * time.Millisecond)
+		for {
+			select {
+			case <-stop:
+				return
+			case <-ticker.C:
+				fmt.Printf("[DEBUG] Current counts: %v\n", sendCount)
+			}
+		}
+	}()
+
+	// run test
+	time.Sleep(testDuration)
+	close(stop)
+	wg.Wait()
+
+	// ---- final report ----
+	fmt.Println("=== FINAL COUNTS ===")
+	fmt.Println(sendCount)
+
+	// ---- fairness check ----
+	total := uint64(0)
+	for _, c := range sendCount {
+		total += c
+	}
+	avg := total / streams
+	tolerance := avg / 4 // 25%
+
+	for sid, c := range sendCount {
+		if c < avg-tolerance || c > avg+tolerance {
+			t.Errorf("stream %d unfair: got %d, avg %d", sid, c, avg)
+		}
+	}
+}
+
+func TestShaperQueue_FastWriteSlowRead(t *testing.T) {
+	rand.Seed(time.Now().UnixNano())
+
+	const (
+		streams      = 10
+		duration     = 10 * time.Second
+		producerWait = 1 * time.Microsecond  // super fast writing
+		consumerWait = 15 * time.Millisecond // super slow reading
+	)
+
+	sq := NewShaperQueue()
+
+	sendCount := make([]uint64, streams)
+	stop := make(chan struct{})
+	var wg sync.WaitGroup
+
+	// Producers: extremely fast writers
+	for sid := 0; sid < streams; sid++ {
+		sid := sid
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			seq := uint32(0)
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+
+				sq.Push(writeRequest{
+					frame: Frame{sid: uint32(sid)},
+					seq:   seq,
+				})
+				seq++
+				time.Sleep(producerWait)
+			}
+		}()
+	}
+
+	// Consumer: very slow reader
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+			}
+
+			req, ok := sq.Pop()
+			if ok {
+				sendCount[req.frame.sid]++
+			}
+
+			time.Sleep(consumerWait)
+		}
+	}()
+
+	// Periodic monitor
+	go func() {
+		ticker := time.NewTicker(500 * time.Millisecond)
+		for {
+			select {
+			case <-stop:
+				return
+			case <-ticker.C:
+				fmt.Printf("[DEBUG] Queue size=%d, counts=%v\n", sq.count, sendCount)
+			}
+		}
+	}()
+
+	// Run test
+	time.Sleep(duration)
+	close(stop)
+	wg.Wait()
+
+	fmt.Printf("=== FINAL ===\ncounts=%v\nqueue remaining=%d\n", sendCount, sq.count)
+
+	// Check fairness
+	total := uint64(0)
+	for _, v := range sendCount {
+		total += v
+	}
+	avg := total / streams
+	tolerance := avg / 3 // allow 33%
+
+	for sid, c := range sendCount {
+		if c < avg-tolerance || c > avg+tolerance {
+			t.Errorf("stream %d unfair: got %d, avg %d", sid, c, avg)
+		}
+	}
+}