Chore: use slice instead of map for system stack udp receiver queue

This commit is contained in:
gVisor bot 2022-03-15 05:19:29 +08:00
parent 4784d739ec
commit e252d96095
2 changed files with 25 additions and 20 deletions

View file

@ -25,7 +25,6 @@ func Start(
tab := newTable() tab := newTable()
udp := &UDP{ udp := &UDP{
calls: map[*call]struct{}{},
device: device, device: device,
buf: [pool.UDPBufferSize]byte{}, buf: [pool.UDPBufferSize]byte{},
} }

View file

@ -21,27 +21,27 @@ type call struct {
type UDP struct { type UDP struct {
closed bool closed bool
lock sync.Mutex
calls map[*call]struct{}
device io.Writer device io.Writer
queueLock sync.Mutex
queue []*call
bufLock sync.Mutex bufLock sync.Mutex
buf [pool.UDPBufferSize]byte buf [pool.UDPBufferSize]byte
} }
func (u *UDP) ReadFrom(buf []byte) (int, net.Addr, net.Addr, error) { func (u *UDP) ReadFrom(buf []byte) (int, net.Addr, net.Addr, error) {
u.lock.Lock() u.queueLock.Lock()
defer u.lock.Unlock() defer u.queueLock.Unlock()
for !u.closed { for !u.closed {
c := &call{ c := &call{
cond: sync.NewCond(&u.lock), cond: sync.NewCond(&u.queueLock),
buf: buf, buf: buf,
n: -1, n: -1,
source: nil, source: nil,
destination: nil, destination: nil,
} }
u.calls[c] = struct{}{} u.queue = append(u.queue, c)
c.cond.Wait() c.cond.Wait()
@ -54,6 +54,10 @@ func (u *UDP) ReadFrom(buf []byte) (int, net.Addr, net.Addr, error) {
} }
func (u *UDP) WriteTo(buf []byte, local net.Addr, remote net.Addr) (int, error) { func (u *UDP) WriteTo(buf []byte, local net.Addr, remote net.Addr) (int, error) {
if u.closed {
return 0, net.ErrClosed
}
u.bufLock.Lock() u.bufLock.Lock()
defer u.bufLock.Unlock() defer u.bufLock.Unlock()
@ -77,8 +81,9 @@ func (u *UDP) WriteTo(buf []byte, local net.Addr, remote net.Addr) (int, error)
return 0, net.InvalidAddrError("invalid ip version") return 0, net.InvalidAddrError("invalid ip version")
} }
tcpip.SetIPv4(u.buf[:])
ip := tcpip.IPv4Packet(u.buf[:]) ip := tcpip.IPv4Packet(u.buf[:])
tcpip.SetIPv4(ip)
ip.SetHeaderLen(tcpip.IPv4HeaderSize) ip.SetHeaderLen(tcpip.IPv4HeaderSize)
ip.SetTotalLength(tcpip.IPv4HeaderSize + tcpip.UDPHeaderSize + uint16(len(buf))) ip.SetTotalLength(tcpip.IPv4HeaderSize + tcpip.UDPHeaderSize + uint16(len(buf)))
ip.SetTypeOfService(0) ip.SetTypeOfService(0)
@ -102,12 +107,12 @@ func (u *UDP) WriteTo(buf []byte, local net.Addr, remote net.Addr) (int, error)
} }
func (u *UDP) Close() error { func (u *UDP) Close() error {
u.lock.Lock() u.queueLock.Lock()
defer u.lock.Unlock() defer u.queueLock.Unlock()
u.closed = true u.closed = true
for c := range u.calls { for _, c := range u.queue {
c.cond.Signal() c.cond.Signal()
} }
@ -117,14 +122,15 @@ func (u *UDP) Close() error {
func (u *UDP) handleUDPPacket(ip tcpip.IP, pkt tcpip.UDPPacket) { func (u *UDP) handleUDPPacket(ip tcpip.IP, pkt tcpip.UDPPacket) {
var c *call var c *call
u.lock.Lock() u.queueLock.Lock()
for c = range u.calls { if len(u.queue) > 0 {
delete(u.calls, c) idx := len(u.queue) - 1
break c = u.queue[idx]
u.queue = u.queue[:idx]
} }
u.lock.Unlock() u.queueLock.Unlock()
if c != nil { if c != nil {
c.source = &net.UDPAddr{ c.source = &net.UDPAddr{