410 lines
12 KiB
Cython
410 lines
12 KiB
Cython
@cython.no_gc_clear
|
|
@cython.freelist(DEFAULT_FREELIST_SIZE)
|
|
cdef class _UDPSendContext:
|
|
# used to hold additional write request information for uv_write
|
|
|
|
cdef:
|
|
uv.uv_udp_send_t req
|
|
|
|
uv.uv_buf_t uv_buf
|
|
Py_buffer py_buf
|
|
|
|
UDPTransport udp
|
|
|
|
bint closed
|
|
|
|
cdef close(self):
|
|
if self.closed:
|
|
return
|
|
|
|
self.closed = 1
|
|
PyBuffer_Release(&self.py_buf) # void
|
|
self.req.data = NULL
|
|
self.uv_buf.base = NULL
|
|
Py_DECREF(self)
|
|
self.udp = None
|
|
|
|
@staticmethod
|
|
cdef _UDPSendContext new(UDPTransport udp, object data):
|
|
cdef _UDPSendContext ctx
|
|
ctx = _UDPSendContext.__new__(_UDPSendContext)
|
|
ctx.udp = None
|
|
ctx.closed = 1
|
|
|
|
ctx.req.data = <void*> ctx
|
|
Py_INCREF(ctx)
|
|
|
|
PyObject_GetBuffer(data, &ctx.py_buf, PyBUF_SIMPLE)
|
|
ctx.uv_buf.base = <char*>ctx.py_buf.buf
|
|
ctx.uv_buf.len = ctx.py_buf.len
|
|
ctx.udp = udp
|
|
|
|
ctx.closed = 0
|
|
return ctx
|
|
|
|
def __dealloc__(self):
|
|
if UVLOOP_DEBUG:
|
|
if not self.closed:
|
|
raise RuntimeError(
|
|
'open _UDPSendContext is being deallocated')
|
|
self.udp = None
|
|
|
|
|
|
@cython.no_gc_clear
|
|
cdef class UDPTransport(UVBaseTransport):
|
|
def __cinit__(self):
|
|
self._family = uv.AF_UNSPEC
|
|
self.__receiving = 0
|
|
self._address = None
|
|
self.context = Context_CopyCurrent()
|
|
|
|
cdef _init(self, Loop loop, unsigned int family):
|
|
cdef int err
|
|
|
|
self._start_init(loop)
|
|
|
|
self._handle = <uv.uv_handle_t*>PyMem_RawMalloc(sizeof(uv.uv_udp_t))
|
|
if self._handle is NULL:
|
|
self._abort_init()
|
|
raise MemoryError()
|
|
|
|
err = uv.uv_udp_init_ex(loop.uvloop,
|
|
<uv.uv_udp_t*>self._handle,
|
|
family)
|
|
if err < 0:
|
|
self._abort_init()
|
|
raise convert_error(err)
|
|
|
|
if family in (uv.AF_INET, uv.AF_INET6):
|
|
self._family = family
|
|
|
|
self._finish_init()
|
|
|
|
cdef _set_address(self, system.addrinfo *addr):
|
|
self._address = __convert_sockaddr_to_pyaddr(addr.ai_addr)
|
|
|
|
cdef _connect(self, system.sockaddr* addr, size_t addr_len):
|
|
cdef int err
|
|
err = uv.uv_udp_connect(<uv.uv_udp_t*>self._handle, addr)
|
|
if err < 0:
|
|
exc = convert_error(err)
|
|
raise exc
|
|
|
|
cdef open(self, int family, int sockfd):
|
|
if family in (uv.AF_INET, uv.AF_INET6, uv.AF_UNIX):
|
|
self._family = family
|
|
else:
|
|
raise ValueError(
|
|
'cannot open a UDP handle, invalid family {}'.format(family))
|
|
|
|
cdef int err
|
|
err = uv.uv_udp_open(<uv.uv_udp_t*>self._handle,
|
|
<uv.uv_os_sock_t>sockfd)
|
|
|
|
if err < 0:
|
|
exc = convert_error(err)
|
|
raise exc
|
|
|
|
cdef _bind(self, system.sockaddr* addr):
|
|
cdef:
|
|
int err
|
|
int flags = 0
|
|
|
|
self._ensure_alive()
|
|
|
|
err = uv.uv_udp_bind(<uv.uv_udp_t*>self._handle, addr, flags)
|
|
if err < 0:
|
|
exc = convert_error(err)
|
|
raise exc
|
|
|
|
cdef _set_broadcast(self, bint on):
|
|
cdef int err
|
|
|
|
self._ensure_alive()
|
|
|
|
err = uv.uv_udp_set_broadcast(<uv.uv_udp_t*>self._handle, on)
|
|
if err < 0:
|
|
exc = convert_error(err)
|
|
raise exc
|
|
|
|
cdef size_t _get_write_buffer_size(self):
|
|
if self._handle is NULL:
|
|
return 0
|
|
return (<uv.uv_udp_t*>self._handle).send_queue_size
|
|
|
|
cdef bint _is_reading(self):
|
|
return self.__receiving
|
|
|
|
cdef _start_reading(self):
|
|
cdef int err
|
|
|
|
if self.__receiving:
|
|
return
|
|
|
|
self._ensure_alive()
|
|
|
|
err = uv.uv_udp_recv_start(<uv.uv_udp_t*>self._handle,
|
|
__loop_alloc_buffer,
|
|
__uv_udp_on_receive)
|
|
|
|
if err < 0:
|
|
exc = convert_error(err)
|
|
self._fatal_error(exc, True)
|
|
return
|
|
else:
|
|
# UDPTransport must live until the read callback is called
|
|
self.__receiving_started()
|
|
|
|
cdef _stop_reading(self):
|
|
cdef int err
|
|
|
|
if not self.__receiving:
|
|
return
|
|
|
|
self._ensure_alive()
|
|
|
|
err = uv.uv_udp_recv_stop(<uv.uv_udp_t*>self._handle)
|
|
if err < 0:
|
|
exc = convert_error(err)
|
|
self._fatal_error(exc, True)
|
|
return
|
|
else:
|
|
self.__receiving_stopped()
|
|
|
|
cdef inline __receiving_started(self):
|
|
if self.__receiving:
|
|
return
|
|
self.__receiving = 1
|
|
Py_INCREF(self)
|
|
|
|
cdef inline __receiving_stopped(self):
|
|
if not self.__receiving:
|
|
return
|
|
self.__receiving = 0
|
|
Py_DECREF(self)
|
|
|
|
cdef _new_socket(self):
|
|
if self._family not in (uv.AF_INET, uv.AF_INET6, uv.AF_UNIX):
|
|
raise RuntimeError(
|
|
'UDPTransport.family is undefined; '
|
|
'cannot create python socket')
|
|
|
|
fileno = self._fileno()
|
|
return PseudoSocket(self._family, uv.SOCK_DGRAM, 0, fileno)
|
|
|
|
cdef _send(self, object data, object addr):
|
|
cdef:
|
|
_UDPSendContext ctx
|
|
system.sockaddr_storage saddr_st
|
|
system.sockaddr *saddr
|
|
Py_buffer try_pybuf
|
|
uv.uv_buf_t try_uvbuf
|
|
|
|
self._ensure_alive()
|
|
|
|
if self._family not in (uv.AF_INET, uv.AF_INET6, uv.AF_UNIX):
|
|
raise RuntimeError('UDPTransport.family is undefined; cannot send')
|
|
|
|
if addr is None:
|
|
saddr = NULL
|
|
else:
|
|
try:
|
|
__convert_pyaddr_to_sockaddr(self._family, addr,
|
|
<system.sockaddr*>&saddr_st)
|
|
except (ValueError, TypeError):
|
|
raise
|
|
except Exception:
|
|
raise ValueError(
|
|
f'{addr!r}: socket family mismatch or '
|
|
f'a DNS lookup is required')
|
|
saddr = <system.sockaddr*>(&saddr_st)
|
|
|
|
if self._get_write_buffer_size() == 0:
|
|
PyObject_GetBuffer(data, &try_pybuf, PyBUF_SIMPLE)
|
|
try_uvbuf.base = <char*>try_pybuf.buf
|
|
try_uvbuf.len = try_pybuf.len
|
|
err = uv.uv_udp_try_send(<uv.uv_udp_t*>self._handle,
|
|
&try_uvbuf,
|
|
1,
|
|
saddr)
|
|
PyBuffer_Release(&try_pybuf)
|
|
else:
|
|
err = uv.UV_EAGAIN
|
|
|
|
if err == uv.UV_EAGAIN:
|
|
ctx = _UDPSendContext.new(self, data)
|
|
err = uv.uv_udp_send(&ctx.req,
|
|
<uv.uv_udp_t*>self._handle,
|
|
&ctx.uv_buf,
|
|
1,
|
|
saddr,
|
|
__uv_udp_on_send)
|
|
|
|
if err < 0:
|
|
ctx.close()
|
|
|
|
exc = convert_error(err)
|
|
self._fatal_error(exc, True)
|
|
else:
|
|
self._maybe_pause_protocol()
|
|
|
|
else:
|
|
if err < 0:
|
|
exc = convert_error(err)
|
|
self._fatal_error(exc, True)
|
|
else:
|
|
self._on_sent(None, self.context.copy())
|
|
|
|
cdef _on_receive(self, bytes data, object exc, object addr):
|
|
if exc is None:
|
|
run_in_context2(
|
|
self.context, self._protocol.datagram_received, data, addr,
|
|
)
|
|
else:
|
|
run_in_context1(self.context, self._protocol.error_received, exc)
|
|
|
|
cdef _on_sent(self, object exc, object context=None):
|
|
if exc is not None:
|
|
if isinstance(exc, OSError):
|
|
if context is None:
|
|
context = self.context
|
|
run_in_context1(context, self._protocol.error_received, exc)
|
|
else:
|
|
self._fatal_error(
|
|
exc, False, 'Fatal write error on datagram transport')
|
|
|
|
self._maybe_resume_protocol()
|
|
if not self._get_write_buffer_size():
|
|
if self._closing:
|
|
self._schedule_call_connection_lost(None)
|
|
|
|
# === Public API ===
|
|
|
|
def sendto(self, data, addr=None):
|
|
if not data:
|
|
# Replicating asyncio logic here.
|
|
return
|
|
|
|
if self._address:
|
|
if addr not in (None, self._address):
|
|
# Replicating asyncio logic here.
|
|
raise ValueError(
|
|
'Invalid address: must be None or %s' % (self._address,))
|
|
|
|
# Instead of setting addr to self._address below like what asyncio
|
|
# does, we depend on previous uv_udp_connect() to set the address
|
|
addr = None
|
|
|
|
if self._conn_lost:
|
|
# Replicating asyncio logic here.
|
|
if self._conn_lost >= LOG_THRESHOLD_FOR_CONNLOST_WRITES:
|
|
aio_logger.warning('socket.send() raised exception.')
|
|
self._conn_lost += 1
|
|
return
|
|
|
|
self._send(data, addr)
|
|
|
|
|
|
cdef void __uv_udp_on_receive(
|
|
uv.uv_udp_t* handle,
|
|
ssize_t nread,
|
|
const uv.uv_buf_t* buf,
|
|
const system.sockaddr* addr,
|
|
unsigned flags
|
|
) noexcept with gil:
|
|
|
|
if __ensure_handle_data(<uv.uv_handle_t*>handle,
|
|
"UDPTransport receive callback") == 0:
|
|
return
|
|
|
|
cdef:
|
|
UDPTransport udp = <UDPTransport>handle.data
|
|
Loop loop = udp._loop
|
|
bytes data
|
|
object pyaddr
|
|
|
|
# It's OK to free the buffer early, since nothing will
|
|
# be able to touch it until this method is done.
|
|
__loop_free_buffer(loop)
|
|
|
|
if udp._closed:
|
|
# The handle was closed, there is no reason to
|
|
# do any work now.
|
|
udp.__receiving_stopped() # Just in case.
|
|
return
|
|
|
|
if addr is NULL and nread == 0:
|
|
# From libuv docs:
|
|
# addr: struct sockaddr* containing the address
|
|
# of the sender. Can be NULL. Valid for the duration
|
|
# of the callback only.
|
|
# [...]
|
|
# The receive callback will be called with
|
|
# nread == 0 and addr == NULL when there is
|
|
# nothing to read, and with nread == 0 and
|
|
# addr != NULL when an empty UDP packet is
|
|
# received.
|
|
return
|
|
|
|
if addr is NULL:
|
|
pyaddr = None
|
|
elif addr.sa_family == uv.AF_UNSPEC:
|
|
# https://github.com/MagicStack/uvloop/issues/304
|
|
if system.PLATFORM_IS_LINUX:
|
|
pyaddr = None
|
|
else:
|
|
pyaddr = ''
|
|
else:
|
|
try:
|
|
pyaddr = __convert_sockaddr_to_pyaddr(addr)
|
|
except BaseException as exc:
|
|
udp._error(exc, False)
|
|
return
|
|
|
|
if nread < 0:
|
|
exc = convert_error(nread)
|
|
udp._on_receive(None, exc, pyaddr)
|
|
return
|
|
|
|
if nread == 0:
|
|
data = b''
|
|
else:
|
|
data = loop._recv_buffer[:nread]
|
|
|
|
try:
|
|
udp._on_receive(data, None, pyaddr)
|
|
except BaseException as exc:
|
|
udp._error(exc, False)
|
|
|
|
|
|
cdef void __uv_udp_on_send(
|
|
uv.uv_udp_send_t* req,
|
|
int status,
|
|
) noexcept with gil:
|
|
|
|
if req.data is NULL:
|
|
# Shouldn't happen as:
|
|
# - _UDPSendContext does an extra INCREF in its 'init()'
|
|
# - _UDPSendContext holds a ref to the relevant UDPTransport
|
|
aio_logger.error(
|
|
'UVStream.write callback called with NULL req.data, status=%r',
|
|
status)
|
|
return
|
|
|
|
cdef:
|
|
_UDPSendContext ctx = <_UDPSendContext> req.data
|
|
UDPTransport udp = <UDPTransport>ctx.udp
|
|
|
|
ctx.close()
|
|
|
|
if status < 0:
|
|
exc = convert_error(status)
|
|
print(exc)
|
|
else:
|
|
exc = None
|
|
|
|
try:
|
|
udp._on_sent(exc)
|
|
except BaseException as exc:
|
|
udp._error(exc, False)
|