From 66051b6877eadc8ddeaf68a25eb8a68800e65553 Mon Sep 17 00:00:00 2001 From: sebapeti Date: Mon, 25 Mar 2024 15:16:16 +0100 Subject: [PATCH] Improve performance of UDPMux map lookups UDPMux is using a map to lookup addresses of each packets. Unfortunately the key is based on a string and each time we want to check the map, a conversion of the UDP address to string is made (.String()) which is expensive. This CR replace the string key by a binary key called ipPort. This structure contains a netip.Addr field and ipPort could be used as a map key --- errors.go | 2 ++ udp_mux.go | 52 ++++++++++++++++++++++++++++++++++++----------- udp_muxed_conn.go | 28 ++++++++++++++++--------- 3 files changed, 60 insertions(+), 22 deletions(-) diff --git a/errors.go b/errors.go index 46785ed..e39c7cf 100644 --- a/errors.go +++ b/errors.go @@ -133,6 +133,8 @@ var ( errWriteSTUNMessage = errors.New("failed to send STUN message") errWriteSTUNMessageToIceConn = errors.New("failed to write STUN message to ICE connection") errXORMappedAddrTimeout = errors.New("timeout while waiting for XORMappedAddr") + errFailedToCastUDPAddr = errors.New("failed to cast net.Addr to net.UDPAddr") + errInvalidIPAddress = errors.New("invalid ip address") // UDPMuxDefault should not listen on unspecified address, but to keep backward compatibility, don't return error now. // will be used in the future. diff --git a/udp_mux.go b/udp_mux.go index cf01537..dc45458 100644 --- a/udp_mux.go +++ b/udp_mux.go @@ -7,6 +7,7 @@ import ( "errors" "io" "net" + "net/netip" "os" "strings" "sync" @@ -36,7 +37,7 @@ type UDPMuxDefault struct { connsIPv4, connsIPv6 map[string]*udpMuxedConn addressMapMu sync.RWMutex - addressMap map[string]*udpMuxedConn + addressMap map[ipPort]*udpMuxedConn // Buffer pool to recycle buffers for net.UDPAddr encodes/decodes pool *sync.Pool @@ -51,8 +52,9 @@ const maxAddrSize = 512 // UDPMuxParams are parameters for UDPMux. type UDPMuxParams struct { - Logger logging.LeveledLogger - UDPConn net.PacketConn + Logger logging.LeveledLogger + UDPConn net.PacketConn + UDPConnString string // Required for gathering local addresses // in case a un UDPConn is passed which does not @@ -103,9 +105,10 @@ func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault { } } } + params.UDPConnString = params.UDPConn.LocalAddr().String() m := &UDPMuxDefault{ - addressMap: map[string]*udpMuxedConn{}, + addressMap: map[ipPort]*udpMuxedConn{}, params: params, connsIPv4: make(map[string]*udpMuxedConn), connsIPv6: make(map[string]*udpMuxedConn), @@ -142,7 +145,7 @@ func (m *UDPMuxDefault) GetListenAddresses() []net.Addr { // creates the connection if an existing one can't be found func (m *UDPMuxDefault) GetConn(ufrag string, addr net.Addr) (net.PacketConn, error) { // don't check addr for mux using unspecified address - if len(m.localAddrsForUnspecified) == 0 && m.params.UDPConn.LocalAddr().String() != addr.String() { + if len(m.localAddrsForUnspecified) == 0 && m.params.UDPConnString != addr.String() { return nil, errInvalidAddress } @@ -246,7 +249,7 @@ func (m *UDPMuxDefault) writeTo(buf []byte, rAddr net.Addr) (n int, err error) { return m.params.UDPConn.WriteTo(buf, rAddr) } -func (m *UDPMuxDefault) registerConnForAddress(conn *udpMuxedConn, addr string) { +func (m *UDPMuxDefault) registerConnForAddress(conn *udpMuxedConn, addr ipPort) { if m.IsClosed() { return } @@ -260,7 +263,7 @@ func (m *UDPMuxDefault) registerConnForAddress(conn *udpMuxedConn, addr string) } m.addressMap[addr] = conn - m.params.Logger.Debugf("Registered %s for %s", addr, conn.params.Key) + m.params.Logger.Debugf("Registered %s for %s", addr.addr.String(), conn.params.Key) } func (m *UDPMuxDefault) createMuxedConn(key string) *udpMuxedConn { @@ -296,15 +299,20 @@ func (m *UDPMuxDefault) connWorker() { return } - udpAddr, ok := addr.(*net.UDPAddr) + netUDPAddr, ok := addr.(*net.UDPAddr) if !ok { logger.Errorf("Underlying PacketConn did not return a UDPAddr") return } + udpAddr, err := newIPPort(netUDPAddr.IP, uint16(netUDPAddr.Port)) + if err != nil { + logger.Errorf("Failed to create a new IP/Port host pair") + return + } // If we have already seen this address dispatch to the appropriate destination m.addressMapMu.Lock() - destinationConn := m.addressMap[addr.String()] + destinationConn := m.addressMap[udpAddr] m.addressMapMu.Unlock() // If we haven't seen this address before but is a STUN packet lookup by ufrag @@ -325,7 +333,7 @@ func (m *UDPMuxDefault) connWorker() { } ufrag := strings.Split(string(attr), ":")[0] - isIPv6 := udpAddr.IP.To4() == nil + isIPv6 := netUDPAddr.IP.To4() == nil m.mu.Lock() destinationConn, _ = m.getConn(ufrag, isIPv6) @@ -333,11 +341,11 @@ func (m *UDPMuxDefault) connWorker() { } if destinationConn == nil { - m.params.Logger.Tracef("Dropping packet from %s, addr: %s", udpAddr.String(), addr.String()) + m.params.Logger.Tracef("Dropping packet from %s, addr: %s", udpAddr.addr.String(), addr.String()) continue } - if err = destinationConn.writePacket(buf[:n], udpAddr); err != nil { + if err = destinationConn.writePacket(buf[:n], netUDPAddr); err != nil { m.params.Logger.Errorf("Failed to write packet: %v", err) } } @@ -361,3 +369,23 @@ func newBufferHolder(size int) *bufferHolder { buf: make([]byte, size), } } + +type ipPort struct { + addr netip.Addr + port uint16 +} + +// newIPPort create a custom type of address based on netip.Addr and +// port. The underlying ip address passed is converted to IPv6 format +// to simplify ip address handling +func newIPPort(ip net.IP, port uint16) (ipPort, error) { + n, ok := netip.AddrFromSlice(ip.To16()) + if !ok { + return ipPort{}, errInvalidIPAddress + } + + return ipPort{ + addr: n, + port: port, + }, nil +} diff --git a/udp_muxed_conn.go b/udp_muxed_conn.go index e69c307..fb05e23 100644 --- a/udp_muxed_conn.go +++ b/udp_muxed_conn.go @@ -26,7 +26,7 @@ type udpMuxedConnParams struct { type udpMuxedConn struct { params *udpMuxedConnParams // Remote addresses that we have sent to on this conn - addresses []string + addresses []ipPort // Channel holding incoming packets buf *packetio.Buffer @@ -81,9 +81,17 @@ func (c *udpMuxedConn) WriteTo(buf []byte, rAddr net.Addr) (n int, err error) { 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) + netUDPAddr, ok := rAddr.(*net.UDPAddr) + if !ok { + return 0, errFailedToCastUDPAddr + } + + ipAndPort, err := newIPPort(netUDPAddr.IP, uint16(netUDPAddr.Port)) + if err != nil { + return 0, err + } + if !c.containsAddress(ipAndPort) { + c.addAddress(ipAndPort) } return c.params.Mux.writeTo(buf, rAddr) @@ -127,15 +135,15 @@ func (c *udpMuxedConn) isClosed() bool { } } -func (c *udpMuxedConn) getAddresses() []string { +func (c *udpMuxedConn) getAddresses() []ipPort { c.mu.Lock() defer c.mu.Unlock() - addresses := make([]string, len(c.addresses)) + addresses := make([]ipPort, len(c.addresses)) copy(addresses, c.addresses) return addresses } -func (c *udpMuxedConn) addAddress(addr string) { +func (c *udpMuxedConn) addAddress(addr ipPort) { c.mu.Lock() c.addresses = append(c.addresses, addr) c.mu.Unlock() @@ -144,11 +152,11 @@ func (c *udpMuxedConn) addAddress(addr string) { c.params.Mux.registerConnForAddress(c, addr) } -func (c *udpMuxedConn) removeAddress(addr string) { +func (c *udpMuxedConn) removeAddress(addr ipPort) { c.mu.Lock() defer c.mu.Unlock() - newAddresses := make([]string, 0, len(c.addresses)) + newAddresses := make([]ipPort, 0, len(c.addresses)) for _, a := range c.addresses { if a != addr { newAddresses = append(newAddresses, a) @@ -158,7 +166,7 @@ func (c *udpMuxedConn) removeAddress(addr string) { c.addresses = newAddresses } -func (c *udpMuxedConn) containsAddress(addr string) bool { +func (c *udpMuxedConn) containsAddress(addr ipPort) bool { c.mu.Lock() defer c.mu.Unlock() for _, a := range c.addresses {