1 /* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
2 #include "config.h"
3
4 #include <memcached/protocol_binary.h>
5
6 #include <getopt.h>
7 #include <errno.h>
8 #include <stdio.h>
9 #include <string.h>
10 #include <stdlib.h>
11
12 /**
13 * Try to connect to the server
14 * @param host the name of the server
15 * @param port the port to connect to
16 * @return a socket descriptor connected to host:port for success, -1 otherwise
17 */
connect_server(const char * hostname,const char * port)18 static int connect_server(const char *hostname, const char *port)
19 {
20 struct addrinfo *ai = NULL;
21 struct addrinfo hints = { .ai_family = AF_UNSPEC,
22 .ai_protocol = IPPROTO_TCP,
23 .ai_socktype = SOCK_STREAM };
24
25 if (getaddrinfo(hostname, port, &hints, &ai) != 0) {
26 return -1;
27 }
28 int sock = -1;
29 if ((sock = socket(ai->ai_family, ai->ai_socktype,
30 ai->ai_protocol)) != -1) {
31 if (connect(sock, ai->ai_addr, ai->ai_addrlen) == -1) {
32 close(sock);
33 sock = -1;
34 }
35 }
36
37 freeaddrinfo(ai);
38 return sock;
39 }
40
41 /**
42 * Send the chunk of data to the other side, retry if an error occurs
43 * (or terminate the program if retry wouldn't help us)
44 * @param sock socket to write data to
45 * @param buf buffer to send
46 * @param len length of data to send
47 */
retry_send(int sock,const void * buf,size_t len)48 static void retry_send(int sock, const void* buf, size_t len)
49 {
50 off_t offset = 0;
51 const char* ptr = buf;
52
53 do {
54 size_t num_bytes = len - offset;
55 ssize_t nw = send(sock, ptr + offset, num_bytes, 0);
56 if (nw == -1) {
57 if (errno != EINTR) {
58 fprintf(stderr, "Failed to write: %s\n", strerror(errno));
59 close(sock);
60 exit(1);
61 }
62 } else {
63 offset += nw;
64 }
65 } while (offset < len);
66 }
67
68 /**
69 * Receive a fixed number of bytes from the socket.
70 * (Terminate the program if we encounter a hard error...)
71 * @param sock socket to receive data from
72 * @param buf buffer to store data to
73 * @param len length of data to receive
74 */
retry_recv(int sock,void * buf,size_t len)75 static void retry_recv(int sock, void *buf, size_t len) {
76 if (len == 0) {
77 return;
78 }
79 off_t offset = 0;
80 do {
81 ssize_t nr = recv(sock, ((char*)buf) + offset, len - offset, 0);
82 if (nr == -1) {
83 if (errno != EINTR) {
84 fprintf(stderr, "Failed to read: %s\n", strerror(errno));
85 close(sock);
86 exit(1);
87 }
88 } else {
89 if (nr == 0) {
90 fprintf(stderr, "Connection closed\n");
91 close(sock);
92 exit(1);
93 }
94 offset += nr;
95 }
96 } while (offset < len);
97 }
98
99 /**
100 * Print the key value pair
101 * @param key key to print
102 * @param keylen length of key to print
103 * @param val value to print
104 * @param vallen length of value
105 */
print(const char * key,int keylen,const char * val,int vallen)106 static void print(const char *key, int keylen, const char *val, int vallen) {
107 fputs("STAT ", stdout);
108 (void)fwrite(key, keylen, 1, stdout);
109 fputs(" ", stdout);
110 (void)fwrite(val, vallen, 1, stdout);
111 fputs("\n", stdout);
112 fflush(stdout);
113 }
114
115 /**
116 * Request a stat from the server
117 * @param sock socket connected to the server
118 * @param key the name of the stat to receive (NULL == ALL)
119 */
request_stat(int sock,const char * key)120 static void request_stat(int sock, const char *key)
121 {
122 uint32_t buffsize = 0;
123 char *buffer = NULL;
124 uint16_t keylen = 0;
125 if (key != NULL) {
126 keylen = (uint16_t)strlen(key);
127 }
128
129 protocol_binary_request_stats request = {
130 .message.header.request = {
131 .magic = PROTOCOL_BINARY_REQ,
132 .opcode = PROTOCOL_BINARY_CMD_STAT,
133 .keylen = htons(keylen),
134 .bodylen = htonl(keylen)
135 }
136 };
137
138 retry_send(sock, &request, sizeof(request));
139 if (keylen > 0) {
140 retry_send(sock, key, keylen);
141 }
142
143 protocol_binary_response_no_extras response;
144 do {
145 retry_recv(sock, &response, sizeof(response.bytes));
146 if (response.message.header.response.keylen != 0) {
147 uint16_t keylen = ntohs(response.message.header.response.keylen);
148 uint32_t vallen = ntohl(response.message.header.response.bodylen);
149 if (vallen > buffsize) {
150 if ((buffer = realloc(buffer, vallen)) == NULL) {
151 fprintf(stderr, "Failed to allocate memory\n");
152 exit(1);
153 }
154 buffsize = vallen;
155 }
156 retry_recv(sock, buffer, vallen);
157 print(buffer, keylen, buffer + keylen, vallen - keylen);
158 }
159 } while (response.message.header.response.keylen != 0);
160 }
161
162 /**
163 * Program entry point. Connect to a memcached server and use the binary
164 * protocol to retrieve a given set of stats.
165 *
166 * @param argc argument count
167 * @param argv argument vector
168 * @return 0 if success, error code otherwise
169 */
main(int argc,char ** argv)170 int main(int argc, char **argv)
171 {
172 int cmd;
173 const char * const default_ports[] = { "memcache", "11211", NULL };
174 const char *port = NULL;
175 const char *host = NULL;
176 char *ptr;
177
178 /* Initialize the socket subsystem */
179 initialize_sockets();
180
181 while ((cmd = getopt(argc, argv, "h:p:")) != EOF) {
182 switch (cmd) {
183 case 'h' :
184 host = optarg;
185 ptr = strchr(optarg, ':');
186 if (ptr != NULL) {
187 *ptr = '\0';
188 port = ptr + 1;
189 }
190 break;
191 case 'p':
192 port = optarg;
193 break;
194 default:
195 fprintf(stderr,
196 "Usage mcstat [-h host[:port]] [-p port] [statkey]*\n");
197 return 1;
198 }
199 }
200
201 if (host == NULL) {
202 host = "localhost";
203 }
204
205 int sock = -1;
206 if (port == NULL) {
207 int ii = 0;
208 do {
209 port = default_ports[ii++];
210 sock = connect_server(host, port);
211 } while (sock == -1 && default_ports[ii] != NULL);
212 } else {
213 sock = connect_server(host, port);
214 }
215
216 if (sock == -1) {
217 fprintf(stderr, "Failed to connect to memcached server (%s:%s): %s\n",
218 host, port, strerror(errno));
219 return 1;
220 }
221
222 if (optind == argc) {
223 request_stat(sock, NULL);
224 } else {
225 for (int ii = optind; ii < argc; ++ii) {
226 request_stat(sock, argv[ii]);
227 }
228 }
229
230 close(sock);
231
232 return 0;
233 }
234