1 /* 2 * Copyright (C) 2025 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #pragma once 18 19 #include <memory> 20 #include <vector> 21 22 #include <liburing.h> 23 24 /* 25 * IOUringSocketHandler is a helper class for using io_uring with a socket. 26 * 27 * Typical usage from a given thread: 28 * 29 * As a one time setup: 30 * 1. Create an instance of IOUringSocketHandler with the socket file descriptor. 31 * 2. Setup io_uring ring buffer. 32 * 3. Allocate buffers for the ring buffer. 33 * 4. Register buffers with io_uring. 34 * 5. EnqueueMultishotRecvmsg() will submit the SQE to receive the data 35 * 36 * In the I/O path: 37 * 38 * 6. Receive data from the socket through ReceiveData() 39 * 7. Release the buffer to io_uring. 40 * 41 * Note that the thread which sets up the io_uring instance should handle the 42 * I/O through ReceiveData() call. 43 */ 44 45 class IOUringSocketHandler { 46 public: 47 IOUringSocketHandler(int socket_fd); 48 ~IOUringSocketHandler(); 49 50 // Setup io_uring ring buffer 51 // queue_size: The size of the io_uring submission queue. 52 // Determines the maximum number of outstanding I/O requests. 53 // return: true on success, false on failure (e.g., if io_uring_setup fails). 54 // 55 // This function initializes the io_uring context and sets up the submission 56 // and completion queues. It prepares the io_uring instance for I/O operations. 57 // A larger queue_size allows for more concurrent I/O operations but consumes 58 // more memory. 59 bool SetupIoUring(int queue_size); 60 61 // Allocate 'num_buffers' of size 'buf_size' 62 // 63 // num_buffers: The number of buffers to allocate - Should be power of 2. 64 // buf_size: The size of each buffer in bytes. 65 // 66 // This function allocates a set of buffers that will be used for I/O operations 67 // with io_uring. These buffers are typically used to hold data that is read from 68 // or written to files or sockets. The allocated buffers are managed internally 69 // and are later registered with io_uring. 70 // 71 // The num_buffers will be the payload for the caller. Internally, it 72 // allocates additional metadata: 73 // a: sizeof(struct ucred) + sizeof(struct cmsghdr) 74 // b: sizeof(struct io_uring_recvmsg_out) 75 // This allows sender to send the ucred credential information if required. 76 // 77 // This function also registers the allocated buffers with the io_uring instance. 78 // Registering buffers allows the kernel to access them directly, avoiding the need 79 // to copy data between user space and kernel space during I/O operations. This 80 // improves performance. 81 // 82 // Please see additional details on how num_buffers will be used 83 // by the io_uring: https://man7.org/linux/man-pages/man3/io_uring_setup_buf_ring.3.html 84 bool AllocateAndRegisterBuffers(size_t num_buffers, size_t buf_size); 85 86 // Free up registered buffers with the io_uring instance. 87 // 88 // All the buffers allocated using AllocateAndRegisterBuffers() API will be 89 // freed and de-registered. Callers can then call 90 // AllocateAndRegisterBuffers() to re-register new set of bufferes with the 91 // ring. 92 void DeRegisterBuffers(); 93 94 // ARM io_uring recvmsg opcode 95 // 96 // return: true on success, false on failure (e.g., if submission queue is full). 97 // 98 // This function enqueues a "multishot recvmsg" operation into the io_uring submission queue. 99 // Multishot recvmsg allows receiving multiple messages from a socket with a single 100 // io_uring submission. The function prepares the submission queue 101 // entry (SQE) for the recvmsg operation. 102 bool EnqueueMultishotRecvmsg(); 103 104 // Release the buffer to io_uring 105 // 106 // This function releases a buffer back to the io_uring subsystem after it has been 107 // used for an I/O operation. This makes the buffer available for reuse in subsequent 108 // I/O operations. 109 // 110 // Additionally, when the buffer is released, a check is done to see if 111 // there are more CQE entries available. If not, EnqueueMultishotRecvmsg() 112 // is invoked so that the SQE submission is done for receiving next set of 113 // I/O. 114 void ReleaseBuffer(); 115 116 // Receive payload data of size payload_len. Additionally, receive 117 // credential data. 118 // 119 // payload: A pointer to a void pointer. This will be set to point to the received 120 // payload data. 121 // 122 // payload_len: A reference to a size_t. This will be set to the length of the 123 // received payload data. 124 // 125 // cred: A pointer to a struct ucred pointer. This will be set to point to the 126 // user credentials associated with the received data (if available). 127 // If the sender doesn't have credential information in the payload, 128 // then nullptr will be returned. 129 // 130 // This function retrieves the data received from a recvmsg operation. It extracts the payload 131 // data and its length, as well as the user credentials associated with the sender. The 132 // caller is responsible for freeing the allocated memory for the payload and credentials 133 // when they are no longer needed. 134 void ReceiveData(void** payload, size_t& payload_len, struct ucred** cred); 135 136 // check if io_uring is supported 137 // 138 // return: true if io_uring is supported by the kernel, false otherwise. 139 // 140 // This function checks if the io_uring feature is supported by the underlying Linux kernel. 141 static bool isIouringEnabled(); 142 143 private: 144 static bool isIouringSupportedByKernel(); 145 // Register buffers with io_uring 146 // 147 // return: true on success, false on failure (e.g., if io_uring_register_buffers fails). 148 // 149 // This function registers the previously allocated buffers with the io_uring instance. 150 // Registering buffers allows the kernel to access them directly, avoiding the need 151 // to copy data between user space and kernel space during I/O operations. This 152 // improves performance. 153 bool RegisterBuffers(); 154 155 struct uring_context { 156 struct io_uring ring; 157 }; 158 // Socket fd 159 int socket_; 160 std::unique_ptr<uring_context> mCtx; 161 std::vector<std::unique_ptr<uint8_t[]>> buffers_; 162 struct msghdr msg; 163 int control_len_; 164 size_t num_buffers_ = 0; 165 int buffer_size_; 166 int active_buffer_id_ = -1; 167 struct io_uring_cqe* cqe; 168 // A constant buffer group id as we don't support multiple buffer groups 169 // yet. 170 const int bgid_ = 7; 171 struct io_uring_buf_ring* br_; 172 bool registered_buffers_ = false; 173 bool ring_setup_ = false; 174 }; 175