From 7e5750c9d64cb1ea2d61cdb1cd559802f435cfd6 Mon Sep 17 00:00:00 2001 From: Oxan van Leeuwen Date: Thu, 2 Feb 2023 23:13:25 +0100 Subject: [PATCH] Keep track of transmitted data per client This prevents dropped data if a write to a client is cut short. Move default port number to config schema --- components/stream_server/__init__.py | 13 +++++- components/stream_server/stream_server.cpp | 52 +++++++++++++++++----- components/stream_server/stream_server.h | 15 ++++++- 3 files changed, 68 insertions(+), 12 deletions(-) diff --git a/components/stream_server/__init__.py b/components/stream_server/__init__.py index b54d653..64f2f58 100644 --- a/components/stream_server/__init__.py +++ b/components/stream_server/__init__.py @@ -16,7 +16,7 @@ import esphome.codegen as cg import esphome.config_validation as cv from esphome.components import uart -from esphome.const import CONF_ID, CONF_PORT +from esphome.const import CONF_ID, CONF_PORT, CONF_BUFFER_SIZE # ESPHome doesn't know the Stream abstraction yet, so hardcode to use a UART for now. @@ -28,11 +28,21 @@ MULTI_CONF = True StreamServerComponent = cg.global_ns.class_("StreamServerComponent", cg.Component) + +def validate_buffer_size(buffer_size): + if buffer_size & (buffer_size - 1) != 0: + raise cv.Invalid("Buffer size must be a power of two.") + return buffer_size + + CONFIG_SCHEMA = ( cv.Schema( { cv.GenerateID(): cv.declare_id(StreamServerComponent), cv.Optional(CONF_PORT): cv.port, + cv.Optional(CONF_BUFFER_SIZE, default=128): cv.All( + cv.positive_int, validate_buffer_size + ), } ) .extend(cv.COMPONENT_SCHEMA) @@ -44,6 +54,7 @@ async def to_code(config): var = cg.new_Pvariable(config[CONF_ID]) if CONF_PORT in config: cg.add(var.set_port(config[CONF_PORT])) + cg.add(var.set_buffer_size(config[CONF_BUFFER_SIZE])) await cg.register_component(var, config) await uart.register_uart_device(var, config) diff --git a/components/stream_server/stream_server.cpp b/components/stream_server/stream_server.cpp index 8652e89..d54a5de 100644 --- a/components/stream_server/stream_server.cpp +++ b/components/stream_server/stream_server.cpp @@ -30,6 +30,9 @@ using namespace esphome; void StreamServerComponent::setup() { ESP_LOGCONFIG(TAG, "Setting up stream server..."); + // The make_unique() wrapper doesn't like arrays, so initialize the unique_ptr directly. + this->buf_ = std::unique_ptr{new uint8_t[this->buf_size_]}; + struct sockaddr_storage bind_addr; socklen_t bind_addrlen = socket::set_sockaddr_any(reinterpret_cast(&bind_addr), sizeof(bind_addr), htons(this->port_)); @@ -42,6 +45,7 @@ void StreamServerComponent::setup() { void StreamServerComponent::loop() { this->accept(); this->read(); + this->flush(); this->write(); this->cleanup(); } @@ -65,7 +69,7 @@ void StreamServerComponent::accept() { socket->setblocking(false); std::string identifier = socket->getpeername(); - this->clients_.emplace_back(std::move(socket), identifier); + this->clients_.emplace_back(std::move(socket), identifier, this->buf_head_); ESP_LOGD(TAG, "New client connected from %s", identifier.c_str()); } @@ -76,13 +80,41 @@ void StreamServerComponent::cleanup() { } void StreamServerComponent::read() { - int len; - while ((len = this->stream_->available()) > 0) { - char buf[128]; - len = std::min(len, 128); - this->stream_->read_array(reinterpret_cast(buf), len); - for (const Client &client : this->clients_) - client.socket->write(buf, len); + bool first_iteration = true; + int available; + while ((available = this->stream_->available()) > 0) { + // Write until the tail is encountered, or wraparound of the ring buffer if that happens before. + size_t max = std::min(this->buf_ahead(this->buf_head_), this->buf_tail_ + this->buf_size_ - this->buf_head_); + if (max == 0) { + // Only warn on the first iteration, the finite buffer size is also used as a throttling mechanism to avoid + // blocking here for too long when a large amount of data comes in. + if (first_iteration) + ESP_LOGW(TAG, "Incoming bytes available in stream, but outgoing buffer is full!"); + break; + } + + size_t len = std::min(available, max); + this->stream_->read_array(&this->buf_[this->buf_index(this->buf_head_)], len); + this->buf_head_ += len; + first_iteration = false; + } +} + +void StreamServerComponent::flush() { + this->buf_tail_ = this->buf_head_; + for (Client &client : this->clients_) { + if (client.position == this->buf_head_) + continue; + + // Split the write into two parts: from the current position to the end of the ring buffer, and from the start + // of the ring buffer until the head. The second part might be zero if no wraparound is necessary. + struct iovec iov[2]; + iov[0].iov_base = &this->buf_[this->buf_index(client.position)]; + iov[0].iov_len = std::min(this->buf_head_ - client.position, this->buf_ahead(client.position)); + iov[1].iov_base = &this->buf_[0]; + iov[1].iov_len = this->buf_head_ - (client.position + iov[0].iov_len); + client.position += client.socket->writev(iov, 2); + this->buf_tail_ = std::min(this->buf_tail_, client.position); } } @@ -101,5 +133,5 @@ void StreamServerComponent::write() { } } -StreamServerComponent::Client::Client(std::unique_ptr socket, std::string identifier) - : socket(std::move(socket)), identifier{identifier} {} +StreamServerComponent::Client::Client(std::unique_ptr socket, std::string identifier, size_t position) + : socket(std::move(socket)), identifier{identifier}, position{position} {} diff --git a/components/stream_server/stream_server.h b/components/stream_server/stream_server.h index 227ea22..49dd93a 100644 --- a/components/stream_server/stream_server.h +++ b/components/stream_server/stream_server.h @@ -29,6 +29,7 @@ public: StreamServerComponent() = default; explicit StreamServerComponent(esphome::uart::UARTComponent *stream) : stream_{stream} {} void set_uart_parent(esphome::uart::UARTComponent *parent) { this->stream_ = parent; } + void set_buffer_size(size_t size) { this->buf_size_ = size; } void setup() override; void loop() override; @@ -43,17 +44,29 @@ protected: void accept(); void cleanup(); void read(); + void flush(); void write(); + size_t buf_index(size_t pos) { return pos & (this->buf_size_ - 1); } + /// Return the number of consecutive elements that are ahead of @p pos in memory. + size_t buf_ahead(size_t pos) { return (pos | (this->buf_size_ - 1)) - pos + 1; } + struct Client { - Client(std::unique_ptr socket, std::string identifier); + Client(std::unique_ptr socket, std::string identifier, size_t position); std::unique_ptr socket{nullptr}; std::string identifier{}; bool disconnected{false}; + size_t position{0}; }; esphome::uart::UARTComponent *stream_{nullptr}; + size_t buf_size_; + + std::unique_ptr buf_{}; + size_t buf_head_{0}; + size_t buf_tail_{0}; + std::unique_ptr socket_{}; uint16_t port_{6638}; std::vector clients_{};