git.sesse.net Git - cubemap/blob - main.cpp

   1 #include <assert.h>
   2 #include <errno.h>
   3 #include <getopt.h>
   4 #include <signal.h>
   5 #include <stddef.h>
   6 #include <stdio.h>
   7 #include <stdlib.h>
   8 #include <string.h>
   9 #include <sys/time.h>
  10 #include <sys/wait.h>
  11 #include <unistd.h>
  12 #include <map>
  13 #include <set>
  14 #include <string>
  15 #include <utility>
  16 #include <vector>
  17
  18 #include "acceptor.h"
  19 #include "config.h"
  20 #include "input.h"
  21 #include "log.h"
  22 #include "markpool.h"
  23 #include "serverpool.h"
  24 #include "state.pb.h"
  25 #include "stats.h"
  26 #include "util.h"
  27 #include "version.h"
  28
  29 using namespace std;
  30
  31 ServerPool *servers = NULL;
  32 volatile bool hupped = false;
  33 volatile bool stopped = false;
  34
  35 void hup(int signum)
  36 {
  37         hupped = true;
  38         if (signum == SIGINT) {
  39                 stopped = true;
  40         }
  41 }
  42
  43 CubemapStateProto collect_state(const timeval &serialize_start,
  44                                 const vector<Acceptor *> acceptors,
  45                                 const vector<Input *> inputs,
  46                                 ServerPool *servers)
  47 {
  48         CubemapStateProto state = servers->serialize();  // Fills streams() and clients().
  49         state.set_serialize_start_sec(serialize_start.tv_sec);
  50         state.set_serialize_start_usec(serialize_start.tv_usec);
  51
  52         for (size_t i = 0; i < acceptors.size(); ++i) {
  53                 state.add_acceptors()->MergeFrom(acceptors[i]->serialize());
  54         }
  55
  56         for (size_t i = 0; i < inputs.size(); ++i) {
  57                 state.add_inputs()->MergeFrom(inputs[i]->serialize());
  58         }
  59
  60         return state;
  61 }
  62
  63 // Find all port statements in the configuration file, and create acceptors for htem.
  64 vector<Acceptor *> create_acceptors(
  65         const Config &config,
  66         map<int, Acceptor *> *deserialized_acceptors)
  67 {
  68         vector<Acceptor *> acceptors;
  69         for (unsigned i = 0; i < config.acceptors.size(); ++i) {
  70                 const AcceptorConfig &acceptor_config = config.acceptors[i];
  71                 Acceptor *acceptor = NULL;
  72                 map<int, Acceptor *>::iterator deserialized_acceptor_it =
  73                         deserialized_acceptors->find(acceptor_config.port);
  74                 if (deserialized_acceptor_it != deserialized_acceptors->end()) {
  75                         acceptor = deserialized_acceptor_it->second;
  76                         deserialized_acceptors->erase(deserialized_acceptor_it);
  77                 } else {
  78                         int server_sock = create_server_socket(acceptor_config.port, TCP_SOCKET);
  79                         acceptor = new Acceptor(server_sock, acceptor_config.port);
  80                 }
  81                 acceptor->run();
  82                 acceptors.push_back(acceptor);
  83         }
  84
  85         // Close all acceptors that are no longer in the configuration file.
  86         for (map<int, Acceptor *>::iterator acceptor_it = deserialized_acceptors->begin();
  87              acceptor_it != deserialized_acceptors->end();
  88              ++acceptor_it) {
  89                 acceptor_it->second->close_socket();
  90                 delete acceptor_it->second;
  91         }
  92
  93         return acceptors;
  94 }
  95
  96 // Find all streams in the configuration file, and create inputs for them.
  97 vector<Input *> create_inputs(const Config &config,
  98                               map<string, Input *> *deserialized_inputs)
  99 {
 100         vector<Input *> inputs;
 101         for (unsigned i = 0; i < config.streams.size(); ++i) {
 102                 const StreamConfig &stream_config = config.streams[i];
 103                 if (stream_config.src.empty()) {
 104                         continue;
 105                 }
 106
 107                 string stream_id = stream_config.stream_id;
 108                 string src = stream_config.src;
 109
 110                 Input *input = NULL;
 111                 map<string, Input *>::iterator deserialized_input_it =
 112                         deserialized_inputs->find(stream_id);
 113                 if (deserialized_input_it != deserialized_inputs->end()) {
 114                         input = deserialized_input_it->second;
 115                         if (input->get_url() != src) {
 116                                 log(INFO, "Stream '%s' has changed URL from '%s' to '%s', restarting input.",
 117                                         stream_id.c_str(), input->get_url().c_str(), src.c_str());
 118                                 input->close_socket();
 119                                 delete input;
 120                                 input = NULL;
 121                         }
 122                         deserialized_inputs->erase(deserialized_input_it);
 123                 }
 124                 if (input == NULL) {
 125                         input = create_input(stream_id, src);
 126                         if (input == NULL) {
 127                                 log(ERROR, "did not understand URL '%s', clients will not get any data.",
 128                                         src.c_str());
 129                                 continue;
 130                         }
 131                 }
 132                 input->run();
 133                 inputs.push_back(input);
 134         }
 135         return inputs;
 136 }
 137
 138 void create_streams(const Config &config,
 139                     const set<string> &deserialized_stream_ids,
 140                     map<string, Input *> *deserialized_inputs)
 141 {
 142         vector<MarkPool *> mark_pools;  // FIXME: leak
 143         for (unsigned i = 0; i < config.mark_pools.size(); ++i) {
 144                 const MarkPoolConfig &mp_config = config.mark_pools[i];
 145                 mark_pools.push_back(new MarkPool(mp_config.from, mp_config.to));
 146         }
 147
 148         set<string> expecting_stream_ids = deserialized_stream_ids;
 149         for (unsigned i = 0; i < config.streams.size(); ++i) {
 150                 const StreamConfig &stream_config = config.streams[i];
 151                 if (deserialized_stream_ids.count(stream_config.stream_id) == 0) {
 152                         servers->add_stream(stream_config.stream_id, stream_config.backlog_size);
 153                 } else {
 154                         servers->set_backlog_size(stream_config.stream_id, stream_config.backlog_size);
 155                 }
 156                 expecting_stream_ids.erase(stream_config.stream_id);
 157
 158                 if (stream_config.mark_pool != -1) {
 159                         servers->set_mark_pool(stream_config.stream_id,
 160                                                mark_pools[stream_config.mark_pool]);
 161                 }
 162         }
 163
 164         // Warn about any servers we've lost.
 165         // TODO: Make an option (delete=yes?) to actually shut down streams.
 166         for (set<string>::const_iterator stream_it = expecting_stream_ids.begin();
 167              stream_it != expecting_stream_ids.end();
 168              ++stream_it) {
 169                 string stream_id = *stream_it;
 170                 log(WARNING, "stream '%s' disappeared from the configuration file. "
 171                              "It will not be deleted, but clients will not get any new inputs.",
 172                              stream_id.c_str());
 173                 if (deserialized_inputs->count(stream_id) != 0) {
 174                         delete (*deserialized_inputs)[stream_id];
 175                         deserialized_inputs->erase(stream_id);
 176                 }
 177         }
 178 }
 179
 180 void open_logs(const vector<LogConfig> &log_destinations)
 181 {
 182         for (size_t i = 0; i < log_destinations.size(); ++i) {
 183                 if (log_destinations[i].type == LogConfig::LOG_TYPE_FILE) {
 184                         add_log_destination_file(log_destinations[i].filename);
 185                 } else if (log_destinations[i].type == LogConfig::LOG_TYPE_CONSOLE) {
 186                         add_log_destination_console();
 187                 } else if (log_destinations[i].type == LogConfig::LOG_TYPE_SYSLOG) {
 188                         add_log_destination_syslog();
 189                 } else {
 190                         assert(false);
 191                 }
 192         }
 193         start_logging();
 194 }
 195
 196 bool dry_run_config(const std::string &argv0, const std::string &config_filename)
 197 {
 198         char *argv0_copy = strdup(argv0.c_str());
 199         char *config_filename_copy = strdup(config_filename.c_str());
 200
 201         pid_t pid = fork();
 202         switch (pid) {
 203         case -1:
 204                 log_perror("fork()");
 205                 free(argv0_copy);
 206                 free(config_filename_copy);
 207                 return false;
 208         case 0:
 209                 // Child.
 210                 execlp(argv0_copy, argv0_copy, "--test-config", config_filename_copy, NULL);
 211                 log_perror(argv0_copy);
 212                 _exit(1);
 213         default:
 214                 // Parent.
 215                 break;
 216         }
 217
 218         free(argv0_copy);
 219         free(config_filename_copy);
 220
 221         int status;
 222         pid_t err;
 223         do {
 224                 err = waitpid(pid, &status, 0);
 225         } while (err == -1 && errno == EINTR);
 226
 227         if (err == -1) {
 228                 log_perror("waitpid()");
 229                 return false;
 230         }
 231
 232         return (WIFEXITED(status) && WEXITSTATUS(status) == 0);
 233 }
 234
 235 int main(int argc, char **argv)
 236 {
 237         signal(SIGHUP, hup);
 238         signal(SIGINT, hup);
 239         signal(SIGPIPE, SIG_IGN);
 240
 241         // Parse options.
 242         int state_fd = -1;
 243         bool test_config = false;
 244         for ( ;; ) {
 245                 static const option long_options[] = {
 246                         { "state", required_argument, 0, 's' },
 247                         { "test-config", no_argument, 0, 't' },
 248                 };
 249                 int option_index = 0;
 250                 int c = getopt_long (argc, argv, "s:t", long_options, &option_index);
 251
 252                 if (c == -1) {
 253                         break;
 254                 }
 255                 switch (c) {
 256                 case 's':
 257                         state_fd = atoi(optarg);
 258                         break;
 259                 case 't':
 260                         test_config = true;
 261                         break;
 262                 default:
 263                         assert(false);
 264                 }
 265         }
 266
 267         string config_filename = "cubemap.config";
 268         if (optind < argc) {
 269                 config_filename = argv[optind++];
 270         }
 271
 272         Config config;
 273         if (!parse_config(config_filename, &config)) {
 274                 exit(1);
 275         }
 276         if (test_config) {
 277                 exit(0);
 278         }
 279
 280 start:
 281         // Open logs as soon as possible.
 282         open_logs(config.log_destinations);
 283
 284         log(INFO, "Cubemap " SERVER_VERSION " starting.");
 285         servers = new ServerPool(config.num_servers);
 286
 287         CubemapStateProto loaded_state;
 288         struct timeval serialize_start;
 289         set<string> deserialized_stream_ids;
 290         map<string, Input *> deserialized_inputs;
 291         map<int, Acceptor *> deserialized_acceptors;
 292         if (state_fd != -1) {
 293                 log(INFO, "Deserializing state from previous process...");
 294                 string serialized;
 295                 if (!read_tempfile(state_fd, &serialized)) {
 296                         exit(1);
 297                 }
 298                 if (!loaded_state.ParseFromString(serialized)) {
 299                         log(ERROR, "Failed deserialization of state.");
 300                         exit(1);
 301                 }
 302
 303                 serialize_start.tv_sec = loaded_state.serialize_start_sec();
 304                 serialize_start.tv_usec = loaded_state.serialize_start_usec();
 305
 306                 // Deserialize the streams.
 307                 for (int i = 0; i < loaded_state.streams_size(); ++i) {
 308                         servers->add_stream_from_serialized(loaded_state.streams(i));
 309                         deserialized_stream_ids.insert(loaded_state.streams(i).stream_id());
 310                 }
 311
 312                 // Deserialize the inputs. Note that we don't actually add them to any state yet.
 313                 for (int i = 0; i < loaded_state.inputs_size(); ++i) {
 314                         deserialized_inputs.insert(make_pair(
 315                                 loaded_state.inputs(i).stream_id(),
 316                                 create_input(loaded_state.inputs(i))));
 317                 }
 318
 319                 // Deserialize the acceptors.
 320                 for (int i = 0; i < loaded_state.acceptors_size(); ++i) {
 321                         deserialized_acceptors.insert(make_pair(
 322                                 loaded_state.acceptors(i).port(),
 323                                 new Acceptor(loaded_state.acceptors(i))));
 324                 }
 325
 326                 log(INFO, "Deserialization done.");
 327         }
 328
 329         // Find all streams in the configuration file, and create them.
 330         create_streams(config, deserialized_stream_ids, &deserialized_inputs);
 331
 332         vector<Acceptor *> acceptors = create_acceptors(config, &deserialized_acceptors);
 333         vector<Input *> inputs = create_inputs(config, &deserialized_inputs);
 334
 335         // All deserialized inputs should now have been taken care of, one way or the other.
 336         assert(deserialized_inputs.empty());
 337
 338         // Put back the existing clients. It doesn't matter which server we
 339         // allocate them to, so just do round-robin. However, we need to add
 340         // them after the mark pools have been set up.
 341         for (int i = 0; i < loaded_state.clients_size(); ++i) {
 342                 servers->add_client_from_serialized(loaded_state.clients(i));
 343         }
 344
 345         servers->run();
 346
 347         // Start writing statistics.
 348         StatsThread *stats_thread = NULL;
 349         if (!config.stats_file.empty()) {
 350                 stats_thread = new StatsThread(config.stats_file, config.stats_interval);
 351                 stats_thread->run();
 352         }
 353
 354         struct timeval server_start;
 355         gettimeofday(&server_start, NULL);
 356         if (state_fd != -1) {
 357                 // Measure time from we started deserializing (below) to now, when basically everything
 358                 // is up and running. This is, in other words, a conservative estimate of how long our
 359                 // “glitch” period was, not counting of course reconnects if the configuration changed.
 360                 double glitch_time = server_start.tv_sec - serialize_start.tv_sec +
 361                         1e-6 * (server_start.tv_usec - serialize_start.tv_usec);
 362                 log(INFO, "Re-exec happened in approx. %.0f ms.", glitch_time * 1000.0);
 363         }
 364
 365         while (!hupped) {
 366                 usleep(100000);
 367         }
 368
 369         // OK, we've been HUPed. Time to shut down everything, serialize, and re-exec.
 370         gettimeofday(&serialize_start, NULL);
 371
 372         if (stats_thread != NULL) {
 373                 stats_thread->stop();
 374         }
 375         for (size_t i = 0; i < acceptors.size(); ++i) {
 376                 acceptors[i]->stop();
 377         }
 378         for (size_t i = 0; i < inputs.size(); ++i) {
 379                 inputs[i]->stop();
 380         }
 381         servers->stop();
 382
 383         CubemapStateProto state;
 384         if (stopped) {
 385                 log(INFO, "Shutting down.");
 386         } else {
 387                 log(INFO, "Serializing state and re-execing...");
 388                 state = collect_state(
 389                         serialize_start, acceptors, inputs, servers);
 390                 string serialized;
 391                 state.SerializeToString(&serialized);
 392                 state_fd = make_tempfile(serialized);
 393                 if (state_fd == -1) {
 394                         exit(1);
 395                 }
 396         }
 397         delete servers;
 398         shut_down_logging();
 399
 400         if (stopped) {
 401                 exit(0);
 402         }
 403
 404         // OK, so the signal was SIGHUP. Check that the new config is okay, then exec the new binary.
 405         if (!dry_run_config(argv[0], config_filename)) {
 406                 open_logs(config.log_destinations);
 407                 log(ERROR, "%s --test-config failed. Restarting old version instead of new.", argv[0]);
 408                 hupped = false;
 409                 shut_down_logging();
 410                 goto start;
 411         }
 412
 413         char buf[16];
 414         sprintf(buf, "%d", state_fd);
 415
 416         for ( ;; ) {
 417                 execlp(argv[0], argv[0], config_filename.c_str(), "--state", buf, NULL);
 418                 open_logs(config.log_destinations);
 419                 log_perror("execlp");
 420                 log(ERROR, "re-exec of %s failed. Waiting 0.2 seconds and trying again...", argv[0]);
 421                 shut_down_logging();
 422                 usleep(200000);
 423         }
 424 }