From 81df2852562604f2c6cdde2dc00f146649745b26 Mon Sep 17 00:00:00 2001
From: Ashish Kumar Singh <ashish@dyte.in>
Date: Wed, 13 Dec 2023 10:09:04 +0000
Subject: [PATCH] feat: only check address at ice connecting stage

---
 udp_muxed_conn.go | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/udp_muxed_conn.go b/udp_muxed_conn.go
index e69c307..843e779 100644
--- a/udp_muxed_conn.go
+++ b/udp_muxed_conn.go
@@ -8,12 +8,17 @@ import (
 	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pion/logging"
 	"github.com/pion/transport/v3/packetio"
 )
 
+const (
+	iceConnectedTimeout = 25 * time.Second
+)
+
 type udpMuxedConnParams struct {
 	Mux       *UDPMuxDefault
 	AddrPool  *sync.Pool
@@ -33,6 +38,9 @@ type udpMuxedConn struct {
 	closedChan chan struct{}
 	closeOnce  sync.Once
 	mu         sync.Mutex
+
+	startAt      time.Time
+	iceConnected atomic.Bool
 }
 
 func newUDPMuxedConn(params *udpMuxedConnParams) *udpMuxedConn {
@@ -40,6 +48,7 @@ func newUDPMuxedConn(params *udpMuxedConnParams) *udpMuxedConn {
 		params:     params,
 		buf:        packetio.NewBuffer(),
 		closedChan: make(chan struct{}),
+		startAt:    time.Now(),
 	}
 
 	return p
@@ -80,10 +89,18 @@ func (c *udpMuxedConn) WriteTo(buf []byte, rAddr net.Addr) (n int, err error) {
 	if c.isClosed() {
 		return 0, io.ErrClosedPipe
 	}
-	// Each time we write to a new address, we'll register it with the mux
-	addr := rAddr.String()
-	if !c.containsAddress(addr) {
-		c.addAddress(addr)
+
+	// Only check the address at the ICE connecting stage to reduce the check cost
+	if !c.iceConnected.Load() {
+		if time.Since(c.startAt) > iceConnectedTimeout {
+			c.iceConnected.Store(true)
+		} else {
+			// Each time we write to a new address, we'll register it with the mux
+			addr := rAddr.String()
+			if !c.containsAddress(addr) {
+				c.addAddress(addr)
+			}
+		}
 	}
 
 	return c.params.Mux.writeTo(buf, rAddr)