diff --git a/src/uvw/handle.hpp b/src/uvw/handle.hpp index 8b949a42..b1c223f2 100644 --- a/src/uvw/handle.hpp +++ b/src/uvw/handle.hpp @@ -243,8 +243,6 @@ public: uv_fileno(as_uv_handle(), &fd); return fd; } - - std::function allocator = details::default_allocator; }; } // namespace uvw diff --git a/src/uvw/stream.h b/src/uvw/stream.h index ec76a040..d49b9154 100644 --- a/src/uvw/stream.h +++ b/src/uvw/stream.h @@ -157,7 +157,7 @@ public: #else using base::base; #endif - + using allocator = std::function; /** * @brief Shutdowns the outgoing (write) side of a duplex stream. * @@ -459,6 +459,24 @@ public: size_t write_queue_size() const noexcept { return uv_stream_get_write_queue_size(as_uv_stream()); } + + /** + * @return Allocator for data_event buffers + */ + allocator &get_allocator() { + return readBufferAllocator; + } + + /** + * Sets allocator for data_event buffers or resets to default + * @param alloc allocator for data_event buffers + */ + void set_allocator(const allocator &alloc = details::default_allocator) { + readBufferAllocator = alloc; + } + +private: + allocator readBufferAllocator = details::default_allocator; }; } // namespace uvw diff --git a/src/uvw/udp.cpp b/src/uvw/udp.cpp index b61f1694..8c57e125 100644 --- a/src/uvw/udp.cpp +++ b/src/uvw/udp.cpp @@ -213,4 +213,10 @@ UVW_INLINE size_t udp_handle::send_queue_count() const noexcept { return uv_udp_get_send_queue_count(raw()); } +UVW_INLINE udp_handle::allocator &udp_handle::get_allocator() { + return readBufferAllocator; +} +UVW_INLINE void udp_handle::set_allocator(const udp_handle::allocator &alloc) { + readBufferAllocator = alloc; +} } // namespace uvw diff --git a/src/uvw/udp.h b/src/uvw/udp.h index c6b276d5..006930f6 100644 --- a/src/uvw/udp.h +++ b/src/uvw/udp.h @@ -87,6 +87,7 @@ public: using udp_flags = details::uvw_udp_flags; using ipv4 = uvw::ipv4; using ipv6 = uvw::ipv6; + using allocator = std::function; explicit udp_handle(loop::token token, std::shared_ptr ref, unsigned int f = {}); @@ -533,6 +534,17 @@ public: */ size_t send_queue_count() const noexcept; + /** + * @return Allocator for data_event buffers + */ + allocator &get_allocator(); + + /** + * Sets allocator for data_event buffers or resets to default + * @param alloc allocator for data_event buffers + */ + void set_allocator(const allocator &alloc = details::default_allocator); + private: enum { DEFAULT, @@ -540,6 +552,8 @@ private: } tag{DEFAULT}; unsigned int flags{}; + + allocator readBufferAllocator = details::default_allocator; }; } // namespace uvw diff --git a/src/uvw/util.h b/src/uvw/util.h index 13cfe7a3..52c746d8 100644 --- a/src/uvw/util.h +++ b/src/uvw/util.h @@ -279,7 +279,7 @@ void default_allocator(uv_buf_t *buf, std::size_t suggested, T &) { template void alloc_callback(uv_handle_t *hndl, std::size_t suggested, uv_buf_t *buf) { T &ref = *(static_cast(hndl->data)); - ref.allocator(buf, suggested, ref); + (ref.get_allocator())(buf, suggested, ref); } sockaddr ip_addr(const char *addr, unsigned int port); diff --git a/test/uvw/tcp.cpp b/test/uvw/tcp.cpp index 20178e0e..9189b303 100644 --- a/test/uvw/tcp.cpp +++ b/test/uvw/tcp.cpp @@ -175,9 +175,9 @@ TEST(TCP, CustomAllocator) { ASSERT_EQ(data, ptr); sock.close(); }); - socket->allocator = [ptr](uv_buf_t *b, size_t suggested, uvw::tcp_handle &h) { + socket->set_allocator([ptr](uv_buf_t *b, size_t suggested, uvw::tcp_handle &h) { *b = uv_buf_init(ptr, 5); - }; + }); ASSERT_EQ(0, handle.accept(*socket)); ASSERT_EQ(0, socket->read()); }); diff --git a/test/uvw/udp.cpp b/test/uvw/udp.cpp index e015be4a..57032d28 100644 --- a/test/uvw/udp.cpp +++ b/test/uvw/udp.cpp @@ -141,9 +141,9 @@ TEST(UDP, CustomAllocator) { ASSERT_EQ(data, ptr); handle.close(); }); - server->allocator = [ptr](uv_buf_t *b, size_t suggested, uvw::udp_handle &h) { + server->set_allocator([ptr](uv_buf_t *b, size_t suggested, uvw::udp_handle &h) { *b = uv_buf_init(ptr, 5); - }; + }); client->on([](const uvw::send_event &, uvw::udp_handle &handle) { handle.close();