1 /* This is a simple yet efficient WebSocket server benchmark much like WRK */
2
3 #define _BSD_SOURCE
4 #include <endian.h>
5 #include <stdint.h>
6
7 #include <libusockets.h>
8 int SSL;
9
10 #include <stdio.h>
11 #include <stdlib.h>
12 #include <string.h>
13
14 /* Whatever type we selected (compressed or not) */
15 unsigned char *web_socket_request;
16 int web_socket_request_size;
17
18 char *upgrade_request;
19 int upgrade_request_length;
20
21 /* Compressed message */
22 unsigned char web_socket_request_deflate[13] = {
23 130 | 64, 128 | 7,
24 0, 0, 0, 0,
25 0xf2, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00
26 };
27
28 /* Not compressed */
29 unsigned char web_socket_request_text_small[26] = {130, 128 | 20, 1, 2, 3, 4};
30 unsigned int web_socket_request_text_size = 26;
31 unsigned char *web_socket_request_text = web_socket_request_text_small;
32
33 /* Called to swap from small text message to big text message */
init_big_message(unsigned int size)34 void init_big_message(unsigned int size) {
35 if (size < 65536) {
36 printf("Error: message size must be bigger\n");
37 exit(0);
38 }
39
40 web_socket_request_text_size = size + 6 + 8;
41
42 web_socket_request_text = malloc(web_socket_request_text_size);
43 web_socket_request_text[0] = 130;
44 web_socket_request_text[1] = 255;
45 uint64_t msg_size = htobe64(size);
46 memcpy(&web_socket_request_text[2], &msg_size, 8);
47 web_socket_request_text[10] = 1;
48 web_socket_request_text[10] = 2;
49 web_socket_request_text[10] = 3;
50 web_socket_request_text[10] = 4;
51 }
52
53 char request_deflate[] = "GET / HTTP/1.1\r\n"
54 "Upgrade: websocket\r\n"
55 "Connection: Upgrade\r\n"
56 "Sec-WebSocket-Key: x3JJHMbDL1EzLkh9GBhXDw==\r\n"
57 "Sec-WebSocket-Extensions: permessage-deflate; client_max_window_bits\r\n"
58 "Host: server.example.com\r\n"
59 "Sec-WebSocket-Version: 13\r\n\r\n";
60
61 char request_text[] = "GET / HTTP/1.1\r\n"
62 "Upgrade: websocket\r\n"
63 "Connection: Upgrade\r\n"
64 "Sec-WebSocket-Key: x3JJHMbDL1EzLkh9GBhXDw==\r\n"
65 //"Sec-WebSocket-Extensions: permessage-deflate; client_max_window_bits\r\n"
66 "Host: server.example.com\r\n"
67 "Sec-WebSocket-Version: 13\r\n\r\n";
68 char *host;
69 int port;
70 int connections;
71
72 int responses;
73
74 struct http_socket {
75 /* How far we have streamed our websocket request */
76 int offset;
77
78 /* How far we have streamed our upgrade request */
79 int upgrade_offset;
80
81 /* Whether or not we have received the upgrade response */
82 int is_upgraded;
83
84 /* How many bytes we expect to be echoed back to us before we consider the echo done */
85 int outstanding_bytes;
86 };
87
88 /* We don't need any of these */
on_wakeup(struct us_loop_t * loop)89 void on_wakeup(struct us_loop_t *loop) {
90
91 }
92
on_pre(struct us_loop_t * loop)93 void on_pre(struct us_loop_t *loop) {
94
95 }
96
97 /* This is not HTTP POST, it is merely an event emitted post loop iteration */
on_post(struct us_loop_t * loop)98 void on_post(struct us_loop_t *loop) {
99
100 }
101
next_connection(struct us_socket_t * s)102 void next_connection(struct us_socket_t *s) {
103 /* We could wait with this until properly upgraded */
104 if (--connections) {
105 us_socket_context_connect(SSL, us_socket_context(SSL, s), host, port, NULL, 0, sizeof(struct http_socket));
106 } else {
107 printf("Running benchmark now...\n");
108
109 us_socket_timeout(SSL, s, LIBUS_TIMEOUT_GRANULARITY);
110 }
111 }
112
on_http_socket_writable(struct us_socket_t * s)113 struct us_socket_t *on_http_socket_writable(struct us_socket_t *s) {
114 struct http_socket *http_socket = (struct http_socket *) us_socket_ext(SSL, s);
115
116 /* Are we still not upgraded yet? */
117 if (http_socket->upgrade_offset < upgrade_request_length) {
118 http_socket->upgrade_offset += us_socket_write(SSL, s, upgrade_request + http_socket->upgrade_offset, upgrade_request_length - http_socket->upgrade_offset, 0);
119
120 /* Now we should be */
121 if (http_socket->upgrade_offset == upgrade_request_length) {
122 next_connection(s);
123 }
124 } else {
125 /* Stream whatever is remaining of the request */
126 http_socket->offset += us_socket_write(SSL, s, (char *) web_socket_request + http_socket->offset, web_socket_request_size - http_socket->offset, 0);
127 }
128
129 return s;
130 }
131
on_http_socket_close(struct us_socket_t * s,int code,void * reason)132 struct us_socket_t *on_http_socket_close(struct us_socket_t *s, int code, void *reason) {
133
134 printf("Closed!\n");
135
136 return s;
137 }
138
on_http_socket_end(struct us_socket_t * s)139 struct us_socket_t *on_http_socket_end(struct us_socket_t *s) {
140 return us_socket_close(SSL, s, 0, NULL);
141 }
142
on_http_socket_data(struct us_socket_t * s,char * data,int length)143 struct us_socket_t *on_http_socket_data(struct us_socket_t *s, char *data, int length) {
144 /* Get socket extension and the socket's context's extension */
145 struct http_socket *http_socket = (struct http_socket *) us_socket_ext(SSL, s);
146
147 if (http_socket->is_upgraded) {
148
149 /* If we are upgraded we now count to see if we receive the corect echo */
150 http_socket->outstanding_bytes -= length;
151
152 if (http_socket->outstanding_bytes == 0) {
153 /* We got exactly the correct amount of bytes back, send another message */
154 http_socket->offset = us_socket_write(SSL, s, (char *) web_socket_request, web_socket_request_size, 0);
155 http_socket->outstanding_bytes = web_socket_request_size - 4;
156
157 /* Increase stats */
158 responses++;
159 } else if (http_socket->outstanding_bytes < 0) {
160 /* This should never happen */
161 printf("ERROR: outstanding bytes negative!");
162 exit(0);
163 }
164 } else {
165 /* We assume the last 4 bytes will be delivered in one chunk */
166 if (length >= 4 && memcmp(data + length - 4, "\r\n\r\n", 4) == 0) {
167 /* We are upgraded so start sending the message for echoing */
168 http_socket->offset = us_socket_write(SSL, s, (char *) web_socket_request, web_socket_request_size, 0);
169
170 /* Server will echo back the same message minus 4 bytes for mask */
171 http_socket->outstanding_bytes = web_socket_request_size - 4;
172 http_socket->is_upgraded = 1;
173 }
174 }
175
176 return s;
177 }
178
on_http_socket_open(struct us_socket_t * s,int is_client,char * ip,int ip_length)179 struct us_socket_t *on_http_socket_open(struct us_socket_t *s, int is_client, char *ip, int ip_length) {
180 struct http_socket *http_socket = (struct http_socket *) us_socket_ext(SSL, s);
181
182 /* Reset offsets */
183 http_socket->offset = 0;
184
185 http_socket->is_upgraded = 0;
186
187 /* Send an upgrade request */
188 http_socket->upgrade_offset = us_socket_write(SSL, s, upgrade_request, upgrade_request_length, 0);
189 if (http_socket->upgrade_offset == upgrade_request_length) {
190 next_connection(s);
191 }
192
193 return s;
194 }
195
on_http_socket_timeout(struct us_socket_t * s)196 struct us_socket_t *on_http_socket_timeout(struct us_socket_t *s) {
197 /* Print current statistics */
198 printf("Msg/sec: %f\n", ((float)responses) / LIBUS_TIMEOUT_GRANULARITY);
199
200 responses = 0;
201 us_socket_timeout(SSL, s, LIBUS_TIMEOUT_GRANULARITY);
202
203 return s;
204 }
205
main(int argc,char ** argv)206 int main(int argc, char **argv) {
207
208 /* Parse host and port */
209 if (argc != 6 && argc != 7) {
210 printf("Usage: connections host port ssl deflate [size_mb]\n");
211 return 0;
212 }
213
214 port = atoi(argv[3]);
215 host = malloc(strlen(argv[2]) + 1);
216 memcpy(host, argv[2], strlen(argv[2]) + 1);
217 connections = atoi(argv[1]);
218 SSL = atoi(argv[4]);
219 if (atoi(argv[5])) {
220 /* Set up deflate */
221 web_socket_request = web_socket_request_deflate;
222 web_socket_request_size = sizeof(web_socket_request_deflate);
223
224 upgrade_request = request_deflate;
225 upgrade_request_length = sizeof(request_deflate) - 1;
226 } else {
227 /* Only if we are NOT using defalte can we support testing with 100mb for now */
228 if (argc == 7) {
229 int size_mb = atoi(argv[6]);
230 printf("Using message size of %d MB\n", size_mb);
231 /* Size has to be in MB since the minimal size is 64kb */
232 init_big_message(size_mb * 1024 * 1024);
233 }
234
235 web_socket_request = web_socket_request_text;
236 web_socket_request_size = web_socket_request_text_size;
237
238 upgrade_request = request_text;
239 upgrade_request_length = sizeof(request_text) - 1;
240 }
241
242 /* Create the event loop */
243 struct us_loop_t *loop = us_create_loop(0, on_wakeup, on_pre, on_post, 0);
244
245 /* Create a socket context for HTTP */
246 struct us_socket_context_options_t options = {};
247 struct us_socket_context_t *http_context = us_create_socket_context(SSL, loop, 0, options);
248
249 /* Set up event handlers */
250 us_socket_context_on_open(SSL, http_context, on_http_socket_open);
251 us_socket_context_on_data(SSL, http_context, on_http_socket_data);
252 us_socket_context_on_writable(SSL, http_context, on_http_socket_writable);
253 us_socket_context_on_close(SSL, http_context, on_http_socket_close);
254 us_socket_context_on_timeout(SSL, http_context, on_http_socket_timeout);
255 us_socket_context_on_end(SSL, http_context, on_http_socket_end);
256
257 /* Start making HTTP connections */
258 us_socket_context_connect(SSL, http_context, host, port, NULL, 0, sizeof(struct http_socket));
259
260 us_loop_run(loop);
261 }
262