]> git.sesse.net Git - cubemap/blob - server.cpp
9517b5480e5f55ce6c5decee9ab6ad726edf3167
[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                         bytes_to_send = stream->backlog_size;
367                 }
368
369                 // See if we need to split across the circular buffer.
370                 bool more_data = false;
371                 if ((client->stream_pos % stream->backlog_size) + bytes_to_send > stream->backlog_size) {
372                         bytes_to_send = stream->backlog_size - (client->stream_pos % stream->backlog_size);
373                         more_data = true;
374                 }
375
376                 ssize_t ret;
377                 do {
378                         loff_t offset = client->stream_pos % stream->backlog_size;
379                         ret = sendfile(client->sock, stream->data_fd, &offset, bytes_to_send);
380                 } while (ret == -1 && errno == EINTR);
381
382                 if (ret == -1 && errno == EAGAIN) {
383                         // We're out of socket space, so return; epoll will wake us up
384                         // when there is more room.
385                         // This is postcondition #4.
386                         return;
387                 }
388                 if (ret == -1) {
389                         // Error, close; postcondition #1.
390                         perror("sendfile");
391                         close_client(client);
392                         return;
393                 }
394                 client->stream_pos += ret;
395
396                 if (client->stream_pos == stream->bytes_received) {
397                         // We don't have any more data for this client, so put it to sleep.
398                         // This is postcondition #3.
399                         stream->put_client_to_sleep(client);
400                 } else if (more_data) {
401                         goto sending_data_again;
402                 }
403                 break;
404         }
405         default:
406                 assert(false);
407         }
408 }
409
410 int Server::parse_request(Client *client)
411 {
412         vector<string> lines = split_lines(client->request);
413         if (lines.empty()) {
414                 return 400;  // Bad request (empty).
415         }
416
417         vector<string> request_tokens = split_tokens(lines[0]);
418         if (request_tokens.size() < 2) {
419                 return 400;  // Bad request (empty).
420         }
421         if (request_tokens[0] != "GET") {
422                 return 400;  // Should maybe be 405 instead?
423         }
424         if (streams.count(request_tokens[1]) == 0) {
425                 return 404;  // Not found.
426         }
427
428         client->stream_id = request_tokens[1];
429         client->stream = find_stream(client->stream_id);
430         if (client->stream->mark_pool != NULL) {
431                 client->fwmark = client->stream->mark_pool->get_mark();
432         } else {
433                 client->fwmark = 0;  // No mark.
434         }
435         if (setsockopt(client->sock, SOL_SOCKET, SO_MARK, &client->fwmark, sizeof(client->fwmark)) == -1) {                          
436                 if (client->fwmark != 0) {
437                         perror("setsockopt(SO_MARK)");
438                 }
439         }
440         client->request.clear();
441
442         return 200;  // OK!
443 }
444
445 void Server::construct_header(Client *client)
446 {
447         client->header_or_error = find_stream(client->stream_id)->header;
448
449         // Switch states.
450         client->state = Client::SENDING_HEADER;
451
452         epoll_event ev;
453         ev.events = EPOLLOUT | EPOLLET | EPOLLRDHUP;
454         ev.data.u64 = 0;  // Keep Valgrind happy.
455         ev.data.fd = client->sock;
456
457         if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, client->sock, &ev) == -1) {
458                 perror("epoll_ctl(EPOLL_CTL_MOD)");
459                 exit(1);
460         }
461 }
462         
463 void Server::construct_error(Client *client, int error_code)
464 {
465         char error[256];
466         snprintf(error, 256, "HTTP/1.0 %d Error\r\nContent-type: text/plain\r\n\r\nSomething went wrong. Sorry.\r\n",
467                 error_code);
468         client->header_or_error = error;
469
470         // Switch states.
471         client->state = Client::SENDING_ERROR;
472
473         epoll_event ev;
474         ev.events = EPOLLOUT | EPOLLET | EPOLLRDHUP;
475         ev.data.u64 = 0;  // Keep Valgrind happy.
476         ev.data.fd = client->sock;
477
478         if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, client->sock, &ev) == -1) {
479                 perror("epoll_ctl(EPOLL_CTL_MOD)");
480                 exit(1);
481         }
482 }
483
484 template<class T>
485 void delete_from(vector<T> *v, T elem)
486 {
487         typename vector<T>::iterator new_end = remove(v->begin(), v->end(), elem);
488         v->erase(new_end, v->end());
489 }
490         
491 void Server::close_client(Client *client)
492 {
493         if (epoll_ctl(epoll_fd, EPOLL_CTL_DEL, client->sock, NULL) == -1) {
494                 perror("epoll_ctl(EPOLL_CTL_DEL)");
495                 exit(1);
496         }
497
498         // This client could be sleeping, so we'll need to fix that. (Argh, O(n).)
499         if (client->stream != NULL) {
500                 delete_from(&client->stream->sleeping_clients, client);
501                 delete_from(&client->stream->to_process, client);
502                 if (client->stream->mark_pool != NULL) {
503                         int fwmark = client->fwmark;
504                         client->stream->mark_pool->release_mark(fwmark);
505                 }
506         }
507
508         // Bye-bye!
509         int ret;
510         do {
511                 ret = close(client->sock);
512         } while (ret == -1 && errno == EINTR);
513
514         if (ret == -1) {
515                 perror("close");
516         }
517
518         clients.erase(client->sock);
519 }
520         
521 Stream *Server::find_stream(const string &stream_id)
522 {
523         map<string, Stream *>::iterator it = streams.find(stream_id);
524         assert(it != streams.end());
525         return it->second;
526 }
527
528 void Server::process_queued_data()
529 {
530         MutexLock lock(&queued_data_mutex);
531
532         for (size_t i = 0; i < queued_add_clients.size(); ++i) {
533                 add_client(queued_add_clients[i]);
534         }
535         queued_add_clients.clear();     
536         
537         for (map<string, string>::iterator queued_it = queued_data.begin();
538              queued_it != queued_data.end();
539              ++queued_it) {
540                 Stream *stream = find_stream(queued_it->first);
541                 stream->add_data(queued_it->second.data(), queued_it->second.size());
542                 stream->wake_up_all_clients();
543         }
544         queued_data.clear();
545 }