From 70092d3a81f3b3f669d885cd8b0d21bd37671ca9 Mon Sep 17 00:00:00 2001
From: Luke Curley <kixelated@gmail.com>
Date: Fri, 1 Mar 2019 16:26:58 -0800
Subject: [PATCH] Optimize Read by buffering

Increases the Read performance to acceptable levels. Packet loss went
from roughly 8% to 0.5% with this change.
---
 agent.go     | 24 ++++++++++++++++--------
 candidate.go | 15 +++++++--------
 transport.go | 11 ++---------
 3 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/agent.go b/agent.go
index 12ffd59..f1a8723 100644
--- a/agent.go
+++ b/agent.go
@@ -13,6 +13,7 @@ import (
 	"errors"
 
 	"github.com/pions/stun"
+	"github.com/pions/transport/packetio"
 	"github.com/pions/webrtc/internal/util"
 )
 
@@ -25,6 +26,9 @@ const (
 
 	// defaultConnectionTimeout used to declare a connection dead
 	defaultConnectionTimeout = 30 * time.Second
+
+	// the number of bytes that can be buffered before we start to error
+	maxBufferSize = 1000 * 1000 // 1MB
 )
 
 // Agent represents the ICE agent
@@ -72,19 +76,13 @@ type Agent struct {
 	selectedPair *candidatePair
 	validPairs   []*candidatePair
 
-	// Channel for reading
-	rcvCh chan *bufIn
+	buffer *packetio.Buffer
 
 	// State for closing
 	done chan struct{}
 	err  atomicError
 }
 
-type bufIn struct {
-	buf  []byte
-	size chan int
-}
-
 func (a *Agent) ok() error {
 	select {
 	case <-a.done:
@@ -142,12 +140,17 @@ func NewAgent(config *AgentConfig) (*Agent, error) {
 		localPwd:    util.RandSeq(32),
 		taskChan:    make(chan task),
 		onConnected: make(chan struct{}),
-		rcvCh:       make(chan *bufIn),
+		buffer:      packetio.NewBuffer(),
 		done:        make(chan struct{}),
 		portmin:     config.PortMin,
 		portmax:     config.PortMax,
 	}
 
+	// Make sure the buffer doesn't grow indefinitely.
+	// NOTE: We actually won't get anywhere close to this limit.
+	// SRTP will constantly read from the endpoint and drop packets if it's full.
+	a.buffer.SetLimitSize(maxBufferSize)
+
 	// connectionTimeout used to declare a connection dead
 	if config.ConnectionTimeout == nil {
 		a.connectionTimeout = defaultConnectionTimeout
@@ -595,6 +598,11 @@ func (a *Agent) Close() error {
 			}
 			delete(agent.remoteCandidates, net)
 		}
+
+		err := a.buffer.Close()
+		if err != nil {
+			iceLog.Warnf("failed to close buffer: %v", err)
+		}
 	})
 	if err != nil {
 		return err
diff --git a/candidate.go b/candidate.go
index 71a9655..3a74bd4 100644
--- a/candidate.go
+++ b/candidate.go
@@ -140,7 +140,8 @@ func (c *Candidate) recvLoop() {
 		}
 
 		if stun.IsSTUN(buffer[:n]) {
-			m, err := stun.NewMessage(buffer[:n])
+			var m *stun.Message
+			m, err = stun.NewMessage(buffer[:n])
 			if err != nil {
 				iceLog.Warnf("Failed to handle decode ICE from %s to %s: %v", c.addr(), srcAddr, err)
 				continue
@@ -154,7 +155,7 @@ func (c *Candidate) recvLoop() {
 
 			continue
 		} else {
-			err := c.agent.run(func(agent *Agent) {
+			err = c.agent.run(func(agent *Agent) {
 				agent.noSTUNSeen(c, srcAddr)
 			})
 			if err != nil {
@@ -162,12 +163,10 @@ func (c *Candidate) recvLoop() {
 			}
 		}
 
-		select {
-		case bufin := <-c.agent.rcvCh:
-			copy(bufin.buf, buffer[:n]) // TODO: avoid copy in common case?
-			bufin.size <- n
-		case <-c.closeCh:
-			return
+		// NOTE This will return packetio.ErrFull if the buffer ever manages to fill up.
+		_, err = c.agent.buffer.Write(buffer[:n])
+		if err != nil {
+			iceLog.Warnf("failed to write packet")
 		}
 	}
 }
diff --git a/transport.go b/transport.go
index 5eda283..a0c529a 100644
--- a/transport.go
+++ b/transport.go
@@ -61,15 +61,7 @@ func (c *Conn) Read(p []byte) (int, error) {
 		return 0, err
 	}
 
-	resN := make(chan int)
-
-	select {
-	case c.agent.rcvCh <- &bufIn{p, resN}:
-		n := <-resN
-		return n, nil
-	case <-c.agent.done:
-		return 0, c.agent.getErr()
-	}
+	return c.agent.buffer.Read(p)
 }
 
 // Write implements the Conn Write method.
@@ -87,6 +79,7 @@ func (c *Conn) Write(p []byte) (int, error) {
 	if err != nil {
 		return 0, err
 	}
+
 	return pair.Write(p)
 }