Skip to content

Commit

Permalink
CLEANUP: Refactor the method of building UDP headers
Browse files Browse the repository at this point in the history
  • Loading branch information
ing-eoking committed Mar 28, 2024
1 parent b1cae74 commit d4cd2da
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 83 deletions.
123 changes: 42 additions & 81 deletions memcached.c
Original file line number Diff line number Diff line change
Expand Up @@ -407,7 +407,7 @@ static int add_msghdr(conn *c)

if (IS_UDP(c->transport)) {
/* Leave room for the UDP header, which we'll fill in later. */
return add_iov(c, NULL, UDP_HEADER_SIZE);
return add_iov(c, c->hdrbuf, UDP_HEADER_SIZE);
}

return 0;
Expand Down Expand Up @@ -880,11 +880,6 @@ static void conn_cleanup(conn *c)
c->write_and_free = 0;
}

if (c->hdrbuf) {
free(c->hdrbuf);
c->hdrbuf = NULL;
}

if (c->sasl_conn) {
sasl_dispose(&c->sasl_conn);
c->sasl_conn = NULL;
Expand Down Expand Up @@ -1414,41 +1409,21 @@ static void einfo_append_tail_crlf(eitem_info *einfo)
}

/*
* Constructs a set of UDP headers and attaches them to the outgoing messages.
* Constructs an UDP header and attaches it to the outgoing message.
*/
static int build_udp_headers(conn *c)
static void build_udp_header(conn *c, struct iovec *header)
{
assert(c != NULL);
unsigned char *hdr;

if (c->msgused > c->hdrsize) {
void *new_hdrbuf;
if (c->hdrbuf)
new_hdrbuf = realloc(c->hdrbuf, c->msgused * 2 * UDP_HEADER_SIZE);
else
new_hdrbuf = malloc(c->msgused * 2 * UDP_HEADER_SIZE);
if (! new_hdrbuf)
return -1;
c->hdrbuf = (unsigned char *)new_hdrbuf;
c->hdrsize = c->msgused * 2;
}

hdr = c->hdrbuf;
for (int i = 0; i < c->msgused; i++) {
c->msglist[i].msg_iov[0].iov_base = (void*)hdr;
c->msglist[i].msg_iov[0].iov_len = UDP_HEADER_SIZE;
*hdr++ = c->request_id / 256;
*hdr++ = c->request_id % 256;
*hdr++ = i / 256;
*hdr++ = i % 256;
*hdr++ = c->msgused / 256;
*hdr++ = c->msgused % 256;
*hdr++ = 0;
*hdr++ = 0;
assert((void *) hdr == (caddr_t)c->msglist[i].msg_iov[0].iov_base + UDP_HEADER_SIZE);
}

return 0;
unsigned char *hdr = c->hdrbuf;
*hdr++ = c->request_id / 256;
*hdr++ = c->request_id % 256;
*hdr++ = c->msgcurr / 256;
*hdr++ = c->msgcurr % 256;
*hdr++ = c->msgused / 256;
*hdr++ = c->msgused % 256;
*hdr++ = 0;
*hdr++ = 0;
header->iov_base = (void*)c->hdrbuf;
}

static void pipe_state_clear(conn *c)
Expand Down Expand Up @@ -2076,8 +2051,7 @@ out_mop_get_response(conn *c, bool delete, struct elems_result *eresultp)
} else {
sprintf(respptr, "END\r\n");
}
if ((add_iov(c, respptr, strlen(respptr)) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, respptr, strlen(respptr)) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -2241,8 +2215,7 @@ out_bop_trim_response(conn *c, void *elem, uint32_t flags)

if ((add_iov(c, respptr, resplen) != 0) ||
(add_iov_einfo_ascii(c, &c->einfo) != 0) ||
(add_iov(c, "TRIMMED\r\n", 9) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0))
(add_iov(c, "TRIMMED\r\n", 9) != 0))
{
ret = ENGINE_ENOMEM; break;
}
Expand Down Expand Up @@ -2527,8 +2500,7 @@ static void process_bop_mget_complete(conn *c)
}
if (k == c->coll_numkeys) {
sprintf(resultptr, "END\r\n");
if ((add_iov(c, resultptr, strlen(resultptr)) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, resultptr, strlen(resultptr)) != 0) {
ret = ENGINE_ENOMEM;
}
}
Expand Down Expand Up @@ -2652,8 +2624,7 @@ out_bop_smget_old_response(conn *c, token_t *key_tokens,
} else {
sprintf(respptr, (duplicated ? "DUPLICATED\r\n" : "END\r\n"));
}
if ((add_iov(c, respptr, strlen(respptr)) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, respptr, strlen(respptr)) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -2857,8 +2828,7 @@ out_bop_smget_response(conn *c, token_t *key_tokens, smget_result_t *smresp)
}

