1 /* This benchmark establishes _connections_ number of WebSocket
2    clients, then iteratively performs the following:
3 
4    1. Send one message for every client.
5    2. Wait for the quadratic (_connections_^2) amount of responses from the server.
6    3. Once received all expected bytes, repeat by going to step 1.
7 
8    Every 4 seconds we print the current average "iterations per second".
9    */
10 
11 #include <libusockets.h>
12 int SSL;
13 
14 #include <stdio.h>
15 #include <stdlib.h>
16 #include <string.h>
17 
18 unsigned char web_socket_request[26] = {130, 128 | 20, 1, 2, 3, 4};
19 
20 char request[] = "GET / HTTP/1.1\r\n"
21                  "Upgrade: websocket\r\n"
22                  "Connection: Upgrade\r\n"
23                  "Sec-WebSocket-Key: x3JJHMbDL1EzLkh9GBhXDw==\r\n"
24                  "Host: server.example.com\r\n"
25                  "Sec-WebSocket-Version: 13\r\n\r\n";
26 char *host;
27 int port;
28 int connections;
29 
30 int satisfied_sockets;
31 int iterations;
32 
33 struct http_socket {
34     /* How far we have streamed our websocket request */
35     int offset;
36 
37     /* How far we have streamed our upgrade request */
38     int upgrade_offset;
39 
40     /* Are we upgraded? */
41     int is_upgraded;
42 
43     /* Bytes received */
44     int bytes_received;
45 };
46 
47 /* We track upgraded websockets */
48 void **web_sockets;
49 int num_web_sockets;
50 
51 /* We don't need any of these */
noop(struct us_loop_t * loop)52 void noop(struct us_loop_t *loop) {
53 
54 }
55 
start_iteration()56 void start_iteration() {
57     for (int i = 0; i < num_web_sockets; i++) {
58         struct us_socket_t *s = (struct us_socket_t *) web_sockets[i];
59         struct http_socket *http_socket = (struct http_socket *) us_socket_ext(SSL, s);
60 
61         http_socket->offset = us_socket_write(SSL, s, (char *) web_socket_request, sizeof(web_socket_request), 0);
62     }
63 }
64 
next_connection(struct us_socket_t * s)65 void next_connection(struct us_socket_t *s) {
66     /* Add this connection to our array */
67     web_sockets[num_web_sockets++] = s;
68 
69     /* We could wait with this until properly upgraded */
70     if (--connections) {
71         us_socket_context_connect(SSL, us_socket_context(SSL, s), host, port, NULL, 0, sizeof(struct http_socket));
72     } else {
73         printf("Running benchmark now...\n");
74         start_iteration();
75 
76         us_socket_timeout(SSL, s, LIBUS_TIMEOUT_GRANULARITY);
77     }
78 }
79 
on_http_socket_writable(struct us_socket_t * s)80 struct us_socket_t *on_http_socket_writable(struct us_socket_t *s) {
81     struct http_socket *http_socket = (struct http_socket *) us_socket_ext(SSL, s);
82 
83     /* Are we still not upgraded yet? */
84     if (http_socket->upgrade_offset < sizeof(request) - 1) {
85         http_socket->upgrade_offset += us_socket_write(SSL, s, request + http_socket->upgrade_offset, sizeof(request) - 1 - http_socket->upgrade_offset, 0);
86     } else {
87         /* Stream whatever is remaining of the request */
88         http_socket->offset += us_socket_write(SSL, s, (char *) web_socket_request + http_socket->offset, sizeof(web_socket_request) - http_socket->offset, 0);
89     }
90 
91     return s;
92 }
93 
on_http_socket_close(struct us_socket_t * s,int code,void * reason)94 struct us_socket_t *on_http_socket_close(struct us_socket_t *s, int code, void *reason) {
95 
96     printf("Client was disconnected, exiting!\n");
97     exit(-1);
98 
99     return s;
100 }
101 
on_http_socket_end(struct us_socket_t * s)102 struct us_socket_t *on_http_socket_end(struct us_socket_t *s) {
103     return us_socket_close(SSL, s, 0, NULL);
104 }
105 
on_http_socket_data(struct us_socket_t * s,char * data,int length)106 struct us_socket_t *on_http_socket_data(struct us_socket_t *s, char *data, int length) {
107     /* Get socket extension and the socket's context's extension */
108     struct http_socket *http_socket = (struct http_socket *) us_socket_ext(SSL, s);
109 
110     /* Are we already upgraded? */
111     if (http_socket->is_upgraded) {
112         http_socket->bytes_received += length;
113 
114         if (http_socket->bytes_received == (sizeof(web_socket_request) - 4) * num_web_sockets) {
115             satisfied_sockets++;
116             http_socket->bytes_received = 0;
117 
118             if (satisfied_sockets == num_web_sockets) {
119                 iterations++;
120                 satisfied_sockets = 0;
121 
122                 start_iteration();
123             }
124         }
125     } else {
126         /* We assume the server is not sending anything immediately following upgrade and that we get rnrn in one chunk */
127         if (length >= 4 && data[length - 1] == '\n' && data[length - 2] == '\r' && data[length - 3] == '\n' && data[length - 4] == '\r') {
128             http_socket->is_upgraded = 1;
129             next_connection(s);
130         }
131     }
132 
133     return s;
134 }
135 
on_http_socket_open(struct us_socket_t * s,int is_client,char * ip,int ip_length)136 struct us_socket_t *on_http_socket_open(struct us_socket_t *s, int is_client, char *ip, int ip_length) {
137     struct http_socket *http_socket = (struct http_socket *) us_socket_ext(SSL, s);
138 
139     /* Reset offsets */
140     http_socket->offset = 0;
141     http_socket->is_upgraded = 0;
142     http_socket->bytes_received = 0;
143 
144     /* Send an upgrade request */
145     http_socket->upgrade_offset = us_socket_write(SSL, s, request, sizeof(request) - 1, 0);
146 
147     return s;
148 }
149 
on_http_socket_timeout(struct us_socket_t * s)150 struct us_socket_t *on_http_socket_timeout(struct us_socket_t *s) {
151     /* Print current statistics */
152     printf("Iterations/second (%d clients): %f\n", num_web_sockets, ((float)iterations) / LIBUS_TIMEOUT_GRANULARITY);
153 
154     iterations = 0;
155     us_socket_timeout(SSL, s, LIBUS_TIMEOUT_GRANULARITY);
156 
157     return s;
158 }
159 
main(int argc,char ** argv)160 int main(int argc, char **argv) {
161 
162     /* Parse host and port */
163     if (argc != 5) {
164         printf("Usage: connections host port ssl\n");
165         return 0;
166     }
167 
168     port = atoi(argv[3]);
169     host = malloc(strlen(argv[2]) + 1);
170     memcpy(host, argv[2], strlen(argv[2]) + 1);
171     connections = atoi(argv[1]);
172     SSL = atoi(argv[4]);
173 
174     /* Allocate room for every socket */
175     web_sockets = (void **) malloc(sizeof(void *) * connections);
176 
177     /* Create the event loop */
178     struct us_loop_t *loop = us_create_loop(0, noop, noop, noop, 0);
179 
180     /* Create a socket context for HTTP */
181     struct us_socket_context_options_t options = {};
182     struct us_socket_context_t *http_context = us_create_socket_context(SSL, loop, 0, options);
183 
184     /* Set up event handlers */
185     us_socket_context_on_open(SSL, http_context, on_http_socket_open);
186     us_socket_context_on_data(SSL, http_context, on_http_socket_data);
187     us_socket_context_on_writable(SSL, http_context, on_http_socket_writable);
188     us_socket_context_on_close(SSL, http_context, on_http_socket_close);
189     us_socket_context_on_timeout(SSL, http_context, on_http_socket_timeout);
190     us_socket_context_on_end(SSL, http_context, on_http_socket_end);
191 
192     /* Start making HTTP connections */
193     us_socket_context_connect(SSL, http_context, host, port, NULL, 0, sizeof(struct http_socket));
194 
195     us_loop_run(loop);
196 }
197