Update benchmark script to run a few other cases.
[bluesky.git] / simplestore / server.c
1 /* A very simple AWS-like server, without any authentication.  This is meant
2  * performance testing in a local environment.  The server offers weak
3  * guarantees on data durability--data is stored directly to the file system
4  * without synchronization, so data might be lost in a crash.  This should
5  * offer good performance though for benchmarking.
6  *
7  * Protocol: Each request is a whitespace-separated (typically, a single space)
8  * list of parameters terminated by a newline character.  The response is a
9  * line containing a response code and a body length (again white-separated and
10  * newline-terminated) followed by the response body.  Requests:
11  *   GET filename offset length
12  *   PUT filename length
13  *   DELETE filename
14  *   LIST prefix
15  */
16
17 #include <stdio.h>
18 #include <stdlib.h>
19 #include <assert.h>
20 #include <errno.h>
21 #include <fcntl.h>
22 #include <signal.h>
23 #include <string.h>
24 #include <time.h>
25 #include <unistd.h>
26 #include <netinet/in.h>
27 #include <sys/socket.h>
28 #include <sys/stat.h>
29 #include <sys/types.h>
30 #include <sys/wait.h>
31
32 #define SIMPLESTORE_PORT 9541
33
34 /* Maximum number of connections the server will queue waiting for us to call
35  * accept(). */
36 #define TCP_BACKLOG 32
37
38 #define WHITESPACE " \t\r\n"
39 #define MAX_ARGS 4
40
41 /* Maximum length of a command that we will accept. */
42 #define MAX_CMD_LEN 4096
43
44 #define MAX_OBJECT_SIZE (256 << 20)
45
46 enum command { GET, PUT, LIST, DELETE };
47
48 void write_data(int fd, const char *buf, size_t len)
49 {
50     while (len > 0) {
51         ssize_t bytes = write(fd, buf, len);
52         if (bytes < 0) {
53             if (errno == EINTR)
54                 continue;
55             exit(1);
56         }
57         buf += bytes;
58         len -= bytes;
59     }
60 }
61
62 /* SIGCHLD handler, used to clean up processes that handle the connections. */
63 static void sigchld_handler(int signal)
64 {
65     int pid;
66     int reaped = 0;
67
68     while ((pid = waitpid(WAIT_ANY, NULL, WNOHANG)) > 0) {
69         reaped++;
70     }
71
72     if (pid < 0) {
73         if (errno == ECHILD && reaped) {
74             /* Don't print an error for the caes that we successfully cleaned
75              * up after all children. */
76         } else {
77             perror("waitpid");
78         }
79     }
80 }
81
82 /* Return a text representation of a socket address.  Returns a pointer to a
83  * static buffer so it is non-reentrant. */
84 const char *sockname(struct sockaddr_in *addr, socklen_t len)
85 {
86     static char buf[128];
87     if (len < sizeof(struct sockaddr_in))
88         return NULL;
89     if (addr->sin_family != AF_INET)
90         return NULL;
91
92     uint32_t ip = ntohl(addr->sin_addr.s_addr);
93     sprintf(buf, "%d.%d.%d.%d:%d",
94             (int)((ip >> 24) & 0xff),
95             (int)((ip >> 16) & 0xff),
96             (int)((ip >> 8) & 0xff),
97             (int)(ip & 0xff),
98             ntohs(addr->sin_port));
99
100     return buf;
101 }
102
103 /* Convert a path from a client to the actual filename used.  Returns a string
104  * that must be freed. */
105 char *normalize_path(const char *in)
106 {
107     char *out = malloc(2 * strlen(in) + 1);
108     int i, j;
109     for (i = 0, j = 0; in[i] != '\0'; i++) {
110         if (in[i] == '/') {
111             out[j++] = '_';
112             out[j++] = 's';
113         } else if (in[i] == '_') {
114             out[j++] = '_';
115             out[j++] = 'u';
116         } else {
117             out[j++] = in[i];
118         }
119     }
120     out[j++] = '\0';
121     return out;
122 }
123
124 void cmd_get(int fd, char *path, size_t start, ssize_t len)
125 {
126     char buf[65536];
127
128     char *response = "-1\n";
129     int file = open(path, O_RDONLY);
130     if (file < 0) {
131         write_data(fd, response, strlen(response));
132         return;
133     }
134
135     struct stat statbuf;
136     if (fstat(file, &statbuf) < 0) {
137         write_data(fd, response, strlen(response));
138         return;
139     }
140
141     size_t datalen = statbuf.st_size;
142     if (start > 0) {
143         if (start >= datalen) {
144             datalen = 0;
145         } else {
146             lseek(file, start, SEEK_SET);
147             datalen -= start;
148         }
149     }
150     if (len > 0 && len < datalen) {
151         datalen = len;
152     }
153
154     sprintf(buf, "%zd\n", datalen);
155     write_data(fd, buf, strlen(buf));
156
157     while (datalen > 0) {
158         size_t needed = datalen > sizeof(buf) ? sizeof(buf) : datalen;
159         ssize_t bytes = read(file, buf, needed);
160         if (bytes < 0 && errno == EINTR)
161             continue;
162         if (bytes <= 0) {
163             /* Error reading necessary data, but we already told the client the
164              * file size so pad the data to the client with null bytes. */
165             memset(buf, 0, needed);
166             bytes = needed;
167         }
168         write_data(fd, buf, bytes);
169         datalen -= bytes;
170     }
171
172     close(file);
173 }
174
175 void cmd_put(int fd, char *path, char *buf,
176              size_t object_size, size_t buf_used)
177 {
178     while (buf_used < object_size) {
179         ssize_t bytes;
180
181         bytes = read(fd, buf + buf_used, object_size - buf_used);
182         if (bytes < 0) {
183             if (errno == EINTR)
184                 continue;
185             else
186                 exit(1);
187         }
188
189         if (bytes == 0)
190             exit(1);
191
192         assert(bytes <= object_size - buf_used);
193         buf_used += bytes;
194
195         continue;
196     }
197
198     /* printf("Got %zd bytes for object '%s'\n", buf_used, path); */
199
200     char *response = "-1\n";
201     int file = open(path, O_WRONLY|O_CREAT|O_TRUNC, 0600);
202     if (file >= 0) {
203         write_data(file, buf, object_size);
204         response = "0\n";
205         close(file);
206     }
207
208     write_data(fd, response, strlen(response));
209 }
210
211 /* The core handler for processing requests from the client.  This can be
212  * single-threaded since each connection is handled in a separate process. */
213 void handle_connection(int fd)
214 {
215     char cmdbuf[MAX_CMD_LEN];
216     size_t buflen = 0;
217
218     while (1) {
219         /* Keep reading data until reaching a newline, so that a complete
220          * command can be parsed. */
221         if (buflen == 0 || memchr(cmdbuf, '\n', buflen) == NULL) {
222             ssize_t bytes;
223
224             if (buflen == MAX_CMD_LEN) {
225                 /* Command is too long and thus malformed; close the
226                  * connection. */
227                 return;
228             }
229
230             bytes = read(fd, cmdbuf + buflen, MAX_CMD_LEN - buflen);
231             if (bytes < 0) {
232                 if (errno == EINTR)
233                     continue;
234                 else
235                     return;
236             }
237             if (bytes == 0)
238                 return;
239
240             assert(bytes <= MAX_CMD_LEN - buflen);
241             buflen += bytes;
242
243             continue;
244         }
245
246         size_t cmdlen = (char *)memchr(cmdbuf, '\n', buflen) - cmdbuf + 1;
247         cmdbuf[cmdlen - 1] = '\0';
248         char *token;
249         int arg_count;
250         char *args[MAX_ARGS];
251         int arg_int[MAX_ARGS];
252         for (token = strtok(cmdbuf, WHITESPACE), arg_count = 0;
253              token != NULL;
254              token = strtok(NULL, WHITESPACE), arg_count++)
255         {
256             args[arg_count] = token;
257             arg_int[arg_count] = atoi(token);
258         }
259
260         if (arg_count < 2) {
261             return;
262         }
263         char *path = normalize_path(args[1]);
264         enum command cmd;
265         if (strcmp(args[0], "GET") == 0 && arg_count == 4) {
266             cmd = GET;
267         } else if (strcmp(args[0], "PUT") == 0 && arg_count == 3) {
268             cmd = PUT;
269         } else if (strcmp(args[0], "DELETE") == 0 && arg_count == 2) {
270             cmd = DELETE;
271         } else {
272             return;
273         }
274
275         if (cmdlen < buflen)
276             memmove(cmdbuf, cmdbuf + cmdlen, buflen - cmdlen);
277         buflen -= cmdlen;
278
279         switch (cmd) {
280         case GET:
281             cmd_get(fd, path, arg_int[2], arg_int[3]);
282             break;
283         case PUT: {
284             size_t object_size = arg_int[2];
285             if (object_size > MAX_OBJECT_SIZE)
286                 return;
287             char *data_buf = malloc(object_size);
288             if (data_buf == NULL)
289                 return;
290             size_t data_buflen = buflen > object_size ? object_size : buflen;
291             if (data_buflen > 0)
292                 memcpy(data_buf, cmdbuf, data_buflen);
293             if (data_buflen < buflen) {
294                 memmove(cmdbuf, cmdbuf + data_buflen, buflen - data_buflen);
295                 buflen -= cmdlen;
296             } else {
297                 buflen = 0;
298             }
299             cmd_put(fd, path, data_buf, object_size, data_buflen);
300             break;
301         }
302         case DELETE:
303             //cmd_delete(fd, path);
304             break;
305         default:
306             return;
307         }
308
309         free(path);
310     }
311 }
312
313 /* Create a listening TCP socket on a new address (we do not use a fixed port).
314  * Return the file descriptor of the listening socket. */
315 int server_init()
316 {
317     int fd;
318     struct sockaddr_in server_addr;
319     socklen_t addr_len = sizeof(server_addr);
320
321     fd = socket(PF_INET, SOCK_STREAM, 0);
322     if (fd < 0) {
323         perror("socket");
324         exit(1);
325     }
326
327     int val = 1;
328     setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val));
329
330     server_addr.sin_family = AF_INET;
331     server_addr.sin_addr.s_addr = htonl(INADDR_ANY);
332     server_addr.sin_port = htons(SIMPLESTORE_PORT);
333     if (bind(fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
334         perror("bind");
335         exit(1);
336     }
337
338     if (listen(fd, TCP_BACKLOG) < 0) {
339         perror("listen");
340         exit(1);
341     }
342
343     if (getsockname(fd, (struct sockaddr *)&server_addr, &addr_len) < 0) {
344         perror("getsockname");
345         exit(1);
346     }
347
348     printf("Server listening on %s ...\n", sockname(&server_addr, addr_len));
349     fflush(stdout);
350
351     return fd;
352 }
353
354 /* Process-based server main loop.  Wait for a connection, accept it, fork a
355  * child process to handle the connection, and repeat.  Child processes are
356  * reaped in the SIGCHLD handler. */
357 void server_main(int listen_fd)
358 {
359     struct sigaction handler;
360
361     /* Install signal handler for SIGCHLD. */
362     handler.sa_handler = sigchld_handler;
363     sigemptyset(&handler.sa_mask);
364     handler.sa_flags = SA_RESTART;
365     if (sigaction(SIGCHLD, &handler, NULL) < 0) {
366         perror("sigaction");
367         exit(1);
368     }
369
370     while (1) {
371         struct sockaddr_in client_addr;
372         socklen_t addr_len = sizeof(client_addr);
373         int client_fd = accept(listen_fd, (struct sockaddr *)&client_addr,
374                                &addr_len);
375         int pid;
376
377         /* Very simple error handling.  Exit if something goes wrong.  Later,
378          * might want to look into not killing off current connections abruptly
379          * if we encounter an error in the accept(). */
380         if (client_fd < 0) {
381             if (errno == EINTR)
382                 continue;
383
384             perror("accept");
385             exit(1);
386         }
387
388         printf("Accepted connection from %s ...\n",
389                sockname(&client_addr, addr_len));
390         fflush(stdout);
391
392         pid = fork();
393         if (pid < 0) {
394             perror("fork");
395         } else if (pid == 0) {
396             handle_connection(client_fd);
397             printf("Closing connection %s ...\n",
398                    sockname(&client_addr, addr_len));
399             close(client_fd);
400             exit(0);
401         }
402
403         close(client_fd);
404     }
405 }
406
407 /* Print a help message describing command-line options to stdout. */
408 static void usage(char *argv0)
409 {
410     printf("Usage: %s [options...]\n"
411            "A simple key-value storage server.\n", argv0);
412 }
413
414 int main(int argc, char *argv[])
415 {
416     int fd;
417     //int i;
418     int display_help = 0, cmdline_error = 0;
419
420 #if 0
421     for (i = 1; i < argc; i++) {
422         if (strcmp(argv[i], "-help") == 0) {
423             display_help = 1;
424         } else if (strcmp(argv[i], "-document_root") == 0) {
425             i++;
426             if (i >= argc) {
427                 fprintf(stderr,
428                         "Error: Argument to -document_root expected.\n");
429                 cmdline_error = 1;
430             } else {
431                 document_root = argv[i];
432             }
433         } else if (strcmp(argv[i], "-port") == 0) {
434             i++;
435             if (i >= argc) {
436                 fprintf(stderr,
437                         "Error: Expected port number after -port.\n");
438                 cmdline_error = 1;
439             } else {
440                 server_port = atoi(argv[i]);
441                 if (server_port < 1 || server_port > 65535) {
442                     fprintf(stderr,
443                             "Error: Port must be between 1 and 65535.\n");
444                     cmdline_error = 1;
445                 }
446             }
447         } else if (strcmp(argv[i], "-mime_types") == 0) {
448             i++;
449             if (i >= argc) {
450                 fprintf(stderr,
451                         "Error: Argument to -mime_types expected.\n");
452                 cmdline_error = 1;
453             } else {
454                 mime_types_file = argv[i];
455             }
456         } else if (strcmp(argv[i], "-log") == 0) {
457             i++;
458             if (i >= argc) {
459                 fprintf(stderr,
460                         "Error: Argument to -log expected.\n");
461                 cmdline_error = 1;
462             } else {
463                 log_fname = argv[i];
464             }
465         } else {
466             fprintf(stderr, "Error: Unrecognized option '%s'\n", argv[i]);
467             cmdline_error = 1;
468         }
469     }
470 #endif
471
472     /* Display a help message if requested, or let the user know to look at the
473      * help message if there was an error parsing the command line.  In either
474      * case, the server never starts. */
475     if (display_help) {
476         usage(argv[0]);
477         exit(0);
478     } else if (cmdline_error) {
479         fprintf(stderr, "Run '%s -help' for a summary of options.\n", argv[0]);
480         exit(2);
481     }
482
483     fd = server_init();
484     if (fd >= 0) {
485         server_main(fd);
486     }
487
488     return 0;
489 }