From 452ebe27aec89c973f2b8385729d9ce95b7c026c Mon Sep 17 00:00:00 2001 From: Jeff Epler Date: Sat, 5 Nov 2022 11:24:36 -0500 Subject: [PATCH] socketpool: make socket objects selectable .. which will lead to them being usable in async contexts, pending relevant changes in asyncio --- .../espressif/common-hal/socketpool/Socket.c | 24 +++++++++++ .../common-hal/socketpool/Socket.c | 43 +++++++++++++++++++ shared-bindings/socketpool/Socket.c | 35 +++++++++++++-- shared-bindings/socketpool/Socket.h | 3 +- 4 files changed, 100 insertions(+), 5 deletions(-) diff --git a/ports/espressif/common-hal/socketpool/Socket.c b/ports/espressif/common-hal/socketpool/Socket.c index 53db7ab393..ce177f3821 100644 --- a/ports/espressif/common-hal/socketpool/Socket.c +++ b/ports/espressif/common-hal/socketpool/Socket.c @@ -553,3 +553,27 @@ mp_uint_t common_hal_socketpool_socket_sendto(socketpool_socket_obj_t *self, void common_hal_socketpool_socket_settimeout(socketpool_socket_obj_t *self, uint32_t timeout_ms) { self->timeout_ms = timeout_ms; } + +bool common_hal_socketpool_readable(socketpool_socket_obj_t *self) { + struct timeval immediate = {0, 0}; + + fd_set fds; + FD_ZERO(&fds); + FD_SET(self->num, &fds); + int num_triggered = select(self->num + 1, &fds, NULL, &fds, &immediate); + + // including returning true in the error case + return num_triggered != 0; +} + +bool common_hal_socketpool_writable(socketpool_socket_obj_t *self) { + struct timeval immediate = {0, 0}; + + fd_set fds; + FD_ZERO(&fds); + FD_SET(self->num, &fds); + int num_triggered = select(self->num + 1, NULL, &fds, &fds, &immediate); + + // including returning true in the error case + return num_triggered != 0; +} diff --git a/ports/raspberrypi/common-hal/socketpool/Socket.c b/ports/raspberrypi/common-hal/socketpool/Socket.c index a099ef9f79..440fe56d42 100644 --- a/ports/raspberrypi/common-hal/socketpool/Socket.c +++ b/ports/raspberrypi/common-hal/socketpool/Socket.c @@ -1163,3 +1163,46 @@ mp_uint_t common_hal_socketpool_socket_sendto(socketpool_socket_obj_t *socket, void common_hal_socketpool_socket_settimeout(socketpool_socket_obj_t *self, uint32_t timeout_ms) { self->timeout = timeout_ms; } + +bool common_hal_socketpool_readable(socketpool_socket_obj_t *self) { + + MICROPY_PY_LWIP_ENTER; + + bool result = self->incoming.pbuf != NULL; + + if (self->state == STATE_PEER_CLOSED) { + result = true; + } + + if (self->type == SOCKETPOOL_SOCK_STREAM && self->pcb.tcp->state == LISTEN) { + struct tcp_pcb *volatile *incoming_connection = &lwip_socket_incoming_array(self)[self->incoming.connection.iget]; + result = (incoming_connection != NULL); + } + + MICROPY_PY_LWIP_EXIT; + + return result; +} + +bool common_hal_socketpool_writable(socketpool_socket_obj_t *self) { + bool result = false; + + MICROPY_PY_LWIP_ENTER; + + switch (self->type) { + case SOCKETPOOL_SOCK_STREAM: { + result = tcp_sndbuf(self->pcb.tcp) != 0; + break; + } + case SOCKETPOOL_SOCK_DGRAM: + #if MICROPY_PY_LWIP_SOCK_RAW + case SOCKETPOOL_SOCK_RAW: + #endif + result = true; + break; + } + + MICROPY_PY_LWIP_EXIT; + + return result; +} diff --git a/shared-bindings/socketpool/Socket.c b/shared-bindings/socketpool/Socket.c index 89b23125c5..7a60463c4b 100644 --- a/shared-bindings/socketpool/Socket.c +++ b/shared-bindings/socketpool/Socket.c @@ -30,13 +30,14 @@ #include #include -#include "shared/runtime/context_manager_helpers.h" -#include "py/objtuple.h" -#include "py/objlist.h" -#include "py/runtime.h" #include "py/mperrno.h" +#include "py/objlist.h" +#include "py/objtuple.h" +#include "py/runtime.h" +#include "py/stream.h" #include "shared/netutils/netutils.h" +#include "shared/runtime/context_manager_helpers.h" #include "shared/runtime/interrupt_char.h" //| class Socket: @@ -422,6 +423,31 @@ STATIC const mp_rom_map_elem_t socketpool_socket_locals_dict_table[] = { STATIC MP_DEFINE_CONST_DICT(socketpool_socket_locals_dict, socketpool_socket_locals_dict_table); +STATIC mp_uint_t socket_ioctl(mp_obj_t self_in, mp_uint_t request, mp_uint_t arg, int *errcode) { + socketpool_socket_obj_t *self = MP_OBJ_TO_PTR(self_in); + mp_uint_t ret; + if (request == MP_STREAM_POLL) { + mp_uint_t flags = arg; + ret = 0; + if ((flags & MP_STREAM_POLL_RD) && common_hal_socketpool_readable(self) > 0) { + ret |= MP_STREAM_POLL_RD; + } + if ((flags & MP_STREAM_POLL_WR) && common_hal_socketpool_writable(self)) { + ret |= MP_STREAM_POLL_WR; + } + } else { + *errcode = MP_EINVAL; + ret = MP_STREAM_ERROR; + } + return ret; +} + +STATIC const mp_stream_p_t socket_stream_p = { + MP_PROTO_IMPLEMENT(MP_QSTR_protocol_stream) + .ioctl = socket_ioctl, + .is_text = false, +}; + const mp_obj_type_t socketpool_socket_type = { { &mp_type_type }, .flags = MP_TYPE_FLAG_EXTENDED, @@ -429,5 +455,6 @@ const mp_obj_type_t socketpool_socket_type = { .locals_dict = (mp_obj_dict_t *)&socketpool_socket_locals_dict, MP_TYPE_EXTENDED_FIELDS( .unary_op = mp_generic_unary_op, + .protocol = &socket_stream_p, ) }; diff --git a/shared-bindings/socketpool/Socket.h b/shared-bindings/socketpool/Socket.h index c6c2a66630..cf5a97a428 100644 --- a/shared-bindings/socketpool/Socket.h +++ b/shared-bindings/socketpool/Socket.h @@ -46,6 +46,8 @@ mp_uint_t common_hal_socketpool_socket_send(socketpool_socket_obj_t *self, const mp_uint_t common_hal_socketpool_socket_sendto(socketpool_socket_obj_t *self, const char *host, size_t hostlen, uint32_t port, const uint8_t *buf, uint32_t len); void common_hal_socketpool_socket_settimeout(socketpool_socket_obj_t *self, uint32_t timeout_ms); +bool common_hal_socketpool_readable(socketpool_socket_obj_t *self); +bool common_hal_socketpool_writable(socketpool_socket_obj_t *self); // Non-allocating versions for internal use. int socketpool_socket_accept(socketpool_socket_obj_t *self, uint8_t *ip, uint32_t *port); @@ -53,5 +55,4 @@ void socketpool_socket_close(socketpool_socket_obj_t *self); int socketpool_socket_send(socketpool_socket_obj_t *self, const uint8_t *buf, uint32_t len); int socketpool_socket_recv_into(socketpool_socket_obj_t *self, const uint8_t *buf, uint32_t len); - #endif // MICROPY_INCLUDED_SHARED_BINDINGS_SOCKETPOOL_SOCKET_H