replay: handle abort in getRemoteFileSize (#23427)

commaai · Jan 6, 2022 · 86d7307 · 86d7307
1 parent e0338fd
commit 86d7307
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 12 deletions.
diff --git a/selfdrive/ui/replay/filereader.cc b/selfdrive/ui/replay/filereader.cc
@@ -39,7 +39,7 @@ std::string FileReader::download(const std::string &url, std::atomic<bool> *abor
     if (!result.empty()) {
       return result;
     }
-    if (i != max_retries_) {
+    if (i != max_retries_ && !(abort && *abort)) {
       std::cout << "download failed, retrying " << i + 1 << std::endl;
     }
   }

diff --git a/selfdrive/ui/replay/main.cc b/selfdrive/ui/replay/main.cc
@@ -11,12 +11,14 @@
 
 const QString DEMO_ROUTE = "4cf7a6ad03080c90|2021-09-29--13-46-36";
 struct termios oldt = {};
+Replay *replay = nullptr;
 
 void sigHandler(int s) {
   std::signal(s, SIG_DFL);
   if (oldt.c_lflag) {
     tcsetattr(STDIN_FILENO, TCSANOW, &oldt);
   }
+  replay->stop();
   qApp->quit();
 }
 
@@ -143,7 +145,7 @@ int main(int argc, char *argv[]) {
       replay_flags |= flag;
     }
   }
-  Replay *replay = new Replay(route, allow, block, nullptr, replay_flags, parser.value("data_dir"), &app);
+  replay = new Replay(route, allow, block, nullptr, replay_flags, parser.value("data_dir"), &app);
   if (!replay->load()) {
     return 0;
   }

diff --git a/selfdrive/ui/replay/replay.cc b/selfdrive/ui/replay/replay.cc
@@ -34,10 +34,15 @@ Replay::Replay(QString route, QStringList allow, QStringList block, SubMaster *s
   qRegisterMetaType<FindFlag>("FindFlag");
   connect(this, &Replay::seekTo, this, &Replay::doSeek);
   connect(this, &Replay::seekToFlag, this, &Replay::doSeekToFlag);
+  connect(this, &Replay::stop, this, &Replay::doStop);
   connect(this, &Replay::segmentChanged, this, &Replay::queueSegment);
 }
 
 Replay::~Replay() {
+  doStop();
+}
+
+void Replay::doStop() {
   if (!stream_thread_ && segments_.empty()) return;
 
   qDebug() << "shutdown: in progress...";

diff --git a/selfdrive/ui/replay/replay.h b/selfdrive/ui/replay/replay.h
@@ -47,9 +47,11 @@ class Replay : public QObject {
   void segmentChanged();
   void seekTo(int seconds, bool relative);
   void seekToFlag(FindFlag flag);
+  void stop();
 
 protected slots:
   void queueSegment();
+  void doStop();
   void doSeek(int seconds, bool relative);
   void doSeekToFlag(FindFlag flag);
   void segmentLoadFinished(bool sucess);

diff --git a/selfdrive/ui/replay/util.cc b/selfdrive/ui/replay/util.cc
@@ -69,22 +69,28 @@ std::string formattedDataSize(size_t size) {
 
 } // namespace
 
-size_t getRemoteFileSize(const std::string &url) {
+size_t getRemoteFileSize(const std::string &url, std::atomic<bool> *abort) {
   CURL *curl = curl_easy_init();
   if (!curl) return -1;
 
   curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
   curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, dumy_write_cb);
   curl_easy_setopt(curl, CURLOPT_HEADER, 1);
   curl_easy_setopt(curl, CURLOPT_NOBODY, 1);
-  CURLcode res = curl_easy_perform(curl);
-  double content_length = -1;
-  if (res == CURLE_OK) {
-    curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD, &content_length);
-  } else {
-    std::cout << "Download failed: error code: " << res << std::endl;
+
+  CURLM *cm = curl_multi_init();
+  curl_multi_add_handle(cm, curl);
+  int still_running = 1;
+  while (still_running > 0 && !(abort && *abort)) {
+    CURLMcode mc = curl_multi_perform(cm, &still_running);
+    if (!mc) curl_multi_wait(cm, nullptr, 0, 1000, nullptr);
   }
+
+  double content_length = -1;
+  curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD, &content_length);
+  curl_multi_remove_handle(cm, curl);
   curl_easy_cleanup(curl);
+  curl_multi_cleanup(cm);
   return content_length > 0 ? (size_t)content_length : 0;
 }
 
@@ -176,15 +182,15 @@ bool httpDownload(const std::string &url, T &buf, size_t chunk_size, size_t cont
 }
 
 std::string httpGet(const std::string &url, size_t chunk_size, std::atomic<bool> *abort) {
-  size_t size = getRemoteFileSize(url);
+  size_t size = getRemoteFileSize(url, abort);
   if (size == 0) return {};
 
   std::string result(size, '\0');
   return httpDownload(url, result, chunk_size, size, abort) ? result : "";
 }
 
 bool httpDownload(const std::string &url, const std::string &file, size_t chunk_size, std::atomic<bool> *abort) {
-  size_t size = getRemoteFileSize(url);
+  size_t size = getRemoteFileSize(url, abort);
   if (size == 0) return false;
 
   std::ofstream of(file, std::ios::binary | std::ios::out);

diff --git a/selfdrive/ui/replay/util.h b/selfdrive/ui/replay/util.h
@@ -9,6 +9,6 @@ std::string decompressBZ2(const std::string &in, std::atomic<bool> *abort = null
 std::string decompressBZ2(const std::byte *in, size_t in_size, std::atomic<bool> *abort = nullptr);
 void enableHttpLogging(bool enable);
 std::string getUrlWithoutQuery(const std::string &url);
-size_t getRemoteFileSize(const std::string &url);
+size_t getRemoteFileSize(const std::string &url, std::atomic<bool> *abort = nullptr);
 std::string httpGet(const std::string &url, size_t chunk_size = 0, std::atomic<bool> *abort = nullptr);
 bool httpDownload(const std::string &url, const std::string &file, size_t chunk_size = 0, std::atomic<bool> *abort = nullptr);