perf(httpfs): pool HTTP connections and cache remote file sizes

httpfs/src/httpfs.cpp

@@ -21,6 +21,13 @@ namespace httpfs_extension {
 
 using namespace lbug::common;
 
+std::unordered_map<std::string, std::unique_ptr<httplib::Client>>
+    HTTPFileSystem::sharedNoRedirectClients;
+std::mutex HTTPFileSystem::httpPoolMtx;
+std::mutex HTTPFileSystem::httpLockMtx;
+std::unordered_map<std::string, uint64_t> 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<std::mutex> lck{HTTPFileSystem::httpSizeCacheMtx};
+        auto it = HTTPFileSystem::httpSizeCache.find(path);
+        if (it != HTTPFileSystem::httpSizeCache.end()) {
+            length = it->second;
+            return;
+        }
+    }
+
     auto hfs = fileSystem->ptrCast<HTTPFileSystem>();
     initializeClient();
     auto res = hfs->headRequest(this->ptrCast<HTTPFileInfo>(), path, {});
@@ -278,6 +298,10 @@ void HTTPFileInfo::initMetadata() {
             // LCOV_EXCL_STOP
         }
     }
+    if (length > 0) {
+        std::lock_guard<std::mutex> lck{HTTPFileSystem::httpSizeCacheMtx};
+        HTTPFileSystem::httpSizeCache[path] = length;
+    }
 }
 
 void HTTPFileInfo::initialize(main::ClientContext* context) {
@@ -474,6 +498,51 @@ std::unique_ptr<httplib::Client> 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<httplib::Client>(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<std::mutex> lck{httpPoolMtx};
+    auto it = sharedNoRedirectClients.find(host);
+    if (it == sharedNoRedirectClients.end()) {
+        std::unique_ptr<httplib::Client> 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<std::mutex> lck{httpPoolMtx};
+    sharedNoRedirectClients.erase(host);
+    std::unique_ptr<httplib::Client> client{makeNoRedirectClient(host)};
+    auto it = sharedNoRedirectClients.emplace(host, std::move(client)).first;
+    return it->second.get();
+}
+
 std::unique_ptr<httplib::Headers> HTTPFileSystem::getHTTPHeaders(HeaderMap& headerMap) {
     auto headers = std::make_unique<httplib::Headers>();
     for (auto& entry : headerMap) {

httpfs/src/include/httpfs.h

@@ -6,6 +6,7 @@
 #include "httplib.h"
 #include "main/client_context.h"
 #include <list>
+#include <unordered_map>
 
 #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> 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<std::string, std::unique_ptr<httplib::Client>> 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<std::string, uint64_t> httpSizeCache;
+    static std::mutex httpSizeCacheMtx;
 };
 
 } // namespace httpfs_extension

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<httplib::Client> 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<HTTPResponse> 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<HTTPResponse> XetFileSystem::headRequest(common::FileInfo* /*fileInfo*/,
+std::unique_ptr<HTTPResponse> 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<httplib::Result(void)> request(
-        [&]() { return client->Head(hostPath.c_str(), *headers); });
-    std::function<void(void)> 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<HTTPResponse> 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<httplib::Result(void)> request(
+            [&]() { return client->Head(hostPath.c_str(), *headers); });
+        std::function<void(void)> 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<HTTPResponse> XetFileSystem::getRangeRequest(common::FileInfo* /*fileInfo*/,
+std::unique_ptr<HTTPResponse> 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<httplib::Result(void)> request(
-        [&]() { return client->Get(hostPath.c_str(), *headers); });
-    std::function<void(void)> 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<HTTPResponse> 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<httplib::Result(void)> request(
+            [&]() { return client->Get(hostPath.c_str(), *headers); });
+        std::function<void(void)> 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) {