Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
@namchuai
namchuai committed Sep 9, 2024
1 parent 7e98d06 commit 03822a4
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 112 deletions.
176 changes: 76 additions & 100 deletions engine/services/download_service.cc
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
#include <curl/curl.h>
#include <httplib.h>
#include <stdio.h>
#include <trantor/utils/Logger.h>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <thread>

Expand All @@ -15,28 +15,78 @@ void DownloadService::AddDownloadTask(const DownloadTask& task,
tasks.push_back(task);

for (const auto& item : task.items) {
// StartDownloadItem(task.id, item, callback);
try {
auto download_url{item.host + "/" + item.path};
auto file_size = GetFileSize(download_url);

std::cout << "file name: " << item.fileName << std::endl;
std::cout << "file size: " << file_size << std::endl;
} catch (const std::runtime_error& e) {
CTL_WRN(e.what());
}
Download(task.id, item, std::nullopt, callback);
}
}

uint64_t DownloadService::GetFileSize(const std::string& url) const {
CURL* curl;
curl = curl_easy_init();

if (!curl) {
// TODO: throw a custom exception
throw std::runtime_error("Failed to initialize curl");
}

curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
curl_easy_setopt(curl, CURLOPT_NOBODY, 1L);
curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
CURLcode res = curl_easy_perform(curl);

if (res != CURLE_OK) {
throw std::runtime_error("Failed to get file size");
}

curl_off_t content_length = 0;
res = curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD_T,
&content_length);
if (content_length > 0) {
std::cout << "content length: " << content_length << std::endl;
} else {
printf("Content-Length not provided by server.\n");
}
return content_length;
}

void DownloadService::AddAsyncDownloadTask(
const DownloadTask& task, std::optional<DownloadItemCb> callback) {
tasks.push_back(task);

for (const auto& item : task.items) {
// TODO: maybe apply std::async is better?
std::thread([this, task, &callback, item]() {
this->StartDownloadItem(task.id, item, callback);
this->Download(task.id, item, std::nullopt, callback);
}).detach();
}
}

size_t WriteCallback(void* ptr, size_t size, size_t nmemb, FILE* stream) {
size_t written = fwrite(ptr, size, nmemb, stream);
return written;
}

size_t ProgressCallback(void* ptr, curl_off_t totalToDownload,
curl_off_t nowDownloaded, curl_off_t totalUpload,
curl_off_t nowUploaded) {
if (totalToDownload > 0) {
double percentage = (nowDownloaded * 100.0) / totalToDownload;
std::cout << "Downloaded: " << percentage << "%\r" << std::flush;
}
return 0;
}

