<commit_before>// Package udp provides a connection-oriented listener over a UDP PacketConn
package udp

import (
	"context"
	"errors"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pion/dtls/v2/internal/net/deadline"
)

const receiveMTU = 8192

var errClosedListener = errors.New("udp: listener closed")

// Listener augments a connection-oriented Listener over a UDP PacketConn
type Listener struct {
	pConn *net.UDPConn

	accepting atomic.Value // bool
	acceptCh  chan *Conn
	doneCh    chan struct{}
	doneOnce  sync.Once

	connLock sync.Mutex
	conns    map[string]*Conn
	connWG   sync.WaitGroup

	readWG   sync.WaitGroup
	errClose atomic.Value // error
}

// Accept waits for and returns the next connection to the listener.
func (l *Listener) Accept() (*Conn, error) {
	select {
	case c := <-l.acceptCh:
		l.connWG.Add(1)
		return c, nil

	case <-l.doneCh:
		return nil, errClosedListener
	}
}

// Close closes the listener.
// Any blocked Accept operations will be unblocked and return errors.
func (l *Listener) Close() error {
	var err error
	l.doneOnce.Do(func() {
		l.connWG.Done()
		l.accepting.Store(false)
		close(l.doneCh)

		l.connLock.Lock()
		nConns := len(l.conns)
		l.connLock.Unlock()

		if nConns == 0 {
			// Wait if this is the final connection
			l.readWG.Wait()
			if errClose, ok := l.errClose.Load().(error); ok {
				err = errClose
			}
		} else {
			err = nil
		}
	})

	return err
}

// Addr returns the listener's network address.
func (l *Listener) Addr() net.Addr {
	return l.pConn.LocalAddr()
}

// Listen creates a new listener
func Listen(network string, laddr *net.UDPAddr) (*Listener, error) {
	conn, err := net.ListenUDP(network, laddr)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		pConn:    conn,
		acceptCh: make(chan *Conn),
		conns:    make(map[string]*Conn),
		doneCh:   make(chan struct{}),
	}
	l.accepting.Store(true)
	l.connWG.Add(1)
	l.readWG.Add(2) // wait readLoop and Close execution routine

	go l.readLoop()
	go func() {
		l.connWG.Wait()
		if err := l.pConn.Close(); err != nil {
			l.errClose.Store(err)
		}
		l.readWG.Done()
	}()

	return l, nil
}

// readLoop has to tasks:
// 1. Dispatching incoming packets to the correct Conn.
//    It can therefore not be ended until all Conns are closed.
// 2. Creating a new Conn when receiving from a new remote.
func (l *Listener) readLoop() {
	defer l.readWG.Done()
	buf := make([]byte, receiveMTU)

	for {
		n, raddr, err := l.pConn.ReadFrom(buf)
		if err != nil {
			return
		}
		conn, err := l.getConn(raddr)
		if err != nil {
			continue
		}
		cBuf := <-conn.readCh
		n = copy(cBuf, buf[:n])
		conn.sizeCh <- n
	}
}

func (l *Listener) getConn(raddr net.Addr) (*Conn, error) {
	l.connLock.Lock()
	defer l.connLock.Unlock()
	conn, ok := l.conns[raddr.String()]
	if !ok {
		if !l.accepting.Load().(bool) {
			return nil, errClosedListener
		}
		conn = l.newConn(raddr)
		l.conns[raddr.String()] = conn
		l.acceptCh <- conn<commit_msg>Fix udp.Conn deadlock on Close()<commit_after>// Package udp provides a connection-oriented listener over a UDP PacketConn
package udp

import (
	"context"
	"errors"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pion/dtls/v2/internal/net/deadline"
)

const receiveMTU = 8192

var errClosedListener = errors.New("udp: listener closed")

// Listener augments a connection-oriented Listener over a UDP PacketConn
type Listener struct {
	pConn *net.UDPConn

	accepting atomic.Value // bool
	acceptCh  chan *Conn
	doneCh    chan struct{}
	doneOnce  sync.Once

	connLock sync.Mutex
	conns    map[string]*Conn
	connWG   sync.WaitGroup

	readWG   sync.WaitGroup
	errClose atomic.Value // error
}

// Accept waits for and returns the next connection to the listener.
func (l *Listener) Accept() (*Conn, error) {
	select {
	case c := <-l.acceptCh:
		l.connWG.Add(1)
		return c, nil

	case <-l.doneCh:
		return nil, errClosedListener
	}
}

// Close closes the listener.
// Any blocked Accept operations will be unblocked and return errors.
func (l *Listener) Close() error {
	var err error
	l.doneOnce.Do(func() {
		l.connWG.Done()
		l.accepting.Store(false)
		close(l.doneCh)

		l.connLock.Lock()
		nConns := len(l.conns)
		l.connLock.Unlock()

		if nConns == 0 {
			// Wait if this is the final connection
			l.readWG.Wait()
			if errClose, ok := l.errClose.Load().(error); ok {
				err = errClose
			}
		} else {
			err = nil
		}
	})

	return err
}

// Addr returns the listener's network address.
func (l *Listener) Addr() net.Addr {
	return l.pConn.LocalAddr()
}

// Listen creates a new listener
func Listen(network string, laddr *net.UDPAddr) (*Listener, error) {
	conn, err := net.ListenUDP(network, laddr)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		pConn:    conn,
		acceptCh: make(chan *Conn),
		conns:    make(map[string]*Conn),
		doneCh:   make(chan struct{}),
	}
	l.accepting.Store(true)
	l.connWG.Add(1)
	l.readWG.Add(2) // wait readLoop and Close execution routine

	go l.readLoop()
	go func() {
		l.connWG.Wait()
		if err := l.pConn.Close(); err != nil {
			l.errClose.Store(err)
		}
		l.readWG.Done()
	}()

	return l, nil
}

// readLoop has to tasks:
// 1. Dispatching incoming packets to the correct Conn.
//    It can therefore not be ended until all Conns are closed.
// 2. Creating a new Conn when receiving from a new remote.
func (l *Listener) readLoop() {
	defer l.readWG.Done()
	buf := make([]byte, receiveMTU)

	for {
		n, raddr, err := l.pConn.ReadFrom(buf)
		if err != nil {
			return
		}
		conn, err := l.getConn(raddr)
		if err != nil {
			continue
		}
		select {
		case <-l.doneCh:
			return
		case cBuf := <-conn.readCh:
			n = copy(cBuf, buf[:n])
			conn.sizeCh <- n
		}
	}
}

func (l *Listener) getConn(raddr net.Addr) (*Conn, error) {
	l.connLock.Lock()
	defer l.connLock.Unlock()
	conn, ok := l.conns[raddr.String()]
	if !ok {
		if !l.accepting.Load().(bool) {
			return nil, errClosedListener
		}
		conn = l.newConn(raddr)
		l.conns[raddr.String()] = conn
		select {
		case <-l.doneCh:
			return nil, errClosedListener
		case l.acceptCh <- conn:
		}
{'fraction_non_alphanumeric': 0.13696369636963696, 'fraction_numerical': 0.0036303630363036304, 'mean_word_length': 3.1456908344733243, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4705209', 'n_tokens_mistral': 2456, 'n_tokens_neox': 2258, 'n_words': 639}