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