void DownloadService::Download(
const std::string& download_id, const DownloadItem& download_item,
std::optional<ProgressUpdateCallback> progress_update_callback,
std::optional<DownloadItemCb> download_success_callback) {
std::cout << "NamH download\n";

// TODO: check the local path inside download_item and create it
auto container_folder_path{file_manager_utils::GetContainerFolderPath(
Expand All @@ -46,6 +96,7 @@ void DownloadService::Download(
auto item_folder_path{container_folder_path /
std::filesystem::path(download_id)};
CTL_INF("itemFolderPath: " << item_folder_path.string());
// TODO: should move this out of this function
if (!std::filesystem::exists(item_folder_path)) {
CTL_INF("Creating " << item_folder_path.string());
std::filesystem::create_directory(item_folder_path);
Expand All @@ -56,53 +107,42 @@ void DownloadService::Download(
CTL_INF("Absolute file output: " << output_file_path.string());

CURL* curl;
std::ofstream out_file(output_file_path, std::ios::binary);
FILE* file;
CURLcode res;
curl = curl_easy_init();

size_t (*writeCallbackPtr)(void*, size_t, size_t, FILE*) = &WriteCallback;
size_t (*progressCallbackPtr)(void*, curl_off_t, curl_off_t, curl_off_t,
curl_off_t) = &ProgressCallback;

auto downloadUrl{download_item.host + "/" + download_item.path};
CTL_INF("download ulr" << downloadUrl.c_str());
curl = curl_easy_init();
if (!curl) {
// TODO: throw a custom exception
throw std::runtime_error("Failed to initialize curl");
}

if (!out_file.is_open()) {
file = fopen(output_file_path.c_str(), "wb");
if (!file) {
// TODO: throw a custom exception
throw std::runtime_error("Failed to open output file");
}

auto write_callback = [](void* ptr, size_t size, size_t nmemb,
void* stream) -> size_t {
std::ofstream* out_stream = static_cast<std::ofstream*>(stream);
out_stream->write(static_cast<char*>(ptr), size * nmemb);
return size * nmemb;
};

// TODO: only create when progress_update_callback is not null
auto progress_callback =
[](void* ptr, curl_off_t totalToDownload, curl_off_t nowDownloaded,
curl_off_t totalToUpload, curl_off_t nowUploaded) -> int {
if (totalToDownload > 0) {
double percentage = (nowDownloaded * 100.0) / totalToDownload;
std::cout << "Downloaded: " << percentage << "%\r" << std::flush;
}
return 0;
};

curl_easy_setopt(curl, CURLOPT_URL, download_item.path.c_str());
curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_callback);
curl_easy_setopt(curl, CURLOPT_WRITEDATA, &out_file);
// TODO: handle progress callback here
curl_easy_setopt(curl, CURLOPT_XFERINFOFUNCTION, progress_callback);
curl_easy_setopt(curl, CURLOPT_URL, downloadUrl.c_str());
curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, writeCallbackPtr);
curl_easy_setopt(curl, CURLOPT_WRITEDATA, file);
curl_easy_setopt(curl, CURLOPT_XFERINFOFUNCTION, progressCallbackPtr);
curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
// allow redirect
curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);

res = curl_easy_perform(curl);

if (res != CURLE_OK) {
// TODO: better error message
throw std::runtime_error("Failed to download file");
fprintf(stderr, "curl_easy_perform() failed: %s\n",
curl_easy_strerror(res));
}

fclose(file);
curl_easy_cleanup(curl);
out_file.close();
if (download_success_callback.has_value()) {
// TODO: this should not be here
auto need_parse_gguf =
Expand All @@ -111,67 +151,3 @@ void DownloadService::Download(
need_parse_gguf);
}
}

void DownloadService::StartDownloadItem(
const std::string& downloadId, const DownloadItem& item,
std::optional<DownloadItemCb> callback) {
CTL_INF("Downloading item: " << downloadId);

auto containerFolderPath{file_manager_utils::GetContainerFolderPath(
file_manager_utils::DownloadTypeToString(item.type))};
CTL_INF("Container folder path: " << containerFolderPath.string() << "\n");

auto itemFolderPath{containerFolderPath / std::filesystem::path(downloadId)};
CTL_INF("itemFolderPath: " << itemFolderPath.string());
if (!std::filesystem::exists(itemFolderPath)) {
CTL_INF("Creating " << itemFolderPath.string());
std::filesystem::create_directory(itemFolderPath);
}

auto outputFilePath{itemFolderPath / std::filesystem::path(item.fileName)};
CTL_INF("Absolute file output: " << outputFilePath.string());

uint64_t last = 0;
uint64_t tot = 0;
std::ofstream outputFile(outputFilePath, std::ios::binary);

auto downloadUrl{item.host + "/" + item.path};
CLI_LOG("Downloading url: " << downloadUrl);

httplib::Client client(item.host);

client.set_follow_location(true);
client.Get(
downloadUrl,
[](const httplib::Response& res) {
if (res.status != httplib::StatusCode::OK_200) {
LOG_ERROR << "HTTP error: " << res.reason;
return false;
}
return true;
},
[&](const char* data, size_t data_length) {
tot += data_length;
outputFile.write(data, data_length);
return true;
},
[&item, &last, &outputFile, &callback, outputFilePath, this](
uint64_t current, uint64_t total) {
if (current - last > kUpdateProgressThreshold) {
last = current;
CLI_LOG("Downloading: " << current << " / " << total);
}
if (current == total) {
outputFile.flush();
CLI_LOG("Done download: " << static_cast<double>(total) / 1024 / 1024
<< " MiB");
if (callback.has_value()) {
auto need_parse_gguf =
item.path.find("cortexso") == std::string::npos;
callback.value()(outputFilePath.string(), need_parse_gguf);
}
return false;
}
return true;
});
}
14 changes: 2 additions & 12 deletions engine/services/download_service.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ enum class DownloadStatus {
struct DownloadItem {
std::string id;

// not needed anymore if we using libcurl to download
std::string host;

std::string fileName;
Expand Down Expand Up @@ -55,17 +56,9 @@ class DownloadService {
const DownloadTask& task,
std::optional<DownloadItemCb> callback = std::nullopt);

// TODO: [NamH] implement the following methods
// void removeTask(const std::string &id);
// void registerCallback
// setup folder path at runtime
// register action after downloaded
uint64_t GetFileSize(const std::string& url) const;

private:
void StartDownloadItem(const std::string& downloadId,
const DownloadItem& item,
std::optional<DownloadItemCb> callback = std::nullopt);

void Download(
const std::string& download_id, const DownloadItem& download_item,
std::optional<ProgressUpdateCallback> progress_update_callback =
Expand All @@ -74,7 +67,4 @@ class DownloadService {

// store tasks so we can abort it later
std::vector<DownloadTask> tasks;

// TODO: remove below
const int kUpdateProgressThreshold = 100000000;
};

0 comments on commit 03822a4

Please sign in to comment.