diff --git a/src/provider_common.c b/src/provider_common.c index 6a773b7..8a36d02 100644 --- a/src/provider_common.c +++ b/src/provider_common.c @@ -17,6 +17,15 @@ #include "utils/memutils.h" +/* + * Per-backend curl handle, kept so that its connection pool outlives a + * single request. See provider_do_curl_request() for why. + */ +static CURL *cached_curl = NULL; + +/* Initial response buffer size; embedding responses run to megabytes. */ +#define RESPONSE_BUFFER_INITIAL 32768 + /* * Curl write callback - accumulates response data into a ResponseBuffer. */ @@ -25,10 +34,55 @@ provider_write_callback(void *contents, size_t size, size_t nmemb, void *userp) { size_t realsize = size * nmemb; ResponseBuffer *mem = (ResponseBuffer *) userp; + size_t needed = mem->size + realsize + 1; + + /* + * Grow geometrically rather than to the exact size wanted. curl hands + * the body over in chunks of its own choosing, so sizing precisely meant + * a repalloc per chunk for the whole response: an embedding batch runs to + * megabytes, which is hundreds of them, each liable to copy everything + * received so far. Doubling makes that a handful. + * + * Nothing here may throw. This runs inside curl_easy_perform(), and a + * PostgreSQL error is a longjmp, which would abandon curl's handle and its + * connection mid-transfer. The handle is cached for the life of the + * backend, so every later request would go on reusing the wreckage. On + * failure we record why and return short instead, which curl reports as + * CURLE_WRITE_ERROR through the caller's existing error path. + */ + if (needed > mem->capacity) + { + size_t newcap = (mem->capacity > 0) ? mem->capacity : needed; + char *newdata; + + while (newcap < needed) + newcap *= 2; + + /* + * Doubling can overshoot MaxAllocSize while the response itself would + * still fit, so try the exact size before giving up on it. + */ + if (!AllocSizeIsValid(newcap)) + newcap = needed; + + if (!AllocSizeIsValid(newcap)) + { + mem->grow_failure = RESPONSE_GROW_TOO_LARGE; + return 0; + } - /* repalloc never returns NULL - it throws on allocation failure */ - mem->data = repalloc(mem->data, mem->size + realsize + 1); - /* flawfinder: ignore - buffer was realloced to mem->size + realsize + 1 */ + newdata = repalloc_extended(mem->data, newcap, MCXT_ALLOC_NO_OOM); + if (newdata == NULL) + { + mem->grow_failure = RESPONSE_GROW_NO_MEMORY; + return 0; + } + + mem->data = newdata; + mem->capacity = newcap; + } + + /* flawfinder: ignore - buffer holds at least mem->size + realsize + 1 */ memcpy(&(mem->data[mem->size]), contents, realsize); /* nosemgrep */ mem->size += realsize; mem->data[mem->size] = 0; @@ -366,18 +420,43 @@ provider_do_curl_request(const char *url, const char *auth_header, long response_code; /* Initialize response buffer */ - response_out->data = palloc(1); + response_out->capacity = RESPONSE_BUFFER_INITIAL; + response_out->data = palloc(response_out->capacity); response_out->data[0] = '\0'; response_out->size = 0; - - curl = curl_easy_init(); - if (!curl) + response_out->grow_failure = RESPONSE_GROW_OK; + + /* + * Reuse this backend's handle rather than making a new one per request. + * A fresh handle knows nothing of any connection, so every batch paid for + * a TCP connection and, against a real provider, a TLS handshake as well: + * one or two extra round trips to a host that is usually not nearby, on + * every batch of a table that may have thousands. Keeping the handle + * keeps curl's connection pool with it, so subsequent batches to the same + * endpoint reuse the established connection. + * + * curl_easy_reset() clears the options set below without disturbing that + * pool, which is what makes reuse survive the reset. + * + * The handle is deliberately never cleaned up: it lives as long as the + * backend, and an idle keep-alive connection is closed by the server at + * its own timeout, so nothing is held open indefinitely. + */ + if (cached_curl == NULL) { - *error_msg = pstrdup("Failed to initialize libcurl"); - pfree(response_out->data); - response_out->data = NULL; - return false; + cached_curl = curl_easy_init(); + if (!cached_curl) + { + *error_msg = pstrdup("Failed to initialize libcurl"); + pfree(response_out->data); + response_out->data = NULL; + return false; + } } + else + curl_easy_reset(cached_curl); + + curl = cached_curl; /* Set up headers */ headers = curl_slist_append(headers, "Content-Type: application/json; charset=utf-8"); @@ -399,16 +478,28 @@ provider_do_curl_request(const char *url, const char *auth_header, if (res != CURLE_OK) { - *error_msg = psprintf("curl_easy_perform() failed: %s", - curl_easy_strerror(res)); + /* + * A buffer we could not grow aborts the transfer, so check it first: + * curl only knows the write callback returned short, and would report + * that as a generic write failure. + */ + if (response_out->grow_failure == RESPONSE_GROW_TOO_LARGE) + *error_msg = psprintf("Response from %s is larger than can be " + "buffered (stopped after %zu bytes)", + provider_name, response_out->size); + else if (response_out->grow_failure == RESPONSE_GROW_NO_MEMORY) + *error_msg = psprintf("Out of memory buffering the response from " + "%s (stopped after %zu bytes)", + provider_name, response_out->size); + else + *error_msg = psprintf("curl_easy_perform() failed: %s", + curl_easy_strerror(res)); curl_slist_free_all(headers); - curl_easy_cleanup(curl); return false; } curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &response_code); curl_slist_free_all(headers); - curl_easy_cleanup(curl); if (response_code != 200) { diff --git a/src/provider_common.h b/src/provider_common.h index eb9f61d..fb3e44a 100644 --- a/src/provider_common.h +++ b/src/provider_common.h @@ -13,13 +13,28 @@ #include "pgedge_vectorizer.h" #include +/* + * Why the write callback could not grow the response buffer. Kept apart + * because the two want different advice: one says the provider sent more than + * can be held, the other says this server is short of memory. See + * provider_write_callback() for why the callback may not simply raise. + */ +typedef enum +{ + RESPONSE_GROW_OK = 0, + RESPONSE_GROW_TOO_LARGE, /* would exceed MaxAllocSize */ + RESPONSE_GROW_NO_MEMORY /* allocation refused */ +} ResponseGrowFailure; + /* * Response buffer for libcurl callbacks */ typedef struct { char *data; - size_t size; + size_t size; /* bytes of response held, excluding the terminator */ + size_t capacity; /* bytes allocated in data, including the terminator */ + ResponseGrowFailure grow_failure; } ResponseBuffer; /*