git.m6w6.name Git - m6w6/libmemcached/blob - example/memcached_light.c

   1 /* -*- Mode: C; tab-width: 2; c-basic-offset: 2; indent-tabs-mode: nil -*- */
   2 /**
   3  * What is a library without an example to show you how to use the library?
   4  * This example use both interfaces to implement a small memcached server.
   5  * Please note that this is an exemple on how to use the library, not
   6  * an implementation of a scalable memcached server. If you look closely
   7  * at the example it isn't even multithreaded ;-)
   8  *
   9  * With that in mind, let me give you some pointers into the source:
  10  *   storage.c/h       - Implements the item store for this server and not really
  11  *                       interesting for this example.
  12  *   interface_v0.c    - Shows an implementation of the memcached server by using
  13  *                       the "raw" access to the packets as they arrive
  14  *   interface_v1.c    - Shows an implementation of the memcached server by using
  15  *                       the more "logical" interface.
  16  *   memcached_light.c - This file sets up all of the sockets and run the main
  17  *                       message loop.
  18  *
  19  *
  20  * config.h is included so that I can use the ntohll/htonll on platforms that
  21  * doesn't have that (this is a private function inside libmemcached, so you
  22  * cannot use it directly from libmemcached without special modifications to
  23  * the library)
  24  */
  25
  26 #include "config.h"
  27 #include <assert.h>
  28 #include <sys/types.h>
  29 #include <sys/socket.h>
  30 #include <netdb.h>
  31 #include <netinet/tcp.h>
  32 #include <stdio.h>
  33 #include <unistd.h>
  34 #include <fcntl.h>
  35 #include <errno.h>
  36 #include <stdlib.h>
  37 #include <string.h>
  38 #include <event.h>
  39
  40 #include <libmemcached/protocol_handler.h>
  41 #include <libmemcached/byteorder.h>
  42 #include "storage.h"
  43 #include "memcached_light.h"
  44
  45 extern memcached_binary_protocol_callback_st interface_v0_impl;
  46 extern memcached_binary_protocol_callback_st interface_v1_impl;
  47
  48 static int server_sockets[1024];
  49 static int num_server_sockets= 0;
  50
  51 struct connection
  52 {
  53   void *userdata;
  54   struct event event;
  55 };
  56
  57 /* The default maximum number of connections... (change with -c) */
  58 static int maxconns = 1024;
  59
  60 static struct connection *socket_userdata_map;
  61 static bool verbose= false;
  62 static struct event_base *event_base;
  63
  64 struct options_st {
  65   char *pid_file;
  66   bool has_port;
  67   in_port_t port;
  68 } global_options;
  69
  70 typedef struct options_st options_st;
  71
  72 /**
  73  * Callback for driving a client connection
  74  * @param fd the socket for the client socket
  75  * @param which identifying the event that occurred (not used)
  76  * @param arg the connection structure for the client
  77  */
  78 static void drive_client(int fd, short which, void *arg)
  79 {
  80   (void)which;
  81   struct connection *client= arg;
  82   struct memcached_protocol_client_st* c= client->userdata;
  83   assert(c != NULL);
  84
  85   memcached_protocol_event_t events= memcached_protocol_client_work(c);
  86   if (events & MEMCACHED_PROTOCOL_ERROR_EVENT)
  87   {
  88     memcached_protocol_client_destroy(c);
  89     (void)close(fd);
  90   } else {
  91     short flags = 0;
  92     if (events & MEMCACHED_PROTOCOL_WRITE_EVENT)
  93     {
  94       flags= EV_WRITE;
  95     }
  96
  97     if (events & MEMCACHED_PROTOCOL_READ_EVENT)
  98     {
  99       flags|= EV_READ;
 100     }
 101
 102     event_set(&client->event, fd, flags, drive_client, client);
 103     event_base_set(event_base, &client->event);
 104
 105     if (event_add(&client->event, 0) == -1)
 106     {
 107       (void)fprintf(stderr, "Failed to add event for %d\n", fd);
 108       memcached_protocol_client_destroy(c);
 109       (void)close(fd);
 110     }
 111   }
 112 }
 113
 114 /**
 115  * Callback for accepting new connections
 116  * @param fd the socket for the server socket
 117  * @param which identifying the event that occurred (not used)
 118  * @param arg the connection structure for the server
 119  */
 120 static void accept_handler(int fd, short which, void *arg)
 121 {
 122   (void)which;
 123   struct connection *server= arg;
 124   /* accept new client */
 125   struct sockaddr_storage addr;
 126   socklen_t addrlen= sizeof(addr);
 127   int sock= accept(fd, (struct sockaddr *)&addr, &addrlen);
 128
 129   if (sock == -1)
 130   {
 131     perror("Failed to accept client");
 132     return ;
 133   }
 134
 135   if (sock >= maxconns)
 136   {
 137     (void)fprintf(stderr, "Client outside socket range (specified with -c)\n");
 138     (void)close(sock);
 139     return ;
 140   }
 141
 142   struct memcached_protocol_client_st* c;
 143   c= memcached_protocol_create_client(server->userdata, sock);
 144   if (c == NULL)
 145   {
 146     (void)fprintf(stderr, "Failed to create client\n");
 147     (void)close(sock);
 148   }
 149   else
 150   {
 151     struct connection *client = &socket_userdata_map[sock];
 152     client->userdata= c;
 153
 154     event_set(&client->event, sock, EV_READ, drive_client, client);
 155     event_base_set(event_base, &client->event);
 156     if (event_add(&client->event, 0) == -1)
 157     {
 158       (void)fprintf(stderr, "Failed to add event for %d\n", sock);
 159       memcached_protocol_client_destroy(c);
 160       (void)close(sock);
 161     }
 162   }
 163 }
 164
 165 /**
 166  * Create a socket and bind it to a specific port number
 167  * @param port the port number to bind to
 168  */
 169 static int server_socket(const char *port)
 170 {
 171   struct addrinfo *ai;
 172   struct addrinfo hints= { .ai_flags= AI_PASSIVE,
 173                            .ai_family= AF_UNSPEC,
 174                            .ai_socktype= SOCK_STREAM };
 175
 176   int error= getaddrinfo("127.0.0.1", port, &hints, &ai);
 177   if (error != 0)
 178   {
 179     if (error != EAI_SYSTEM)
 180       fprintf(stderr, "getaddrinfo(): %s\n", gai_strerror(error));
 181     else
 182       perror("getaddrinfo()");
 183
 184     return 1;
 185   }
 186
 187   struct linger ling= {0, 0};
 188
 189   for (struct addrinfo *next= ai; next; next= next->ai_next)
 190   {
 191     int sock= socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
 192     if (sock == -1)
 193     {
 194       perror("Failed to create socket");
 195       continue;
 196     }
 197
 198     int flags= fcntl(sock, F_GETFL, 0);
 199     if (flags == -1)
 200     {
 201       perror("Failed to get socket flags");
 202       close(sock);
 203       continue;
 204     }
 205
 206     if ((flags & O_NONBLOCK) != O_NONBLOCK)
 207     {
 208       if (fcntl(sock, F_SETFL, flags | O_NONBLOCK) == -1)
 209       {
 210         perror("Failed to set socket to nonblocking mode");
 211         close(sock);
 212         continue;
 213       }
 214     }
 215
 216     flags= 1;
 217     if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (void *)&flags, sizeof(flags)) != 0)
 218       perror("Failed to set SO_REUSEADDR");
 219
 220     if (setsockopt(sock, SOL_SOCKET, SO_KEEPALIVE, (void *)&flags, sizeof(flags)) != 0)
 221       perror("Failed to set SO_KEEPALIVE");
 222
 223     if (setsockopt(sock, SOL_SOCKET, SO_LINGER, (void *)&ling, sizeof(ling)) != 0)
 224       perror("Failed to set SO_LINGER");
 225
 226     if (setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (void *)&flags, sizeof(flags)) != 0)
 227       perror("Failed to set TCP_NODELAY");
 228
 229     if (bind(sock, next->ai_addr, next->ai_addrlen) == -1)
 230     {
 231       if (errno != EADDRINUSE)
 232       {
 233         perror("bind()");
 234         freeaddrinfo(ai);
 235       }
 236       close(sock);
 237       continue;
 238     }
 239
 240     if (listen(sock, 1024) == -1)
 241     {
 242       perror("listen()");
 243       close(sock);
 244       continue;
 245     }
 246
 247     server_sockets[num_server_sockets++]= sock;
 248   }
 249
 250   freeaddrinfo(ai);
 251
 252   return (num_server_sockets > 0) ? 0 : 1;
 253 }
 254
 255 /**
 256  * Convert a command code to a textual string
 257  * @param cmd the comcode to convert
 258  * @return a textual string with the command or NULL for unknown commands
 259  */
 260 static const char* comcode2str(uint8_t cmd)
 261 {
 262   static const char * const text[] = {
 263     "GET", "SET", "ADD", "REPLACE", "DELETE",
 264     "INCREMENT", "DECREMENT", "QUIT", "FLUSH",
 265     "GETQ", "NOOP", "VERSION", "GETK", "GETKQ",
 266     "APPEND", "PREPEND", "STAT", "SETQ", "ADDQ",
 267     "REPLACEQ", "DELETEQ", "INCREMENTQ", "DECREMENTQ",
 268     "QUITQ", "FLUSHQ", "APPENDQ", "PREPENDQ"
 269   };
 270
 271   if (cmd <= PROTOCOL_BINARY_CMD_PREPENDQ)
 272     return text[cmd];
 273
 274   return NULL;
 275 }
 276
 277 /**
 278  * Print out the command we are about to execute
 279  */
 280 static void pre_execute(const void *cookie __attribute__((unused)),
 281                         protocol_binary_request_header *header __attribute__((unused)))
 282 {
 283   if (verbose)
 284   {
 285     const char *cmd= comcode2str(header->request.opcode);
 286     if (cmd != NULL)
 287       fprintf(stderr, "pre_execute from %p: %s\n", cookie, cmd);
 288     else
 289       fprintf(stderr, "pre_execute from %p: 0x%02x\n", cookie, header->request.opcode);
 290   }
 291 }
 292
 293 /**
 294  * Print out the command we just executed
 295  */
 296 static void post_execute(const void *cookie __attribute__((unused)),
 297                          protocol_binary_request_header *header __attribute__((unused)))
 298 {
 299   if (verbose)
 300   {
 301     const char *cmd= comcode2str(header->request.opcode);
 302     if (cmd != NULL)
 303       fprintf(stderr, "post_execute from %p: %s\n", cookie, cmd);
 304     else
 305       fprintf(stderr, "post_execute from %p: 0x%02x\n", cookie, header->request.opcode);
 306   }
 307 }
 308
 309 /**
 310  * Callback handler for all unknown commands.
 311  * Send an unknown command back to the client
 312  */
 313 static protocol_binary_response_status unknown(const void *cookie,
 314                                                protocol_binary_request_header *header,
 315                                                memcached_binary_protocol_raw_response_handler response_handler)
 316 {
 317   protocol_binary_response_no_extras response= {
 318     .message= {
 319       .header.response= {
 320         .magic= PROTOCOL_BINARY_RES,
 321         .opcode= header->request.opcode,
 322         .status= htons(PROTOCOL_BINARY_RESPONSE_UNKNOWN_COMMAND),
 323         .opaque= header->request.opaque
 324       }
 325     }
 326   };
 327
 328   return response_handler(cookie, header, (void*)&response);
 329 }
 330
 331 /**
 332  * Program entry point. Bind to the specified port(s) and serve clients
 333  *
 334  * @param argc number of items in the argument vector
 335  * @param argv argument vector
 336  * @return 0 on success, 1 otherwise
 337  */
 338 int main(int argc, char **argv)
 339 {
 340   int cmd;
 341   memcached_binary_protocol_callback_st *interface= &interface_v0_impl;
 342
 343   memset(&global_options, 0, sizeof(global_options));
 344
 345   event_base= event_init();
 346   if (event_base == NULL)
 347   {
 348     fprintf(stderr, "Failed to create an instance of libevent\n");
 349     return 1;
 350   }
 351
 352   /*
 353    * We need to initialize the handlers manually due to a bug in the
 354    * warnings generated by struct initialization in gcc (all the way up to 4.4)
 355    */
 356   initialize_interface_v0_handler();
 357
 358   while ((cmd= getopt(argc, argv, "v1p:P:?hc:")) != EOF)
 359   {
 360     switch (cmd) {
 361     case '1':
 362       interface= &interface_v1_impl;
 363       break;
 364     case 'P':
 365       global_options.pid_file= strdup(optarg);
 366       break;
 367     case 'p':
 368       global_options.has_port= true;
 369       (void)server_socket(optarg);
 370       break;
 371     case 'v':
 372       verbose= true;
 373       break;
 374     case 'c':
 375       maxconns= atoi(optarg);
 376       break;
 377     case 'h':  /* FALLTHROUGH */
 378     case '?':  /* FALLTHROUGH */
 379     default:
 380       (void)fprintf(stderr, "Usage: %s [-p port] [-v] [-1] [-c #clients] [-P pidfile]\n",
 381                     argv[0]);
 382       return 1;
 383     }
 384   }
 385
 386   if (! initialize_storage())
 387   {
 388     /* Error message already printed */
 389     return 1;
 390   }
 391
 392   if (! global_options.has_port)
 393     (void)server_socket("9999");
 394
 395   if (global_options.pid_file)
 396   {
 397     FILE *pid_file;
 398     uint32_t pid;
 399
 400     pid_file= fopen(global_options.pid_file, "w+");
 401
 402     if (pid_file == NULL)
 403     {
 404       perror(strerror(errno));
 405       abort();
 406     }
 407
 408     pid= (uint32_t)getpid();
 409     fprintf(pid_file, "%u\n", pid);
 410     fclose(pid_file);
 411   }
 412
 413   if (num_server_sockets == 0)
 414   {
 415     fprintf(stderr, "I don't have any server sockets\n");
 416     return 1;
 417   }
 418
 419   /*
 420    * Create and initialize the handles to the protocol handlers. I want
 421    * to be able to trace the traffic throught the pre/post handlers, and
 422    * set up a common handler for unknown messages
 423    */
 424   interface->pre_execute= pre_execute;
 425   interface->post_execute= post_execute;
 426   interface->unknown= unknown;
 427
 428   struct memcached_protocol_st *protocol_handle;
 429   if ((protocol_handle= memcached_protocol_create_instance()) == NULL)
 430   {
 431     fprintf(stderr, "Failed to allocate protocol handle\n");
 432     return 1;
 433   }
 434
 435   socket_userdata_map= calloc((size_t)(maxconns), sizeof(struct connection));
 436   if (socket_userdata_map == NULL)
 437   {
 438     fprintf(stderr, "Failed to allocate room for connections\n");
 439     return 1;
 440   }
 441
 442   memcached_binary_protocol_set_callbacks(protocol_handle, interface);
 443   memcached_binary_protocol_set_pedantic(protocol_handle, true);
 444
 445   for (int xx= 0; xx < num_server_sockets; ++xx)
 446   {
 447     struct connection *conn= &socket_userdata_map[server_sockets[xx]];
 448     conn->userdata= protocol_handle;
 449     event_set(&conn->event, server_sockets[xx], EV_READ | EV_PERSIST,
 450               accept_handler, conn);
 451     event_base_set(event_base, &conn->event);
 452     if (event_add(&conn->event, 0) == -1)
 453     {
 454       fprintf(stderr, "Failed to add event for %d\n", server_sockets[xx]);
 455       close(server_sockets[xx]);
 456     }
 457   }
 458
 459   /* Serve all of the clients */
 460   event_base_loop(event_base, 0);
 461
 462   /* NOTREACHED */
 463   return 0;
 464 }