io: add io_fd_block() helper.
[ccan] / ccan / io / benchmarks / run-different-speed.c
1 /* Simulate a server with connections of different speeds.  We count
2  * how many connections complete in 10 seconds. */
3 #include <ccan/io/io.h>
4 #include <ccan/time/time.h>
5 #include <ccan/err/err.h>
6 #include <sys/types.h>
7 #include <sys/socket.h>
8 #include <sys/un.h>
9 #include <stdlib.h>
10 #include <unistd.h>
11 #include <fcntl.h>
12 #include <errno.h>
13 #include <stdio.h>
14 #include <signal.h>
15
16 #define REQUEST_SIZE 1024
17 #define REPLY_SIZE 10240
18 #define NUM_CONNS 500 /* per child */
19 #define NUM_CHILDREN 2
20
21 static unsigned int completed;
22
23 struct client {
24         char request_buffer[REQUEST_SIZE];
25         char reply_buffer[REPLY_SIZE];
26 };
27
28 static struct io_plan write_reply(struct io_conn *conn, struct client *client);
29 static struct io_plan read_request(struct io_conn *conn, struct client *client)
30 {
31         return io_read(client->request_buffer, REQUEST_SIZE,
32                        write_reply, client);
33 }
34
35 /* once we're done, loop again. */
36 static struct io_plan write_complete(struct io_conn *conn, struct client *client)
37 {
38         completed++;
39         return read_request(conn, client);
40 }
41
42 static struct io_plan write_reply(struct io_conn *conn, struct client *client)
43 {
44         return io_write(client->reply_buffer, REPLY_SIZE,
45                         write_complete, client);
46 }
47
48 /* This runs in the child. */
49 static void create_clients(struct sockaddr_un *addr, int waitfd)
50 {
51         struct client data;
52         int i, sock[NUM_CONNS], speed[NUM_CONNS], done[NUM_CONNS], count = 0;
53
54         for (i = 0; i < NUM_CONNS; i++) {
55                 /* Set speed. */
56                 speed[i] = (1 << (random() % 10));
57                 sock[i] = socket(AF_UNIX, SOCK_STREAM, 0);
58                 if (sock[i] < 0)
59                         err(1, "creating socket");
60                 if (connect(sock[i], (void *)addr, sizeof(*addr)) != 0)
61                         err(1, "connecting socket");
62                 /* Make nonblocking. */
63                 io_fd_block(sock[i], false);
64                 done[i] = 0;
65         }
66
67         read(waitfd, &i, 1);
68
69         for (;;) {
70                 for (i = 0; i < NUM_CONNS; i++) {
71                         int ret, bytes = speed[i];
72                         if (done[i] < REQUEST_SIZE) {
73                                 if (REQUEST_SIZE - done[i] < bytes)
74                                         bytes = REQUEST_SIZE - done[i];
75                                 ret = write(sock[i], data.request_buffer,
76                                             bytes);
77                                 if (ret > 0)
78                                         done[i] += ret;
79                                 else if (ret < 0 && errno != EAGAIN)
80                                         goto fail;
81                         } else {
82                                 if (REQUEST_SIZE + REPLY_SIZE - done[i] < bytes)
83                                         bytes = REQUEST_SIZE + REPLY_SIZE
84                                                 - done[i];
85                                 ret = read(sock[i], data.reply_buffer,
86                                             bytes);
87                                 if (ret > 0) {
88                                         done[i] += ret;
89                                         if (done[i] == REQUEST_SIZE + REPLY_SIZE) {
90                                                 count++;
91                                                 done[i] = 0;
92                                         }
93                                 } else if (ret < 0 && errno != EAGAIN)
94                                         goto fail;
95                         }
96                 }
97         }
98 fail:
99         printf("Child did %u\n", count);
100         exit(0);
101 }
102
103 static int timeout[2];
104 static void sigalarm(int sig)
105 {
106         write(timeout[1], "1", 1);
107 }
108
109 static struct io_plan do_timeout(struct io_conn *conn, char *buf)
110 {
111         return io_break(buf, io_idle());
112 }
113
114 int main(int argc, char *argv[])
115 {
116         struct client client;
117         unsigned int i, j;
118         struct sockaddr_un addr;
119         struct timespec start, end;
120         int fd, wake[2];
121         char buf;
122
123         addr.sun_family = AF_UNIX;
124         sprintf(addr.sun_path, "/tmp/run-different-speed.sock.%u", getpid());
125
126         if (pipe(wake) != 0 || pipe(timeout) != 0)
127                 err(1, "Creating pipes");
128
129         fd = socket(AF_UNIX, SOCK_STREAM, 0);
130         if (fd < 0)
131                 err(1, "Creating socket");
132
133         if (bind(fd, (void *)&addr, sizeof(addr)) != 0)
134                 err(1, "Binding to %s", addr.sun_path);
135
136         if (listen(fd, NUM_CONNS) != 0)
137                 err(1, "Listening on %s", addr.sun_path);
138
139         for (i = 0; i < NUM_CHILDREN; i++) {
140                 switch (fork()) {
141                 case -1:
142                         err(1, "forking");
143                 case 0:
144                         close(wake[1]);
145                         create_clients(&addr, wake[0]);
146                         break;
147                 }
148                 for (j = 0; j < NUM_CONNS; j++) {
149                         int ret = accept(fd, NULL, 0);
150                         if (ret < 0)
151                                 err(1, "Accepting fd");
152                         /* For efficiency, we share client structure */
153                         io_new_conn(ret,
154                                     io_read(client.request_buffer, REQUEST_SIZE,
155                                             write_reply, &client));
156                 }
157         }
158
159         io_new_conn(timeout[0], io_read(&buf, 1, do_timeout, &buf));
160
161         close(wake[0]);
162         for (i = 0; i < NUM_CHILDREN; i++)
163                 write(wake[1], "1", 1);
164
165         signal(SIGALRM, sigalarm);
166         alarm(10);
167         start = time_now();
168         io_loop();
169         end = time_now();
170         close(fd);
171
172         printf("%u connections complete (%u ns per conn)\n",
173                completed,
174                (int)time_to_nsec(time_divide(time_sub(end, start), completed)));
175         return 0;
176 }