From dd2b152d8c73f0222b8a17e20f9f94441e3e2cc3 Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Sun, 12 May 2019 17:21:56 +0800
Subject: [PATCH] add uncorkInternal for control of batchSize

---
 batchconn.go      |  4 ++++
 kcp.go            |  1 +
 readloop_linux.go |  5 -----
 sess.go           | 20 ++++++++++++++------
 4 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/batchconn.go b/batchconn.go
index 52ff705..6c30701 100644
--- a/batchconn.go
+++ b/batchconn.go
@@ -2,6 +2,10 @@ package kcp
 
 import "golang.org/x/net/ipv4"
 
+const (
+	batchSize = 16
+)
+
 type batchConn interface {
 	WriteBatch(ms []ipv4.Message, flags int) (int, error)
 	ReadBatch(ms []ipv4.Message, flags int) (int, error)
diff --git a/kcp.go b/kcp.go
index 6f51716..c9177ac 100644
--- a/kcp.go
+++ b/kcp.go
@@ -799,6 +799,7 @@ func (kcp *KCP) flush(ackOnly bool) uint32 {
 		}
 
 		if needsend {
+			current = currentMs()
 			segment.xmit++
 			segment.ts = current
 			segment.wnd = seg.wnd
diff --git a/readloop_linux.go b/readloop_linux.go
index 7b9ad2e..7d2a004 100644
--- a/readloop_linux.go
+++ b/readloop_linux.go
@@ -11,11 +11,6 @@ import (
 	"golang.org/x/net/ipv6"
 )
 
-const (
-	// ReadBatch() message size
-	batchSize = 16
-)
-
 // the read loop for a client session
 func (s *UDPSession) readLoop() {
 	addr, _ := net.ResolveUDPAddr("udp", s.conn.LocalAddr().String())
diff --git a/sess.go b/sess.go
index 25c7b9e..abcf29b 100644
--- a/sess.go
+++ b/sess.go
@@ -304,25 +304,27 @@ func (s *UDPSession) WriteBuffers(v [][]byte) (n int, err error) {
 
 // uncork sends data in txqueue if there is any
 func (s *UDPSession) uncork() {
-	var txqueue []ipv4.Message
 	s.mu.Lock()
-	txqueue = s.txqueue
-	s.txqueue = nil
+	s.uncorkInternal()
 	s.mu.Unlock()
+}
 
-	if len(txqueue) > 0 {
+// uncork sends data in txqueue if there is any
+func (s *UDPSession) uncorkInternal() {
+	if len(s.txqueue) > 0 {
 		if s.l != nil {
 			select {
-			case s.l.chTxQueue <- txqueue:
+			case s.l.chTxQueue <- s.txqueue:
 			case <-s.l.die:
 			}
 		} else {
 			select {
-			case s.chTxQueue <- txqueue:
+			case s.chTxQueue <- s.txqueue:
 			case <-s.die:
 			}
 		}
 	}
+	s.txqueue = nil
 }
 
 // Close closes the connection.
@@ -524,6 +526,9 @@ func (s *UDPSession) output(buf []byte) {
 		msg.Buffers = [][]byte{bts}
 		msg.Addr = s.remote
 		s.txqueue = append(s.txqueue, msg)
+		if len(s.txqueue) >= batchSize {
+			s.uncorkInternal()
+		}
 	}
 
 	for k := range ecc {
@@ -532,6 +537,9 @@ func (s *UDPSession) output(buf []byte) {
 		msg.Buffers = [][]byte{bts}
 		msg.Addr = s.remote
 		s.txqueue = append(s.txqueue, msg)
+		if len(s.txqueue) >= batchSize {
+			s.uncorkInternal()
+		}
 	}
 }