1 /*
2 * twemproxy - A fast and lightweight proxy for memcached protocol.
3 * Copyright (C) 2011 Twitter, Inc.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18 #include <stdlib.h>
19 #include <string.h>
20
21 #include <nc_core.h>
22
23 static uint32_t nfree_mbufq; /* # free mbuf */
24 static struct mhdr free_mbufq; /* free mbuf q */
25
26 static size_t mbuf_chunk_size; /* mbuf chunk size - header + data (const) */
27 static size_t mbuf_offset; /* mbuf offset in chunk (const) */
28
29 static struct mbuf *
_mbuf_get(void)30 _mbuf_get(void)
31 {
32 struct mbuf *mbuf;
33 uint8_t *buf;
34
35 if (!STAILQ_EMPTY(&free_mbufq)) {
36 ASSERT(nfree_mbufq > 0);
37
38 mbuf = STAILQ_FIRST(&free_mbufq);
39 nfree_mbufq--;
40 STAILQ_REMOVE_HEAD(&free_mbufq, next);
41
42 ASSERT(mbuf->magic == MBUF_MAGIC);
43 goto done;
44 }
45
46 buf = nc_alloc(mbuf_chunk_size);
47 if (buf == NULL) {
48 return NULL;
49 }
50
51 /*
52 * mbuf header is at the tail end of the mbuf. This enables us to catch
53 * buffer overrun early by asserting on the magic value during get or
54 * put operations
55 *
56 * <------------- mbuf_chunk_size ------------->
57 * +-------------------------------------------+
58 * | mbuf data | mbuf header |
59 * | (mbuf_offset) | (struct mbuf) |
60 * +-------------------------------------------+
61 * ^ ^ ^ ^^
62 * | | | ||
63 * \ | | |\
64 * mbuf->start \ | | mbuf->end (one byte past valid bound)
65 * mbuf->pos \
66 * \ mbuf
67 * mbuf->last (one byte past valid byte)
68 *
69 */
70 mbuf = (struct mbuf *)(buf + mbuf_offset);
71 mbuf->magic = MBUF_MAGIC;
72
73 done:
74 STAILQ_NEXT(mbuf, next) = NULL;
75 return mbuf;
76 }
77
78 struct mbuf *
mbuf_get(void)79 mbuf_get(void)
80 {
81 struct mbuf *mbuf;
82 uint8_t *buf;
83
84 mbuf = _mbuf_get();
85 if (mbuf == NULL) {
86 return NULL;
87 }
88
89 buf = (uint8_t *)mbuf - mbuf_offset;
90 mbuf->start = buf;
91 mbuf->end = buf + mbuf_offset;
92
93 ASSERT(mbuf->end - mbuf->start == (int)mbuf_offset);
94 ASSERT(mbuf->start < mbuf->end);
95
96 mbuf->pos = mbuf->start;
97 mbuf->last = mbuf->start;
98
99 log_debug(LOG_VVERB, "get mbuf %p", mbuf);
100
101 return mbuf;
102 }
103
104 static void
mbuf_free(struct mbuf * mbuf)105 mbuf_free(struct mbuf *mbuf)
106 {
107 uint8_t *buf;
108
109 log_debug(LOG_VVERB, "put mbuf %p len %d", mbuf, (int)(mbuf->last - mbuf->pos));
110
111 ASSERT(STAILQ_NEXT(mbuf, next) == NULL);
112 ASSERT(mbuf->magic == MBUF_MAGIC);
113
114 buf = (uint8_t *)mbuf - mbuf_offset;
115 nc_free(buf);
116 }
117
118 void
mbuf_put(struct mbuf * mbuf)119 mbuf_put(struct mbuf *mbuf)
120 {
121 log_debug(LOG_VVERB, "put mbuf %p len %d", mbuf, (int)(mbuf->last - mbuf->pos));
122
123 ASSERT(STAILQ_NEXT(mbuf, next) == NULL);
124 ASSERT(mbuf->magic == MBUF_MAGIC);
125
126 nfree_mbufq++;
127 STAILQ_INSERT_HEAD(&free_mbufq, mbuf, next);
128 }
129
130 /*
131 * Rewind the mbuf by discarding any of the read or unread data that it
132 * might hold.
133 */
134 void
mbuf_rewind(struct mbuf * mbuf)135 mbuf_rewind(struct mbuf *mbuf)
136 {
137 mbuf->pos = mbuf->start;
138 mbuf->last = mbuf->start;
139 }
140
141 /*
142 * Return the length of data in mbuf. Mbuf cannot contain more than
143 * 2^32 bytes (4G).
144 */
145 uint32_t
mbuf_length(const struct mbuf * mbuf)146 mbuf_length(const struct mbuf *mbuf)
147 {
148 ASSERT(mbuf->last >= mbuf->pos);
149
150 return (uint32_t)(mbuf->last - mbuf->pos);
151 }
152
153 /*
154 * Return the remaining space size for any new data in mbuf. Mbuf cannot
155 * contain more than 2^32 bytes (4G).
156 */
157 uint32_t
mbuf_size(const struct mbuf * mbuf)158 mbuf_size(const struct mbuf *mbuf)
159 {
160 ASSERT(mbuf->end >= mbuf->last);
161
162 return (uint32_t)(mbuf->end - mbuf->last);
163 }
164
165 /*
166 * Return the maximum available space size for data in any mbuf. Mbuf cannot
167 * contain more than 2^32 bytes (4G).
168 */
169 size_t
mbuf_data_size(void)170 mbuf_data_size(void)
171 {
172 return mbuf_offset;
173 }
174
175 /*
176 * Insert mbuf at the tail of the mhdr Q
177 */
178 void
mbuf_insert(struct mhdr * mhdr,struct mbuf * mbuf)179 mbuf_insert(struct mhdr *mhdr, struct mbuf *mbuf)
180 {
181 STAILQ_INSERT_TAIL(mhdr, mbuf, next);
182 log_debug(LOG_VVERB, "insert mbuf %p len %d", mbuf,
183 (int)(mbuf->last - mbuf->pos));
184 }
185
186 /*
187 * Remove mbuf from the mhdr Q
188 */
189 void
mbuf_remove(struct mhdr * mhdr,struct mbuf * mbuf)190 mbuf_remove(struct mhdr *mhdr, struct mbuf *mbuf)
191 {
192 log_debug(LOG_VVERB, "remove mbuf %p len %d", mbuf,
193 (int)(mbuf->last - mbuf->pos));
194
195 STAILQ_REMOVE(mhdr, mbuf, mbuf, next);
196 STAILQ_NEXT(mbuf, next) = NULL;
197 }
198
199 /*
200 * Copy n bytes from memory area pos to mbuf.
201 *
202 * The memory areas should not overlap and the mbuf should have
203 * enough space for n bytes.
204 */
205 void
mbuf_copy(struct mbuf * mbuf,const uint8_t * pos,size_t n)206 mbuf_copy(struct mbuf *mbuf, const uint8_t *pos, size_t n)
207 {
208 if (n == 0) {
209 return;
210 }
211
212 /* mbuf has space for n bytes */
213 ASSERT(!mbuf_full(mbuf) && n <= mbuf_size(mbuf));
214
215 /* no overlapping copy */
216 ASSERT(pos < mbuf->start || pos >= mbuf->end);
217
218 nc_memcpy(mbuf->last, pos, n);
219 mbuf->last += n;
220 }
221
222 /*
223 * Split mbuf h into h and t by copying data from h to t. Before
224 * the copy, we invoke a precopy handler cb that will copy a predefined
225 * string to the head of t.
226 *
227 * Return new mbuf t, if the split was successful.
228 */
229 struct mbuf *
mbuf_split(struct mhdr * h,uint8_t * pos,mbuf_copy_t cb,void * cbarg)230 mbuf_split(struct mhdr *h, uint8_t *pos, mbuf_copy_t cb, void *cbarg)
231 {
232 struct mbuf *mbuf, *nbuf;
233 size_t size;
234
235 ASSERT(!STAILQ_EMPTY(h));
236
237 mbuf = STAILQ_LAST(h, mbuf, next);
238 ASSERT(pos >= mbuf->pos && pos <= mbuf->last);
239
240 nbuf = mbuf_get();
241 if (nbuf == NULL) {
242 return NULL;
243 }
244
245 if (cb != NULL) {
246 /* precopy nbuf */
247 cb(nbuf, cbarg);
248 }
249
250 /* copy data from mbuf to nbuf */
251 size = (size_t)(mbuf->last - pos);
252 mbuf_copy(nbuf, pos, size);
253
254 /* adjust mbuf */
255 mbuf->last = pos;
256
257 log_debug(LOG_VVERB, "split into mbuf %p len %"PRIu32" and nbuf %p len "
258 "%"PRIu32" copied %zu bytes", mbuf, mbuf_length(mbuf), nbuf,
259 mbuf_length(nbuf), size);
260
261 return nbuf;
262 }
263
264 void
mbuf_init(const struct instance * nci)265 mbuf_init(const struct instance *nci)
266 {
267 nfree_mbufq = 0;
268 STAILQ_INIT(&free_mbufq);
269
270 mbuf_chunk_size = nci->mbuf_chunk_size;
271 mbuf_offset = mbuf_chunk_size - MBUF_HSIZE;
272
273 log_debug(LOG_DEBUG, "mbuf hsize %d chunk size %zu offset %zu length %zu",
274 (int)MBUF_HSIZE, mbuf_chunk_size, mbuf_offset, mbuf_offset);
275 }
276
277 void
mbuf_deinit(void)278 mbuf_deinit(void)
279 {
280 while (!STAILQ_EMPTY(&free_mbufq)) {
281 struct mbuf *mbuf = STAILQ_FIRST(&free_mbufq);
282 mbuf_remove(&free_mbufq, mbuf);
283 mbuf_free(mbuf);
284 nfree_mbufq--;
285 }
286 ASSERT(nfree_mbufq == 0);
287 }
288