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

   1 #include <assert.h>
   2 #include <errno.h>
   3 #include <limits.h>
   4 #include <netinet/in.h>
   5 #include <stdlib.h>
   6 #include <string.h>
   7 #include <sys/types.h>
   8 #include <string>
   9 #include <vector>
  10
  11 #include "log.h"
  12 #include "metacube.h"
  13 #include "state.pb.h"
  14 #include "stream.h"
  15 #include "util.h"
  16
  17 using namespace std;
  18
  19 Stream::Stream(const string &url, size_t backlog_size, Encoding encoding)
  20         : url(url),
  21           encoding(encoding),
  22           data_fd(make_tempfile("")),
  23           backlog_size(backlog_size),
  24           bytes_received(0),
  25           mark_pool(NULL)
  26 {
  27         if (data_fd == -1) {
  28                 exit(1);
  29         }
  30 }
  31
  32 Stream::~Stream()
  33 {
  34         if (data_fd != -1) {
  35                 safe_close(data_fd);
  36         }
  37 }
  38
  39 Stream::Stream(const StreamProto &serialized, int data_fd)
  40         : url(serialized.url()),
  41           http_header(serialized.http_header()),
  42           stream_header(serialized.stream_header()),
  43           encoding(Stream::STREAM_ENCODING_RAW),  // Will be changed later.
  44           data_fd(data_fd),
  45           backlog_size(serialized.backlog_size()),
  46           bytes_received(serialized.bytes_received()),
  47           mark_pool(NULL)
  48 {
  49         if (data_fd == -1) {
  50                 exit(1);
  51         }
  52
  53         // Split old-style headers into HTTP and video headers.
  54         if (!serialized.header().empty()) {
  55                 string header = serialized.header();
  56                 size_t split = header.find("\r\n\r\n");
  57                 if (split == string::npos) {
  58                         http_header = header;
  59                         stream_header = "";
  60                 } else {
  61                         http_header = header.substr(0, split + 2);  // Split off the second \r\n.
  62                         stream_header = header.substr(split, string::npos);
  63                 }
  64         }
  65 }
  66
  67 StreamProto Stream::serialize()
  68 {
  69         StreamProto serialized;
  70         serialized.set_http_header(http_header);
  71         serialized.set_stream_header(stream_header);
  72         serialized.add_data_fds(data_fd);
  73         serialized.set_backlog_size(backlog_size);
  74         serialized.set_bytes_received(bytes_received);
  75         serialized.set_url(url);
  76         data_fd = -1;
  77         return serialized;
  78 }
  79
  80 void Stream::set_backlog_size(size_t new_size)
  81 {
  82         if (backlog_size == new_size) {
  83                 return;
  84         }
  85
  86         string existing_data;
  87         if (!read_tempfile_and_close(data_fd, &existing_data)) {
  88                 exit(1);
  89         }
  90
  91         // Unwrap the data so it's no longer circular.
  92         if (bytes_received <= backlog_size) {
  93                 existing_data.resize(bytes_received);
  94         } else {
  95                 size_t pos = bytes_received % backlog_size;
  96                 existing_data = existing_data.substr(pos, string::npos) +
  97                         existing_data.substr(0, pos);
  98         }
  99
 100         // See if we need to discard data.
 101         if (new_size < existing_data.size()) {
 102                 size_t to_discard = existing_data.size() - new_size;
 103                 existing_data = existing_data.substr(to_discard, string::npos);
 104         }
 105
 106         // Create a new, empty data file.
 107         data_fd = make_tempfile("");
 108         if (data_fd == -1) {
 109                 exit(1);
 110         }
 111         backlog_size = new_size;
 112
 113         // Now cheat a bit by rewinding, and adding all the old data back.
 114         bytes_received -= existing_data.size();
 115         iovec iov;
 116         iov.iov_base = const_cast<char *>(existing_data.data());
 117         iov.iov_len = existing_data.size();
 118
 119         vector<iovec> iovs;
 120         iovs.push_back(iov);
 121         add_data_raw(iovs);
 122 }
 123
 124 void Stream::put_client_to_sleep(Client *client)
 125 {
 126         sleeping_clients.push_back(client);
 127 }
 128
 129 // Return a new set of iovecs that contains only the first <bytes_wanted> bytes of <data>.
 130 vector<iovec> collect_iovecs(const vector<iovec> &data, size_t bytes_wanted)
 131 {
 132         vector<iovec> ret;
 133         size_t max_iovecs = std::min<size_t>(data.size(), IOV_MAX);
 134         for (size_t i = 0; i < max_iovecs && bytes_wanted > 0; ++i) {
 135                 if (data[i].iov_len <= bytes_wanted) {
 136                         // Consume the entire iovec.
 137                         ret.push_back(data[i]);
 138                         bytes_wanted -= data[i].iov_len;
 139                 } else {
 140                         // Take only parts of this iovec.
 141                         iovec iov;
 142                         iov.iov_base = data[i].iov_base;
 143                         iov.iov_len = bytes_wanted;
 144                         ret.push_back(iov);
 145                         bytes_wanted = 0;
 146                 }
 147         }
 148         return ret;
 149 }
 150
 151 // Return a new set of iovecs that contains all of <data> except the first <bytes_wanted> bytes.
 152 vector<iovec> remove_iovecs(const vector<iovec> &data, size_t bytes_wanted)
 153 {
 154         vector<iovec> ret;
 155         size_t i;
 156         for (i = 0; i < data.size() && bytes_wanted > 0; ++i) {
 157                 if (data[i].iov_len <= bytes_wanted) {
 158                         // Consume the entire iovec.
 159                         bytes_wanted -= data[i].iov_len;
 160                 } else {
 161                         // Take only parts of this iovec.
 162                         iovec iov;
 163                         iov.iov_base = reinterpret_cast<char *>(data[i].iov_base) + bytes_wanted;
 164                         iov.iov_len = data[i].iov_len - bytes_wanted;
 165                         ret.push_back(iov);
 166                         bytes_wanted = 0;
 167                 }
 168         }
 169
 170         // Add the rest of the iovecs unchanged.
 171         ret.insert(ret.end(), data.begin() + i, data.end());
 172         return ret;
 173 }
 174
 175 void Stream::add_data_raw(const vector<iovec> &orig_data)
 176 {
 177         vector<iovec> data = orig_data;
 178         while (!data.empty()) {
 179                 size_t pos = bytes_received % backlog_size;
 180
 181                 // Collect as many iovecs as we can before we hit the point
 182                 // where the circular buffer wraps around.
 183                 vector<iovec> to_write = collect_iovecs(data, backlog_size - pos);
 184                 ssize_t ret;
 185                 do {
 186                         ret = pwritev(data_fd, to_write.data(), to_write.size(), pos);
 187                 } while (ret == -1 && errno == EINTR);
 188
 189                 if (ret == -1) {
 190                         log_perror("pwritev");
 191                         // Dazed and confused, but trying to continue...
 192                         return;
 193                 }
 194                 bytes_received += ret;
 195
 196                 // Remove the data that was actually written from the set of iovecs.
 197                 data = remove_iovecs(data, ret);
 198         }
 199 }
 200
 201 void Stream::add_data_deferred(const char *data, size_t bytes)
 202 {
 203         if (encoding == Stream::STREAM_ENCODING_METACUBE) {
 204                 // Add a Metacube block header before the data.
 205                 metacube_block_header hdr;
 206                 memcpy(hdr.sync, METACUBE_SYNC, sizeof(hdr.sync));
 207                 hdr.size = htonl(bytes);
 208                 hdr.flags = htonl(0);
 209
 210                 iovec iov;
 211                 iov.iov_base = new char[bytes + sizeof(hdr)];
 212                 iov.iov_len = bytes + sizeof(hdr);
 213
 214                 memcpy(iov.iov_base, &hdr, sizeof(hdr));
 215                 memcpy(reinterpret_cast<char *>(iov.iov_base) + sizeof(hdr), data, bytes);
 216
 217                 queued_data.push_back(iov);
 218         } else if (encoding == Stream::STREAM_ENCODING_RAW) {
 219                 // Just add the data itself.
 220                 iovec iov;
 221                 iov.iov_base = new char[bytes];
 222                 memcpy(iov.iov_base, data, bytes);
 223                 iov.iov_len = bytes;
 224
 225                 queued_data.push_back(iov);
 226         } else {
 227                 assert(false);
 228         }
 229 }
 230
 231 void Stream::process_queued_data()
 232 {
 233         if (queued_data.empty()) {
 234                 return;
 235         }
 236
 237         add_data_raw(queued_data);
 238         for (size_t i = 0; i < queued_data.size(); ++i) {
 239                 char *data = reinterpret_cast<char *>(queued_data[i].iov_base);
 240                 delete[] data;
 241         }
 242         queued_data.clear();
 243
 244         // We have more data, so wake up all clients.
 245         if (to_process.empty()) {
 246                 swap(sleeping_clients, to_process);
 247         } else {
 248                 to_process.insert(to_process.end(), sleeping_clients.begin(), sleeping_clients.end());
 249                 sleeping_clients.clear();
 250         }
 251 }