quic: prioritise listen connections for reuse

p2p/transport/quicreuse/reuse.go

@@ -76,14 +76,19 @@ type reuse struct {
 	unicast map[string] /* IP.String() */ map[int] /* port */ *reuseConn
 	// global contains connections that are listening on 0.0.0.0 / ::
 	global map[int]*reuseConn
+	// globalFallback contains connections that we've dialed out from. connections
+	// are reused from this map if no connection is available in the global
+	// map. These connections are listening on 0.0.0.0 / ::
+	globalFallback map[int]*reuseConn
 }
 
 func newReuse() *reuse {
 	r := &reuse{
-		unicast:    make(map[string]map[int]*reuseConn),
-		global:     make(map[int]*reuseConn),
-		closeChan:  make(chan struct{}),
-		gcStopChan: make(chan struct{}),
+		unicast:        make(map[string]map[int]*reuseConn),
+		global:         make(map[int]*reuseConn),
+		globalFallback: make(map[int]*reuseConn),
+		closeChan:      make(chan struct{}),
+		gcStopChan:     make(chan struct{}),
 	}
 	go r.gc()
 	return r
@@ -95,6 +100,9 @@ func (r *reuse) gc() {
 		for _, conn := range r.global {
 			conn.Close()
 		}
+		for _, conn := range r.globalFallback {
+			conn.Close()
+		}
 		for _, conns := range r.unicast {
 			for _, conn := range conns {
 				conn.Close()
@@ -119,6 +127,12 @@ func (r *reuse) gc() {
 					delete(r.global, key)
 				}
 			}
+			for key, conn := range r.globalFallback {
+				if conn.ShouldGarbageCollect(now) {
+					conn.Close()
+					delete(r.globalFallback, key)
+				}
+			}
 			for ukey, conns := range r.unicast {
 				for key, conn := range conns {
 					if conn.ShouldGarbageCollect(now) {
@@ -189,6 +203,11 @@ func (r *reuse) dialLocked(network string, source *net.IP) (*reuseConn, error) {
 		return conn, nil
 	}
 
+	// Use a connection we've previously dialed from
+	for _, conn := range r.globalFallback {
+		return conn, nil
+	}
+
 	// We don't have a connection that we can use for dialing.
 	// Dial a new connection from a random port.
 	var addr *net.UDPAddr
@@ -203,29 +222,43 @@ func (r *reuse) dialLocked(network string, source *net.IP) (*reuseConn, error) {
 		return nil, err
 	}
 	rconn := newReuseConn(conn)
-	r.global[conn.LocalAddr().(*net.UDPAddr).Port] = rconn
+	r.globalFallback[conn.LocalAddr().(*net.UDPAddr).Port] = rconn
 	return rconn, nil
 }
 
 func (r *reuse) Listen(network string, laddr *net.UDPAddr) (*reuseConn, error) {
-	conn, err := net.ListenUDP(network, laddr)
-	if err != nil {
-		return nil, err
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+
+	var rconn *reuseConn
+	var localAddr *net.UDPAddr
+
+	// reuse the fallback connection if we've dialed out from this port already
+	if laddr.IP.IsUnspecified() {
+		if _, ok := r.globalFallback[laddr.Port]; ok {
+			rconn = r.globalFallback[laddr.Port]
+			localAddr = rconn.UDPConn.LocalAddr().(*net.UDPAddr)
+		}
+	}
+	if rconn == nil {
+		conn, err := net.ListenUDP(network, laddr)
+		if err != nil {
+			return nil, err
+		}
+		localAddr = conn.LocalAddr().(*net.UDPAddr)
+		rconn = newReuseConn(conn)
 	}
-	localAddr := conn.LocalAddr().(*net.UDPAddr)
 
-	rconn := newReuseConn(conn)
 	rconn.IncreaseCount()
 
-	r.mutex.Lock()
-	defer r.mutex.Unlock()
-
 	// Deal with listen on a global address
 	if localAddr.IP.IsUnspecified() {
 		// The kernel already checked that the laddr is not already listen
 		// so we need not check here (when we create ListenUDP).
 		r.global[localAddr.Port] = rconn
-		return rconn, err
+		// delete the entry from fallback map in case we are reusing this connection
+		delete(r.globalFallback, localAddr.Port)
+		return rconn, nil
 	}
 
 	// Deal with listen on a unicast address
@@ -239,7 +272,7 @@ func (r *reuse) Listen(network string, laddr *net.UDPAddr) (*reuseConn, error) {
 	// The kernel already checked that the laddr is not already listen
 	// so we need not check here (when we create ListenUDP).
 	r.unicast[localAddr.IP.String()][localAddr.Port] = rconn
-	return rconn, err
+	return rconn, nil
 }
 
 func (r *reuse) Close() error {

p2p/transport/quicreuse/reuse_test.go

@@ -26,6 +26,11 @@ func closeAllConns(reuse *reuse) {
 			conn.DecreaseCount()
 		}
 	}
+	for _, conn := range reuse.globalFallback {
+		for conn.GetCount() > 0 {
+			conn.DecreaseCount()
+		}
+	}
 	for _, conns := range reuse.unicast {
 		for _, conn := range conns {
 			for conn.GetCount() > 0 {
@@ -110,6 +115,45 @@ func TestReuseConnectionWhenDialing(t *testing.T) {
 	require.Equal(t, conn.GetCount(), 2)
 }
 
+func TestReuseConnectionWhenListening(t *testing.T) {
+	reuse := newReuse()
+	cleanup(t, reuse)
+
+	raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
+	require.NoError(t, err)
+	conn, err := reuse.Dial("udp4", raddr)
+	require.NoError(t, err)
+	laddr := &net.UDPAddr{IP: net.IPv4zero, Port: conn.UDPConn.LocalAddr().(*net.UDPAddr).Port}
+	lconn, err := reuse.Listen("udp4", laddr)
+	require.NoError(t, err)
+	require.Equal(t, lconn.GetCount(), 2)
+	require.Equal(t, conn.GetCount(), 2)
+}
+
+func TestReuseConnectionWhenDialBeforeListen(t *testing.T) {
+	reuse := newReuse()
+	cleanup(t, reuse)
+
+	// dial any address
+	raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
+	require.NoError(t, err)
+	_, err = reuse.Dial("udp4", raddr)
+	require.NoError(t, err)
+
+	// open a listener
+	laddr := &net.UDPAddr{IP: net.IPv4zero, Port: 10000}
+	lconn, err := reuse.Listen("udp4", laddr)
+	require.NoError(t, err)
+
+	// new dials should go via the listener connection
+	raddr, err = net.ResolveUDPAddr("udp4", "1.1.1.1:1235")
+	require.NoError(t, err)
+	conn, err := reuse.Dial("udp4", raddr)
+	require.NoError(t, err)
+	require.Equal(t, conn, lconn)
+	require.Equal(t, conn.GetCount(), 2)
+}
+
 func TestReuseListenOnSpecificInterface(t *testing.T) {
 	if platformHasRoutingTables() {
 		t.Skip("this test only works on platforms that support routing tables")