/tests/websocket.c

http://github.com/nicolasff/webdis · C · 326 lines · 256 code · 54 blank · 16 comment · 39 complexity · 2c2c6e5110b061eecad24228a8ae5315 MD5 · raw file

  1. /* http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76 */
  2. #include <stdlib.h>
  3. #define _GNU_SOURCE
  4. #include <stdio.h>
  5. #include <string.h>
  6. #include <unistd.h>
  7. #include <pthread.h>
  8. #include <arpa/inet.h>
  9. #include <errno.h>
  10. #include <sys/types.h>
  11. #include <sys/socket.h>
  12. #include <event.h>
  13. struct host_info {
  14. char *host;
  15. short port;
  16. };
  17. /* worker_thread, with counter of remaining messages */
  18. struct worker_thread {
  19. struct host_info *hi;
  20. struct event_base *base;
  21. int msg_target;
  22. int msg_received;
  23. int msg_sent;
  24. int byte_count;
  25. pthread_t thread;
  26. struct evbuffer *buffer;
  27. int got_header;
  28. int verbose;
  29. struct event ev_w;
  30. };
  31. void
  32. process_message(struct worker_thread *wt, size_t sz) {
  33. // printf("process_message\n");
  34. if(wt->msg_received % 10000 == 0) {
  35. printf("thread %u: %8d messages left (got %9d bytes so far).\n",
  36. (unsigned int)wt->thread,
  37. wt->msg_target - wt->msg_received, wt->byte_count);
  38. }
  39. wt->byte_count += sz;
  40. /* decrement read count, and stop receiving when we reach zero. */
  41. wt->msg_received++;
  42. if(wt->msg_received == wt->msg_target) {
  43. event_base_loopexit(wt->base, NULL);
  44. }
  45. }
  46. void
  47. websocket_write(int fd, short event, void *ptr) {
  48. int ret;
  49. struct worker_thread *wt = ptr;
  50. if(event != EV_WRITE) {
  51. return;
  52. }
  53. char message[] = "\x00[\"SET\",\"key\",\"value\"]\xff\x00[\"GET\",\"key\"]\xff";
  54. ret = write(fd, message, sizeof(message)-1);
  55. if(ret != sizeof(message)-1) {
  56. fprintf(stderr, "write on %d failed: %s\n", fd, strerror(errno));
  57. close(fd);
  58. }
  59. wt->msg_sent += 2;
  60. if(wt->msg_sent < wt->msg_target) {
  61. event_set(&wt->ev_w, fd, EV_WRITE, websocket_write, wt);
  62. event_base_set(wt->base, &wt->ev_w);
  63. ret = event_add(&wt->ev_w, NULL);
  64. }
  65. }
  66. static void
  67. websocket_read(int fd, short event, void *ptr) {
  68. char packet[2048], *pos;
  69. int ret, success = 1;
  70. struct worker_thread *wt = ptr;
  71. if(event != EV_READ) {
  72. return;
  73. }
  74. /* read message */
  75. ret = read(fd, packet, sizeof(packet));
  76. pos = packet;
  77. if(ret > 0) {
  78. char *data, *last;
  79. int sz, msg_sz;
  80. if(wt->got_header == 0) { /* first response */
  81. char *frame_start = strstr(packet, "MH"); /* end of the handshake */
  82. if(frame_start == NULL) {
  83. return; /* not yet */
  84. } else { /* start monitoring possible writes */
  85. printf("start monitoring possible writes\n");
  86. evbuffer_add(wt->buffer, frame_start + 2, ret - (frame_start + 2 - packet));
  87. wt->got_header = 1;
  88. event_set(&wt->ev_w, fd, EV_WRITE,
  89. websocket_write, wt);
  90. event_base_set(wt->base, &wt->ev_w);
  91. ret = event_add(&wt->ev_w, NULL);
  92. }
  93. } else {
  94. /* we've had the header already, now bufffer data. */
  95. evbuffer_add(wt->buffer, packet, ret);
  96. }
  97. while(1) {
  98. data = (char*)EVBUFFER_DATA(wt->buffer);
  99. sz = EVBUFFER_LENGTH(wt->buffer);
  100. if(sz == 0) { /* no data */
  101. break;
  102. }
  103. if(*data != 0) { /* missing frame start */
  104. success = 0;
  105. break;
  106. }
  107. last = memchr(data, 0xff, sz); /* look for frame end */
  108. if(!last) {
  109. /* no end of frame in sight. */
  110. break;
  111. }
  112. msg_sz = last - data - 1;
  113. process_message(ptr, msg_sz); /* record packet */
  114. /* drain including frame delimiters (+2 bytes) */
  115. evbuffer_drain(wt->buffer, msg_sz + 2);
  116. }
  117. } else {
  118. printf("ret=%d\n", ret);
  119. success = 0;
  120. }
  121. if(success == 0) {
  122. shutdown(fd, SHUT_RDWR);
  123. close(fd);
  124. event_base_loopexit(wt->base, NULL);
  125. }
  126. }
  127. void*
  128. worker_main(void *ptr) {
  129. char ws_template[] = "GET /.json HTTP/1.1\r\n"
  130. "Host: %s:%d\r\n"
  131. "Connection: Upgrade\r\n"
  132. "Upgrade: WebSocket\r\n"
  133. "Origin: http://%s:%d\r\n"
  134. "Sec-WebSocket-Key1: 18x 6]8vM;54 *(5: { U1]8 z [ 8\r\n"
  135. "Sec-WebSocket-Key2: 1_ tx7X d < nw 334J702) 7]o}` 0\r\n"
  136. "\r\n"
  137. "Tm[K T2u";
  138. struct worker_thread *wt = ptr;
  139. int ret;
  140. int fd;
  141. struct sockaddr_in addr;
  142. char *ws_handshake;
  143. size_t ws_handshake_sz;
  144. /* connect socket */
  145. fd = socket(AF_INET, SOCK_STREAM, 0);
  146. addr.sin_family = AF_INET;
  147. addr.sin_port = htons(wt->hi->port);
  148. memset(&(addr.sin_addr), 0, sizeof(addr.sin_addr));
  149. addr.sin_addr.s_addr = inet_addr(wt->hi->host);
  150. ret = connect(fd, (struct sockaddr*)&addr, sizeof(struct sockaddr));
  151. if(ret != 0) {
  152. fprintf(stderr, "connect: ret=%d: %s\n", ret, strerror(errno));
  153. return NULL;
  154. }
  155. /* initialize worker thread */
  156. wt->base = event_base_new();
  157. wt->buffer = evbuffer_new();
  158. wt->byte_count = 0;
  159. wt->got_header = 0;
  160. /* send handshake */
  161. ws_handshake_sz = sizeof(ws_handshake)
  162. + 2*strlen(wt->hi->host) + 500;
  163. ws_handshake = calloc(ws_handshake_sz, 1);
  164. ws_handshake_sz = (size_t)sprintf(ws_handshake, ws_template,
  165. wt->hi->host, wt->hi->port,
  166. wt->hi->host, wt->hi->port);
  167. ret = write(fd, ws_handshake, ws_handshake_sz);
  168. struct event ev_r;
  169. event_set(&ev_r, fd, EV_READ | EV_PERSIST, websocket_read, wt);
  170. event_base_set(wt->base, &ev_r);
  171. event_add(&ev_r, NULL);
  172. /* go! */
  173. event_base_dispatch(wt->base);
  174. event_base_free(wt->base);
  175. free(ws_handshake);
  176. return NULL;
  177. }
  178. void
  179. usage(const char* argv0, char *host_default, short port_default,
  180. int thread_count_default, int messages_default) {
  181. printf("Usage: %s [options]\n"
  182. "Options are:\n"
  183. "\t-h host\t\t(default = \"%s\")\n"
  184. "\t-p port\t\t(default = %d)\n"
  185. "\t-c threads\t(default = %d)\n"
  186. "\t-n count\t(number of messages per thread, default = %d)\n"
  187. "\t-v\t\t(verbose)\n",
  188. argv0, host_default, (int)port_default,
  189. thread_count_default, messages_default);
  190. }
  191. int
  192. main(int argc, char *argv[]) {
  193. struct timespec t0, t1;
  194. int messages_default = 100000;
  195. int thread_count_default = 4;
  196. short port_default = 7379;
  197. char *host_default = "127.0.0.1";
  198. int msg_target = messages_default;
  199. int thread_count = thread_count_default;
  200. int i, opt;
  201. char *colon;
  202. double total = 0, total_bytes = 0;
  203. int verbose = 0;
  204. struct host_info hi = {host_default, port_default};
  205. struct worker_thread *workers;
  206. /* getopt */
  207. while ((opt = getopt(argc, argv, "h:p:c:n:v")) != -1) {
  208. switch (opt) {
  209. case 'h':
  210. colon = strchr(optarg, ':');
  211. if(!colon) {
  212. size_t sz = strlen(optarg);
  213. hi.host = calloc(1 + sz, 1);
  214. strncpy(hi.host, optarg, sz);
  215. } else {
  216. hi.host = calloc(1+colon-optarg, 1);
  217. strncpy(hi.host, optarg, colon-optarg);
  218. hi.port = (short)atol(colon+1);
  219. }
  220. break;
  221. case 'p':
  222. hi.port = (short)atol(optarg);
  223. break;
  224. case 'c':
  225. thread_count = atoi(optarg);
  226. break;
  227. case 'n':
  228. msg_target = atoi(optarg);
  229. break;
  230. case 'v':
  231. verbose = 1;
  232. break;
  233. default: /* '?' */
  234. usage(argv[0], host_default, port_default,
  235. thread_count_default,
  236. messages_default);
  237. exit(EXIT_FAILURE);
  238. }
  239. }
  240. /* run threads */
  241. workers = calloc(sizeof(struct worker_thread), thread_count);
  242. clock_gettime(CLOCK_MONOTONIC, &t0);
  243. for(i = 0; i < thread_count; ++i) {
  244. workers[i].msg_target = msg_target;
  245. workers[i].hi = &hi;
  246. workers[i].verbose = verbose;
  247. pthread_create(&workers[i].thread, NULL,
  248. worker_main, &workers[i]);
  249. }
  250. /* wait for threads to finish */
  251. for(i = 0; i < thread_count; ++i) {
  252. pthread_join(workers[i].thread, NULL);
  253. total += workers[i].msg_received;
  254. total_bytes += workers[i].byte_count;
  255. }
  256. /* timing */
  257. clock_gettime(CLOCK_MONOTONIC, &t1);
  258. float mili0 = t0.tv_sec * 1000 + t0.tv_nsec / 1000000;
  259. float mili1 = t1.tv_sec * 1000 + t1.tv_nsec / 1000000;
  260. if(total != 0) {
  261. printf("Read %ld messages in %0.2f sec: %0.2f msg/sec (%d MB/sec, %d KB/sec)\n",
  262. (long)total,
  263. (mili1-mili0)/1000.0,
  264. 1000*total/(mili1-mili0),
  265. (int)(total_bytes / (1000*(mili1-mili0))),
  266. (int)(total_bytes / (mili1-mili0)));
  267. return EXIT_SUCCESS;
  268. } else {
  269. printf("No message was read.\n");
  270. return EXIT_FAILURE;
  271. }
  272. }