• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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