Small microoptimization in do_work().
[cubemap] / server.cpp
1 #include <assert.h>
2 #include <errno.h>
3 #include <pthread.h>
4 #include <stdio.h>
5 #include <stdlib.h>
6 #include <sys/epoll.h>
7 #include <sys/sendfile.h>
8 #include <sys/socket.h>
9 #include <sys/types.h>
10 #include <unistd.h>
11 #include <algorithm>
12 #include <map>
13 #include <string>
14 #include <utility>
15 #include <vector>
16
17 #include "markpool.h"
18 #include "mutexlock.h"
19 #include "parse.h"
20 #include "server.h"
21 #include "state.pb.h"
22 #include "stream.h"
23
24 using namespace std;
25
26 Server::Server()
27 {
28         pthread_mutex_init(&mutex, NULL);
29         pthread_mutex_init(&queued_data_mutex, NULL);
30
31         epoll_fd = epoll_create(1024);  // Size argument is ignored.
32         if (epoll_fd == -1) {
33                 perror("epoll_fd");
34                 exit(1);
35         }
36 }
37
38 Server::~Server()
39 {
40         int ret;
41         do {
42                 ret = close(epoll_fd);
43         } while (ret == -1 && errno == EINTR);
44
45         if (ret == -1) {
46                 perror("close(epoll_fd)");
47         }
48 }
49
50 vector<ClientStats> Server::get_client_stats() const
51 {
52         vector<ClientStats> ret;
53
54         MutexLock lock(&mutex);
55         for (map<int, Client>::const_iterator client_it = clients.begin();
56              client_it != clients.end();
57              ++client_it) {
58                 ret.push_back(client_it->second.get_stats());
59         }
60         return ret;
61 }
62
63 void Server::do_work()
64 {
65         for ( ;; ) {
66                 int nfds = epoll_wait(epoll_fd, events, EPOLL_MAX_EVENTS, EPOLL_TIMEOUT_MS);
67                 if (nfds == -1 && errno == EINTR) {
68                         if (should_stop) {
69                                 return;
70                         }
71                         continue;
72                 }
73                 if (nfds == -1) {
74                         perror("epoll_wait");
75                         exit(1);
76                 }
77
78                 MutexLock lock(&mutex);  // We release the mutex between iterations.
79         
80                 process_queued_data();
81
82                 for (int i = 0; i < nfds; ++i) {
83                         int fd = events[i].data.fd;
84                         assert(clients.count(fd) != 0);
85                         Client *client = &clients[fd];
86
87                         if (events[i].events & (EPOLLERR | EPOLLRDHUP | EPOLLHUP)) {
88                                 close_client(client);
89                                 continue;
90                         }
91
92                         process_client(client);
93                 }
94
95                 for (map<string, Stream *>::iterator stream_it = streams.begin();
96                      stream_it != streams.end();
97                      ++stream_it) {
98                         vector<Client *> to_process;
99                         swap(stream_it->second->to_process, to_process);
100                         for (size_t i = 0; i < to_process.size(); ++i) {
101                                 process_client(to_process[i]);
102                         }
103                 }
104
105                 if (should_stop) {
106                         return;
107                 }
108         }
109 }
110
111 CubemapStateProto Server::serialize()
112 {
113         // We don't serialize anything queued, so empty the queues.
114         process_queued_data();
115
116         CubemapStateProto serialized;
117         for (map<int, Client>::const_iterator client_it = clients.begin();
118              client_it != clients.end();
119              ++client_it) {
120                 serialized.add_clients()->MergeFrom(client_it->second.serialize());
121         }
122         for (map<string, Stream *>::const_iterator stream_it = streams.begin();
123              stream_it != streams.end();
124              ++stream_it) {
125                 serialized.add_streams()->MergeFrom(stream_it->second->serialize());
126         }
127         return serialized;
128 }
129
130 void Server::add_client_deferred(int sock)
131 {
132         MutexLock lock(&queued_data_mutex);
133         queued_add_clients.push_back(sock);
134 }
135
136 void Server::add_client(int sock)
137 {
138         clients.insert(make_pair(sock, Client(sock)));
139
140         // Start listening on data from this socket.
141         epoll_event ev;
142         ev.events = EPOLLIN | EPOLLET | EPOLLRDHUP;
143         ev.data.u64 = 0;  // Keep Valgrind happy.
144         ev.data.fd = sock;
145         if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sock, &ev) == -1) {
146                 perror("epoll_ctl(EPOLL_CTL_ADD)");
147                 exit(1);
148         }
149
150         process_client(&clients[sock]);
151 }
152
153 void Server::add_client_from_serialized(const ClientProto &client)
154 {
155         MutexLock lock(&mutex);
156         Stream *stream;
157         map<string, Stream *>::iterator stream_it = streams.find(client.stream_id());
158         if (stream_it == streams.end()) {
159                 stream = NULL;
160         } else {
161                 stream = stream_it->second;
162         }
163         clients.insert(make_pair(client.sock(), Client(client, stream)));
164         Client *client_ptr = &clients[client.sock()];
165
166         // Start listening on data from this socket.
167         epoll_event ev;
168         if (client.state() == Client::READING_REQUEST) {
169                 ev.events = EPOLLIN | EPOLLET | EPOLLRDHUP;
170         } else {
171                 // If we don't have more data for this client, we'll be putting it into
172                 // the sleeping array again soon.
173                 ev.events = EPOLLOUT | EPOLLET | EPOLLRDHUP;
174         }
175         ev.data.u64 = 0;  // Keep Valgrind happy.
176         ev.data.fd = client.sock();
177         if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client.sock(), &ev) == -1) {
178                 perror("epoll_ctl(EPOLL_CTL_ADD)");
179                 exit(1);
180         }
181
182         if (client_ptr->state == Client::SENDING_DATA && 
183             client_ptr->stream_pos == client_ptr->stream->bytes_received) {
184                 client_ptr->stream->put_client_to_sleep(client_ptr);
185         } else {
186                 process_client(client_ptr);
187         }
188 }
189
190 void Server::add_stream(const string &stream_id, size_t backlog_size)
191 {
192         MutexLock lock(&mutex);
193         streams.insert(make_pair(stream_id, new Stream(stream_id, backlog_size)));
194 }
195
196 void Server::add_stream_from_serialized(const StreamProto &stream)
197 {
198         MutexLock lock(&mutex);
199         streams.insert(make_pair(stream.stream_id(), new Stream(stream)));
200 }
201         
202 void Server::set_backlog_size(const std::string &stream_id, size_t new_size)
203 {
204         MutexLock lock(&mutex);
205         assert(streams.count(stream_id) != 0);
206         streams[stream_id]->set_backlog_size(new_size);
207 }
208         
209 void Server::set_header(const string &stream_id, const string &header)
210 {
211         MutexLock lock(&mutex);
212         find_stream(stream_id)->header = header;
213
214         // If there are clients we haven't sent anything to yet, we should give
215         // them the header, so push back into the SENDING_HEADER state.
216         for (map<int, Client>::iterator client_it = clients.begin();
217              client_it != clients.end();
218              ++client_it) {
219                 Client *client = &client_it->second;
220                 if (client->state == Client::SENDING_DATA &&
221                     client->stream_pos == 0) {
222                         construct_header(client);
223                 }
224         }
225 }
226         
227 void Server::set_mark_pool(const std::string &stream_id, MarkPool *mark_pool)
228 {
229         MutexLock lock(&mutex);
230         assert(clients.empty());
231         find_stream(stream_id)->mark_pool = mark_pool;
232 }
233
234 void Server::add_data_deferred(const string &stream_id, const char *data, size_t bytes)
235 {
236         MutexLock lock(&queued_data_mutex);
237         queued_data[stream_id].append(string(data, data + bytes));
238 }
239
240 // See the .h file for postconditions after this function.      
241 void Server::process_client(Client *client)
242 {
243         switch (client->state) {
244         case Client::READING_REQUEST: {
245 read_request_again:
246                 // Try to read more of the request.
247                 char buf[1024];
248                 int ret;
249                 do {
250                         ret = read(client->sock, buf, sizeof(buf));
251                 } while (ret == -1 && errno == EINTR);
252
253                 if (ret == -1 && errno == EAGAIN) {
254                         // No more data right now. Nothing to do.
255                         // This is postcondition #2.
256                         return;
257                 }
258                 if (ret == -1) {
259                         perror("read");
260                         close_client(client);
261                         return;
262                 }
263                 if (ret == 0) {
264                         // OK, the socket is closed.
265                         close_client(client);
266                         return;
267                 }
268
269                 RequestParseStatus status = wait_for_double_newline(&client->request, buf, ret);
270         
271                 switch (status) {
272                 case RP_OUT_OF_SPACE:
273                         fprintf(stderr, "WARNING: fd %d sent overlong request!\n", client->sock);
274                         close_client(client);
275                         return;
276                 case RP_NOT_FINISHED_YET:
277                         // OK, we don't have the entire header yet. Fine; we'll get it later.
278                         // See if there's more data for us.
279                         goto read_request_again;
280                 case RP_EXTRA_DATA:
281                         fprintf(stderr, "WARNING: fd %d had junk data after request!\n", client->sock);
282                         close_client(client);
283                         return;
284                 case RP_FINISHED:
285                         break;
286                 }
287
288                 assert(status == RP_FINISHED);
289
290                 int error_code = parse_request(client);
291                 if (error_code == 200) {
292                         construct_header(client);
293                 } else {
294                         construct_error(client, error_code);
295                 }
296
297                 // We've changed states, so fall through.
298                 assert(client->state == Client::SENDING_ERROR ||
299                        client->state == Client::SENDING_HEADER);
300         }
301         case Client::SENDING_ERROR:
302         case Client::SENDING_HEADER: {
303 sending_header_or_error_again:
304                 int ret;
305                 do {
306                         ret = write(client->sock,
307                                     client->header_or_error.data() + client->header_or_error_bytes_sent,
308                                     client->header_or_error.size() - client->header_or_error_bytes_sent);
309                 } while (ret == -1 && errno == EINTR);
310
311                 if (ret == -1 && errno == EAGAIN) {
312                         // We're out of socket space, so now we're at the “low edge” of epoll's
313                         // edge triggering. epoll will tell us when there is more room, so for now,
314                         // just return.
315                         // This is postcondition #4.
316                         return;
317                 }
318
319                 if (ret == -1) {
320                         // Error! Postcondition #1.
321                         perror("write");
322                         close_client(client);
323                         return;
324                 }
325                 
326                 client->header_or_error_bytes_sent += ret;
327                 assert(client->header_or_error_bytes_sent <= client->header_or_error.size());
328
329                 if (client->header_or_error_bytes_sent < client->header_or_error.size()) {
330                         // We haven't sent all yet. Fine; go another round.
331                         goto sending_header_or_error_again;
332                 }
333
334                 // We're done sending the header or error! Clear it to release some memory.
335                 client->header_or_error.clear();
336
337                 if (client->state == Client::SENDING_ERROR) {
338                         // We're done sending the error, so now close.  
339                         // This is postcondition #1.
340                         close_client(client);
341                         return;
342                 }
343
344                 // Start sending from the end. In other words, we won't send any of the backlog,
345                 // but we'll start sending immediately as we get data.
346                 // This is postcondition #3.
347                 client->state = Client::SENDING_DATA;
348                 client->stream_pos = client->stream->bytes_received;
349                 client->stream->put_client_to_sleep(client);
350                 return;
351         }
352         case Client::SENDING_DATA: {
353 sending_data_again:
354                 // See if there's some data we've lost. Ideally, we should drop to a block boundary,
355                 // but resync will be the mux's problem.
356                 Stream *stream = client->stream;
357                 size_t bytes_to_send = stream->bytes_received - client->stream_pos;
358                 if (bytes_to_send == 0) {
359                         return;
360                 }
361                 if (bytes_to_send > stream->backlog_size) {
362                         fprintf(stderr, "WARNING: fd %d lost %lld bytes, maybe too slow connection\n",
363                                 client->sock,
364                                 (long long int)(bytes_to_send - stream->backlog_size));
365                         client->stream_pos = stream->bytes_received - stream->backlog_size;
366                         client->bytes_lost += bytes_to_send - stream->backlog_size;
367                         ++client->num_loss_events;
368                         bytes_to_send = stream->backlog_size;
369                 }
370
371                 // See if we need to split across the circular buffer.
372                 bool more_data = false;
373                 if ((client->stream_pos % stream->backlog_size) + bytes_to_send > stream->backlog_size) {
374                         bytes_to_send = stream->backlog_size - (client->stream_pos % stream->backlog_size);
375                         more_data = true;
376                 }
377
378                 ssize_t ret;
379                 do {
380                         loff_t offset = client->stream_pos % stream->backlog_size;
381                         ret = sendfile(client->sock, stream->data_fd, &offset, bytes_to_send);
382                 } while (ret == -1 && errno == EINTR);
383
384                 if (ret == -1 && errno == EAGAIN) {
385                         // We're out of socket space, so return; epoll will wake us up
386                         // when there is more room.
387                         // This is postcondition #4.
388                         return;
389                 }
390                 if (ret == -1) {
391                         // Error, close; postcondition #1.
392                         perror("sendfile");
393                         close_client(client);
394                         return;
395                 }
396                 client->stream_pos += ret;
397                 client->bytes_sent += ret;
398
399                 if (client->stream_pos == stream->bytes_received) {
400                         // We don't have any more data for this client, so put it to sleep.
401                         // This is postcondition #3.
402                         stream->put_client_to_sleep(client);
403                 } else if (more_data && ret == bytes_to_send) {
404                         goto sending_data_again;
405                 }
406                 break;
407         }
408         default:
409                 assert(false);
410         }
411 }
412
413 int Server::parse_request(Client *client)
414 {
415         vector<string> lines = split_lines(client->request);
416         if (lines.empty()) {
417                 return 400;  // Bad request (empty).
418         }
419
420         vector<string> request_tokens = split_tokens(lines[0]);
421         if (request_tokens.size() < 2) {
422                 return 400;  // Bad request (empty).
423         }
424         if (request_tokens[0] != "GET") {
425                 return 400;  // Should maybe be 405 instead?
426         }
427         if (streams.count(request_tokens[1]) == 0) {
428                 return 404;  // Not found.
429         }
430
431         client->stream_id = request_tokens[1];
432         client->stream = find_stream(client->stream_id);
433         if (client->stream->mark_pool != NULL) {
434                 client->fwmark = client->stream->mark_pool->get_mark();
435         } else {
436                 client->fwmark = 0;  // No mark.
437         }
438         if (setsockopt(client->sock, SOL_SOCKET, SO_MARK, &client->fwmark, sizeof(client->fwmark)) == -1) {                          
439                 if (client->fwmark != 0) {
440                         perror("setsockopt(SO_MARK)");
441                 }
442         }
443         client->request.clear();
444
445         return 200;  // OK!
446 }
447
448 void Server::construct_header(Client *client)
449 {
450         client->header_or_error = find_stream(client->stream_id)->header;
451
452         // Switch states.
453         client->state = Client::SENDING_HEADER;
454
455         epoll_event ev;
456         ev.events = EPOLLOUT | EPOLLET | EPOLLRDHUP;
457         ev.data.u64 = 0;  // Keep Valgrind happy.
458         ev.data.fd = client->sock;
459
460         if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, client->sock, &ev) == -1) {
461                 perror("epoll_ctl(EPOLL_CTL_MOD)");
462                 exit(1);
463         }
464 }
465         
466 void Server::construct_error(Client *client, int error_code)
467 {
468         char error[256];
469         snprintf(error, 256, "HTTP/1.0 %d Error\r\nContent-type: text/plain\r\n\r\nSomething went wrong. Sorry.\r\n",
470                 error_code);
471         client->header_or_error = error;
472
473         // Switch states.
474         client->state = Client::SENDING_ERROR;
475
476         epoll_event ev;
477         ev.events = EPOLLOUT | EPOLLET | EPOLLRDHUP;
478         ev.data.u64 = 0;  // Keep Valgrind happy.
479         ev.data.fd = client->sock;
480
481         if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, client->sock, &ev) == -1) {
482                 perror("epoll_ctl(EPOLL_CTL_MOD)");
483                 exit(1);
484         }
485 }
486
487 template<class T>
488 void delete_from(vector<T> *v, T elem)
489 {
490         typename vector<T>::iterator new_end = remove(v->begin(), v->end(), elem);
491         v->erase(new_end, v->end());
492 }
493         
494 void Server::close_client(Client *client)
495 {
496         if (epoll_ctl(epoll_fd, EPOLL_CTL_DEL, client->sock, NULL) == -1) {
497                 perror("epoll_ctl(EPOLL_CTL_DEL)");
498                 exit(1);
499         }
500
501         // This client could be sleeping, so we'll need to fix that. (Argh, O(n).)
502         if (client->stream != NULL) {
503                 delete_from(&client->stream->sleeping_clients, client);
504                 delete_from(&client->stream->to_process, client);
505                 if (client->stream->mark_pool != NULL) {
506                         int fwmark = client->fwmark;
507                         client->stream->mark_pool->release_mark(fwmark);
508                 }
509         }
510
511         // Bye-bye!
512         int ret;
513         do {
514                 ret = close(client->sock);
515         } while (ret == -1 && errno == EINTR);
516
517         if (ret == -1) {
518                 perror("close");
519         }
520
521         clients.erase(client->sock);
522 }
523         
524 Stream *Server::find_stream(const string &stream_id)
525 {
526         map<string, Stream *>::iterator it = streams.find(stream_id);
527         assert(it != streams.end());
528         return it->second;
529 }
530
531 void Server::process_queued_data()
532 {
533         MutexLock lock(&queued_data_mutex);
534
535         for (size_t i = 0; i < queued_add_clients.size(); ++i) {
536                 add_client(queued_add_clients[i]);
537         }
538         queued_add_clients.clear();     
539         
540         for (map<string, string>::iterator queued_it = queued_data.begin();
541              queued_it != queued_data.end();
542              ++queued_it) {
543                 Stream *stream = find_stream(queued_it->first);
544                 stream->add_data(queued_it->second.data(), queued_it->second.size());
545                 stream->wake_up_all_clients();
546         }
547         queued_data.clear();
548 }