diff --git a/httpfs/src/httpfs.cpp b/httpfs/src/httpfs.cpp index 0a9b482d..b7052782 100644 --- a/httpfs/src/httpfs.cpp +++ b/httpfs/src/httpfs.cpp @@ -21,6 +21,13 @@ namespace httpfs_extension { using namespace lbug::common; +std::unordered_map> + HTTPFileSystem::sharedNoRedirectClients; +std::mutex HTTPFileSystem::httpPoolMtx; +std::mutex HTTPFileSystem::httpLockMtx; +std::unordered_map HTTPFileSystem::httpSizeCache; +std::mutex HTTPFileSystem::httpSizeCacheMtx; + HTTPResponse::HTTPResponse(httplib::Response& res, std::string url) : code{res.status}, error{res.reason}, url{std::move(url)}, body{res.body} { for (auto& [name, value] : res.headers) { @@ -199,6 +206,19 @@ HTTPFileInfo::HTTPFileInfo(std::string path, FileSystem* fileSystem, int flags, httpConfig{context}, cachedFileInfo{nullptr} {} void HTTPFileInfo::initMetadata() { + // Remote files are immutable during a query; reuse the file size learned by + // an earlier openFile() for the same URL instead of paying a fresh HEAD + + // redirect round trip per open. This is what turns N opens of the same + // parquet file (~450 in a rel scan) into a single HEAD. + { + std::lock_guard lck{HTTPFileSystem::httpSizeCacheMtx}; + auto it = HTTPFileSystem::httpSizeCache.find(path); + if (it != HTTPFileSystem::httpSizeCache.end()) { + length = it->second; + return; + } + } + auto hfs = fileSystem->ptrCast(); initializeClient(); auto res = hfs->headRequest(this->ptrCast(), path, {}); @@ -278,6 +298,10 @@ void HTTPFileInfo::initMetadata() { // LCOV_EXCL_STOP } } + if (length > 0) { + std::lock_guard lck{HTTPFileSystem::httpSizeCacheMtx}; + HTTPFileSystem::httpSizeCache[path] = length; + } } void HTTPFileInfo::initialize(main::ClientContext* context) { @@ -474,6 +498,51 @@ std::unique_ptr HTTPFileSystem::getClient(const std::string& ho return client; } +// Build a no-follow httplib client for `host` that does not automatically +// follow redirects. The Xet resolve endpoint returns 302 Location headers which +// the filesystem follows itself, so redirects must be surfaced rather than +// absorbed (which would also point the pooled connection at the wrong host). +static httplib::Client* makeNoRedirectClient(const std::string& host) { + auto client = std::make_unique(host); + client->set_follow_location(false); + client->set_url_encode(false); + client->set_keep_alive(HTTPParams::DEFAULT_KEEP_ALIVE); +#ifdef CPPHTTPLIB_OPENSSL_SUPPORT + client->enable_server_certificate_verification(false); +#endif + client->set_write_timeout(HTTPParams::DEFAULT_TIMEOUT); + client->set_read_timeout(HTTPParams::DEFAULT_TIMEOUT); + client->set_connection_timeout(HTTPParams::DEFAULT_TIMEOUT); + client->set_decompress(false); + return client.release(); +} + +void HTTPFileSystem::lockHttp() { + httpLockMtx.lock(); +} + +void HTTPFileSystem::unlockHttp() { + httpLockMtx.unlock(); +} + +httplib::Client* HTTPFileSystem::getSharedNoRedirectClient(const std::string& host) { + std::lock_guard lck{httpPoolMtx}; + auto it = sharedNoRedirectClients.find(host); + if (it == sharedNoRedirectClients.end()) { + std::unique_ptr client{makeNoRedirectClient(host)}; + it = sharedNoRedirectClients.emplace(host, std::move(client)).first; + } + return it->second.get(); +} + +httplib::Client* HTTPFileSystem::evictAndGetSharedNoRedirectClient(const std::string& host) { + std::lock_guard lck{httpPoolMtx}; + sharedNoRedirectClients.erase(host); + std::unique_ptr client{makeNoRedirectClient(host)}; + auto it = sharedNoRedirectClients.emplace(host, std::move(client)).first; + return it->second.get(); +} + std::unique_ptr HTTPFileSystem::getHTTPHeaders(HeaderMap& headerMap) { auto headers = std::make_unique(); for (auto& entry : headerMap) { diff --git a/httpfs/src/include/httpfs.h b/httpfs/src/include/httpfs.h index baec31d4..a2998a7b 100644 --- a/httpfs/src/include/httpfs.h +++ b/httpfs/src/include/httpfs.h @@ -6,6 +6,7 @@ #include "httplib.h" #include "main/client_context.h" #include +#include #if defined(_WIN32) #define O_ACCMODE 0x0003 @@ -100,6 +101,27 @@ class HTTPFileSystem : public common::FileSystem { void cleanUP(main::ClientContext* context) override; + // Returns a process-wide shared no-follow httplib client for `host`, creating + // and caching it on first use. Because each client owns one persistent + // keep-alive TLS socket that is reused across requests, sharing a client per + // host (instead of creating a fresh one per file info / per retry) is what + // avoids a TCP+TLS reconnect on every request. Callers must serialize access + // to a pooled client around a request via lockHttp()/unlockHttp(). + static httplib::Client* getSharedNoRedirectClient(const std::string& host); + + // Drops the pooled client for `host` (if present) and returns a freshly + // created one, so a retry after a failed/dead connection does not keep + // reusing the same broken socket for the rest of the run. + static httplib::Client* evictAndGetSharedNoRedirectClient(const std::string& host); + + // Global lock guarding access to any pooled no-follow client. Requests that + // use a pooled client must hold this for the duration of the request. It must + // be released before handling a redirect (which recurses into another + // lockHttp() on the same thread), otherwise the non-recursive mutex + // self-deadlocks. + static void lockHttp(); + static void unlockHttp(); + protected: void readFromFile(common::FileInfo& fileInfo, void* buffer, uint64_t numBytes, uint64_t position) const override; @@ -135,6 +157,22 @@ class HTTPFileSystem : public common::FileSystem { private: std::unique_ptr cachedFileManager; std::mutex cachedFileManagerMtx; + + // Process-wide HTTP connection pool, keyed by host. Each entry is a client + // with one persistent keep-alive TLS socket reused across all requests to + // that host. Guarded by httpPoolMtx. The single httpLockMtx serializes + // access to any pooled client; it must NEVER be held across a redirect + // recursion (which re-enters on the same thread). + static std::unordered_map> sharedNoRedirectClients; + static std::mutex httpPoolMtx; + static std::mutex httpLockMtx; + + // Process-wide cache of remote file sizes keyed by URL. Parquet files are + // immutable during a query, so the very expensive HEAD round trip used to + // learn a file's size only needs to happen once per URL instead of once per + // openFile(). Guarded by httpSizeCacheMtx. + static std::unordered_map httpSizeCache; + static std::mutex httpSizeCacheMtx; }; } // namespace httpfs_extension diff --git a/httpfs/src/xetfs.cpp b/httpfs/src/xetfs.cpp index e5073e98..c89de499 100644 --- a/httpfs/src/xetfs.cpp +++ b/httpfs/src/xetfs.cpp @@ -82,13 +82,6 @@ std::string makeAbsoluteRedirectURL(const std::string& sourceURL, const std::str return host + basePath + location; } -std::unique_ptr getNoRedirectClient(const std::string& host) { - auto client = HTTPFileSystem::getClient(host); - client->set_follow_location(false); - client->set_url_encode(false); - return client; -} - std::unique_ptr synthesizeHeadResponse(const HTTPResponse& response, const std::string& url, const std::string& contentLength) { httplib::Response res; @@ -161,44 +154,86 @@ std::string XetFileSystem::toHuggingFaceURL(const std::string& path) { return buildResolveURLWithExplicitResolve("", segments); } -std::unique_ptr XetFileSystem::headRequest(common::FileInfo* /*fileInfo*/, +std::unique_ptr XetFileSystem::headRequest(common::FileInfo* fileInfo, const std::string& url, HeaderMap headerMap) const { const auto [host, hostPath] = HTTPFileSystem::parseUrl(url); auto headers = getHTTPHeaders(headerMap); - auto client = getNoRedirectClient(host); - std::function request( - [&]() { return client->Head(hostPath.c_str(), *headers); }); - std::function retry([&]() { client = getNoRedirectClient(host); }); + // Send one request to `host` under the shared-connection lock. The lock is + // scoped to this block and released before any redirect recursion below, + // because recursing into another host's request re-enters lockHttp() and a + // non-recursive mutex would self-deadlock if still held. + std::unique_ptr response; + { + HTTPFileSystem::lockHttp(); + struct HttpLockGuard { + ~HttpLockGuard() { HTTPFileSystem::unlockHttp(); } + } httpLockGuard; + + // Reuse the pooled connection for this host if one is healthy, else + // evict the stale one and build a fresh client. + httplib::Client* client = HTTPFileSystem::getSharedNoRedirectClient(host); + + std::function request( + [&]() { return client->Head(hostPath.c_str(), *headers); }); + std::function retry([&]() { + // A persistent connection may have gone stale (rate-limited or closed + // by the server); evict it so the retry gets a fresh socket instead of + // repeatedly timing out against the same dead connection. + client = HTTPFileSystem::evictAndGetSharedNoRedirectClient(host); + }); + + response = runRequestWithRetry(request, url, "HEAD", retry); + } - auto response = runRequestWithRetry(request, url, "HEAD", retry); if (response->code >= 300 && response->code < 400 && response->headers.contains("x-linked-size")) { return synthesizeHeadResponse(*response, url, response->headers["x-linked-size"]); } if (response->code >= 300 && response->code < 400 && response->headers.contains("Location")) { - return headRequest(nullptr, makeAbsoluteRedirectURL(url, response->headers["Location"]), + return headRequest(fileInfo, makeAbsoluteRedirectURL(url, response->headers["Location"]), headerMap); } return response; } -std::unique_ptr XetFileSystem::getRangeRequest(common::FileInfo* /*fileInfo*/, +std::unique_ptr XetFileSystem::getRangeRequest(common::FileInfo* fileInfo, const std::string& url, HeaderMap headerMap, uint64_t fileOffset, char* buffer, uint64_t bufferLen) const { const auto [host, hostPath] = HTTPFileSystem::parseUrl(url); auto headers = getHTTPHeaders(headerMap); headers->insert(std::make_pair("Range", std::format("bytes={}-{}", fileOffset, fileOffset + bufferLen - 1))); - auto client = getNoRedirectClient(host); - std::function request( - [&]() { return client->Get(hostPath.c_str(), *headers); }); - std::function retry([&]() { client = getNoRedirectClient(host); }); + // Send one range request to `host` under the shared-connection lock. The + // lock is scoped to this block and released before any redirect recursion + // below, because recursing into another host's request re-enters lockHttp() + // and a non-recursive mutex would self-deadlock if still held. + std::unique_ptr response; + { + HTTPFileSystem::lockHttp(); + struct HttpLockGuard { + ~HttpLockGuard() { HTTPFileSystem::unlockHttp(); } + } httpLockGuard; + + // Reuse the pooled connection for this host if one is healthy, else + // evict the stale one and build a fresh client. + httplib::Client* client = HTTPFileSystem::getSharedNoRedirectClient(host); + + std::function request( + [&]() { return client->Get(hostPath.c_str(), *headers); }); + std::function retry([&]() { + // A persistent connection may have gone stale (rate-limited or closed + // by the server); evict it so the retry gets a fresh socket instead of + // repeatedly timing out against the same dead connection. + client = HTTPFileSystem::evictAndGetSharedNoRedirectClient(host); + }); + + response = runRequestWithRetry(request, url, "GET Range", retry); + } - auto response = runRequestWithRetry(request, url, "GET Range", retry); if (response->code >= 300 && response->code < 400 && response->headers.contains("Location")) { - return getRangeRequest(nullptr, makeAbsoluteRedirectURL(url, response->headers["Location"]), + return getRangeRequest(fileInfo, makeAbsoluteRedirectURL(url, response->headers["Location"]), headerMap, fileOffset, buffer, bufferLen); } if (response->code >= 400) {