sprintf(respptr, (smresp->duplicated ? "DUPLICATED\r\n" : "END\r\n"));
if ((add_iov(c, respptr, strlen(respptr)) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, respptr, strlen(respptr)) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -3137,8 +3107,7 @@ static void process_mget_complete(conn *c, bool return_cas)
* reliable to add END\r\n to the buffer, because it might not end
* in \r\n. So we send SERVER_ERROR instead.
*/
if ((add_iov(c, "END\r\n", 5) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, "END\r\n", 5) != 0) {
/* Releasing items on ilist and freeing suffixes will be
* performed later by calling out_string() function.
* See conn_write() and conn_mwrite() state.
Expand Down Expand Up @@ -8378,8 +8347,7 @@ static inline void process_get_command(conn *c, token_t *tokens, size_t ntokens,
* reliable to add END\r\n to the buffer, because it might not end
* in \r\n. So we send SERVER_ERROR instead.
*/
if ((add_iov(c, "END\r\n", 5) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, "END\r\n", 5) != 0) {
out_string(c, "SERVER_ERROR out of memory writing get response");
} else {
conn_set_state(c, conn_mwrite);
Expand Down Expand Up @@ -9725,8 +9693,7 @@ static void process_prefixscan_command(conn *c, token_t *tokens, const size_t nt
attrptr += PREFIXSCAN_RESPONSE_ATTR_MAX_LENGTH;
}
if (ret != ENGINE_SUCCESS) break;
if (add_iov(c, "END\r\n", 5) != 0 ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, "END\r\n", 5) != 0) {
ret = ENGINE_ENOMEM;
}
} while (0);
Expand Down Expand Up @@ -9869,8 +9836,7 @@ static void process_keyscan_command(conn *c, token_t *tokens, const size_t ntoke
attrptr += KEYSCAN_RESPONSE_ATTR_MAX_LENGTH;
}
if (ret != ENGINE_SUCCESS) break;
if (add_iov(c, "END\r\n", 5) != 0 ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, "END\r\n", 5) != 0) {
ret = ENGINE_ENOMEM;
}
} while (0);
Expand Down Expand Up @@ -10170,8 +10136,7 @@ out_lop_get_response(conn *c, bool delete, struct elems_result *eresultp)
} else {
sprintf(respptr, "END\r\n");
}
if ((add_iov(c, respptr, strlen(respptr)) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, respptr, strlen(respptr)) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -10562,8 +10527,7 @@ out_sop_get_response(conn *c, bool delete, struct elems_result *eresultp)
} else {
sprintf(respptr, "END\r\n");
}
if ((add_iov(c, respptr, strlen(respptr)) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, respptr, strlen(respptr)) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -10932,8 +10896,7 @@ out_bop_get_response(conn *c, bool delete, struct elems_result *eresultp)
} else {
sprintf(respptr, "%s\r\n", (eresultp->trimmed ? "TRIMMED" : "END"));
}
if ((add_iov(c, respptr, strlen(respptr)) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, respptr, strlen(respptr)) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -11136,8 +11099,7 @@ out_bop_pwg_response(conn *c, int position, struct elems_result *eresultp)
if (ret == ENGINE_ENOMEM) break;

/* make response tail */
if ((add_iov(c, "END\r\n", 5) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, "END\r\n", 5) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -11245,8 +11207,7 @@ out_bop_gbp_response(conn *c, struct elems_result *eresultp)
if (ret == ENGINE_ENOMEM) break;

/* make response tail */
if ((add_iov(c, "END\r\n", 5) != 0) ||
(IS_UDP(c->transport) && build_udp_headers(c) != 0)) {
if (add_iov(c, "END\r\n", 5) != 0) {
ret = ENGINE_ENOMEM; break;
}
} while(0);
Expand Down Expand Up @@ -13495,11 +13456,20 @@ static enum transmit_result transmit(conn *c)
{
assert(c != NULL);

if (c->msgcurr < c->msgused && c->msglist[c->msgcurr].msg_iovlen == 0) {
while (c->msgcurr < c->msgused) {
/* Finished writing the current msg; advance to the next. */
c->msgcurr++;
}
if (c->msgcurr < c->msgused) {
if (c->msglist[c->msgcurr].msg_iovlen == 0) {
c->msgcurr++;
if (c->msgcurr == c->msgused) break;
if (IS_UDP(c->transport)) {
build_udp_header(c, c->msglist[c->msgcurr].msg_iov);
}
}
/* */
if (IS_UDP(c->transport) && c->msgcurr == 0) {
build_udp_header(c, c->iov);
}

ssize_t res;
struct msghdr *m = &c->msglist[c->msgcurr];

Expand Down Expand Up @@ -13546,9 +13516,9 @@ static enum transmit_result transmit(conn *c)
else
conn_set_state(c, conn_closing);
return TRANSMIT_HARD_ERROR;
} else {
return TRANSMIT_COMPLETE;
}

return TRANSMIT_COMPLETE;
}

bool conn_listening(conn *c)
Expand Down Expand Up @@ -13890,15 +13860,6 @@ bool conn_mwrite(conn *c)
/* The response must be reset according to c->aiostat. */
}

if (IS_UDP(c->transport) && c->msgcurr == 0 && build_udp_headers(c) != 0) {
if (settings.verbose > 0) {
mc_logger->log(EXTENSION_LOG_WARNING, c,
"Failed to build UDP headers in conn_mwrite.\n");
}
conn_set_state(c, conn_closing);
return true;
}

/* Clear the ewouldblock so that the next read command from
* the same connection does not falsely block and time out.
*/
Expand Down
3 changes: 1 addition & 2 deletions memcached.h
Original file line number Diff line number Diff line change
Expand Up @@ -383,8 +383,7 @@ struct conn {
int request_id; /* Incoming UDP request ID, if this is a UDP "connection" */
struct sockaddr request_addr; /* Who sent the most recent request */
socklen_t request_addr_size;
unsigned char *hdrbuf; /* udp packet headers */
int hdrsize; /* number of headers' worth of space is allocated */
unsigned char hdrbuf[UDP_HEADER_SIZE]; /* udp packet headers */

/* command pipelining processing fields */
int pipe_state;
Expand Down

0 comments on commit d4cd2da

Please sign in to comment.