C++多线程下载工具代码及制作

#include <iostream>
#include <fstream>
#include <vector>
#include <thread>
#include <mutex>
#include <atomic>
#include <chrono>
#include <curl/curl.h>

using namespace std;

// 全局变量
mutex g_mutex;  // 用于线程同步
atomic<size_t> g_downloaded(0);  // 已下载字节数
size_t g_file_size = 0;  // 文件总大小

// 写入回调函数
size_t write_data(void* ptr, size_t size, size_t nmemb, void* data) {
    ofstream* ofs = static_cast<ofstream*>(data);
    size_t total_size = size * nmemb;
    ofs->write(static_cast<char*>(ptr), total_size);
    
    // 更新下载进度
    g_downloaded += total_size;
    
    return total_size;
}

// 下载单个分块
void download_chunk(const string& url, const string& temp_file, size_t start, size_t end) {
    CURL* curl;
    CURLcode res;
    ofstream ofs(temp_file, ios::binary);
    
    if (!ofs) {
        lock_guard<mutex> lock(g_mutex);
        cerr << "Failed to open file: " << temp_file << endl;
        return;
    }

    curl_global_init(CURL_GLOBAL_DEFAULT);
    curl = curl_easy_init();
    
    if (curl) {
        string range = to_string(start) + "-" + to_string(end);
        
        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
        curl_easy_setopt(curl, CURLOPT_RANGE, range.c_str());
        curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_data);
        curl_easy_setopt(curl, CURLOPT_WRITEDATA, &ofs);
        curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
        curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 1L);
        
        res = curl_easy_perform(curl);
        
        if (res != CURLE_OK) {
            lock_guard<mutex> lock(g_mutex);
            cerr << "Download failed: " << curl_easy_strerror(res) << endl;
        }
        
        curl_easy_cleanup(curl);
    }
    
    curl_global_cleanup();
}

// 获取文件大小
bool get_file_size(const string& url, size_t& file_size) {
    CURL* curl;
    curl_global_init(CURL_GLOBAL_DEFAULT);
    curl = curl_easy_init();
    
    if (curl) {
        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
        curl_easy_setopt(curl, CURLOPT_NOBODY, 1L);  // HEAD请求
        curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
        
        CURLcode res = curl_easy_perform(curl);
        
        if (res == CURLE_OK) {
            curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD, &file_size);
            curl_easy_cleanup(curl);
            curl_global_cleanup();
            return true;
        }
        
        curl_easy_cleanup(curl);
    }
    
    curl_global_cleanup();
    return false;
}

// 显示下载进度
void show_progress() {
    while (g_downloaded < g_file_size) {
        double progress = static_cast<double>(g_downloaded) / g_file_size * 100;
        double speed = g_downloaded / (1024.0 * 1024.0);  // MB
        
        cout << "\rProgress: " << fixed << setprecision(2) << progress << "% ";
        cout << "Downloaded: " << speed << " MB";
        cout.flush();
        
        this_thread::sleep_for(chrono::milliseconds(500));
    }
    
    cout << "\rDownload completed: 100.00%" << endl;
}

// 合并分块文件
void merge_files(const string& output_file, int thread_count) {
    ofstream ofs(output_file, ios::binary);
    
    for (int i = 0; i < thread_count; ++i) {
        string part_file = output_file + ".part" + to_string(i);
        ifstream ifs(part_file, ios::binary);
        
        if (ifs) {
            ofs << ifs.rdbuf();
            ifs.close();
            remove(part_file.c_str());
        }
    }
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        cout << "Usage: " << argv[0] << " <URL> <output_file> [thread_count]" << endl;
        return 1;
    }
    
    string url = argv[1];
    string output_file = argv[2];
    int thread_count = argc > 3 ? stoi(argv[3](@ref) : 4;
    
    // 获取文件大小
    if (!get_file_size(url, g_file_size)) {
        cerr << "Failed to get file size from URL" << endl;
        return 1;
    }
    
    cout << "File size: " << g_file_size / (1024 * 1024) << " MB" << endl;
    cout << "Downloading using " << thread_count << " threads..." << endl;
    
    // 计算分块大小
    size_t chunk_size = g_file_size / thread_count;
    vector<thread> threads;
    
    // 启动进度显示线程
    thread progress_thread(show_progress);
    
    // 启动下载线程
    for (int i = 0; i < thread_count; ++i) {
        size_t start = i * chunk_size;
        size_t end = (i == thread_count - 1) ? g_file_size - 1 : start + chunk_size - 1;
        string temp_file = output_file + ".part" + to_string(i);
        
        threads.emplace_back(download_chunk, url, temp_file, start, end);
    }
    
    // 等待所有线程完成
    for (auto& t : threads) {
        t.join();
    }
    
    // 等待进度线程结束
    progress_thread.join();
    
    // 合并文件
    merge_files(output_file, thread_count);
    
    cout << "File saved as: " << output_file << endl;
    
    return 0;
}