This is an automated email from the ASF dual-hosted git repository.
CalvinKirs pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 70f7dcbe11c [improvement](be) Reuse HTTP connections across AI
function requests (#68816)
70f7dcbe11c is described below
commit 70f7dcbe11c8e142b81536c1dfaaf505ec05437b
Author: linrrarity <[email protected]>
AuthorDate: Wed Oct 14 16:45:02 2026 +0800
[improvement](be) Reuse HTTP connections across AI function requests
(#68816)
### What problem does this PR solve?
Issue Number: None
Problem Summary:
AI scalar functions, EMBED and AI\_AGG recreate HTTP clients for each
request sequence, losing connection reuse between batches and Blocks.
This PR introduces `AIRequestExecutor` with exclusive client borrowing
and a bounded idle cache. Scalar functions retain the cache in their
execution contexts, while AI\_AGG shares it across local groups, flushes
and resets.
It also clears responses between attempts, detaches cancellation
callbacks before returning clients, sets JSON headers for LOCAL
providers without an API key, and fixes `max_retries` to count retries
after the initial request.
Unity compilation issues are fixed by adding missing protobuf and
RuntimeState includes and locally suppressing CLucene conversion
warnings, without changing CMakeLists.
### Release note
AI functions reuse HTTP connections across batches and Blocks.
`max_retries=0` performs one request; `max_retries=3` allows up to four
attempts
---
be/src/exprs/aggregate/aggregate_function_ai_agg.h | 51 +-
be/src/exprs/function/ai/ai_functions.h | 65 +-
be/src/exprs/function/ai/ai_request_executor.cpp | 301 ++++++
be/src/exprs/function/ai/ai_request_executor.h | 88 ++
be/src/service/http/http_client.cpp | 114 +-
be/src/service/http/http_client.h | 11 +
be/src/service/http/http_status.cpp | 4 +-
be/src/service/http/http_status.h | 4 +-
be/test/ai/ai_function_test.cpp | 22 +-
be/test/ai/ai_request_executor_test.cpp | 1087 ++++++++++++++++++++
be/test/service/http/http_client_test.cpp | 497 ++++++++-
.../property/constants/AIProperties.java | 2 +-
.../org/apache/doris/catalog/AIResourceTest.java | 34 +-
.../suites/ai_p0/test_ai_functions.groovy | 49 +-
14 files changed, 2150 insertions(+), 179 deletions(-)
diff --git a/be/src/exprs/aggregate/aggregate_function_ai_agg.h
b/be/src/exprs/aggregate/aggregate_function_ai_agg.h
index d3da899fca8..f7929cb0114 100644
--- a/be/src/exprs/aggregate/aggregate_function_ai_agg.h
+++ b/be/src/exprs/aggregate/aggregate_function_ai_agg.h
@@ -28,9 +28,9 @@
#include "core/types.h"
#include "exprs/aggregate/aggregate_function.h"
#include "exprs/function/ai/ai_adapter.h"
+#include "exprs/function/ai/ai_request_executor.h"
#include "runtime/query_context.h"
#include "runtime/runtime_state.h"
-#include "service/http/http_client.h"
#include "util/string_util.h"
namespace doris {
@@ -152,7 +152,10 @@ public:
}
}
- void set_query_context(QueryContext* context) { _ctx = context; }
+ void set_query_context(QueryContext* context,
std::shared_ptr<AIRequestExecutor> executor) {
+ _ctx = context;
+ _request_executor = std::move(executor);
+ }
const std::string& get_task() const { return _task; }
@@ -162,36 +165,16 @@ public:
private:
Status send_request_to_ai(const std::string& request_body, std::string&
response) const {
- // Mock path for testing
+ // Real providers still use HTTP in BE_TEST; only explicit mocks
bypass it.
#ifdef BE_TEST
const char* test_result = std::getenv("AI_TEST_RESULT");
- response = test_result != nullptr ? test_result : "this is a mock
response";
- return Status::OK();
-#endif
-
- return HttpClient::execute_with_retry(
- _ai_config.max_retries, _ai_config.retry_delay_second,
- [this, &request_body, &response](HttpClient* client) -> Status
{
- return this->do_send_request(client, request_body,
response);
- });
- }
-
- Status do_send_request(HttpClient* client, const std::string& request_body,
- std::string& response) const {
- RETURN_IF_ERROR(client->init(_ai_config.endpoint));
- if (_ctx == nullptr) {
- return Status::InternalError("Query context is null");
+ if (test_result != nullptr || _ai_config.provider_type == "MOCK") {
+ response = test_result != nullptr ? test_result : "this is a mock
response";
+ return Status::OK();
}
+#endif
- int64_t remaining_query_time =
_ctx->get_remaining_query_time_seconds();
- if (remaining_query_time <= 0) {
- return Status::TimedOut("Query timeout exceeded before AI
request");
- }
- client->set_timeout_ms(remaining_query_time * 1000);
-
- RETURN_IF_ERROR(_ai_adapter->set_authentication(client));
-
- return client->execute_post_request(request_body, &response);
+ return _request_executor->execute(request_body, response, _ai_config,
*_ai_adapter, _ctx);
}
// Treat the context window as a soft batching trigger instead of a hard
reject.
@@ -255,6 +238,8 @@ private:
}
QueryContext* _ctx = nullptr;
+ // Local reference only; excluded from serialization and merge().
+ std::shared_ptr<AIRequestExecutor> _request_executor;
AIResource _ai_config;
std::shared_ptr<AIAdapter> _ai_adapter;
std::string _task;
@@ -283,7 +268,8 @@ public:
void create(AggregateDataPtr __restrict place) const override {
new (place) AggregateFunctionAIAggData;
- data(place).set_query_context(_ctx);
+ // Share one cache across groups instead of retaining a client per
group.
+ data(place).set_query_context(_ctx, _request_executor);
}
void add(AggregateDataPtr __restrict place, const IColumn** columns,
ssize_t row_num,
@@ -323,7 +309,7 @@ public:
void reset(AggregateDataPtr place) const override {
data(place).reset();
- data(place).set_query_context(_ctx);
+ data(place).set_query_context(_ctx, _request_executor);
}
void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,
@@ -338,7 +324,8 @@ public:
void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,
Arena&) const override {
data(place).read(buf);
- data(place).set_query_context(_ctx);
+ // A state received from another node uses this evaluator's local
cache.
+ data(place).set_query_context(_ctx, _request_executor);
}
void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn&
to) const override {
@@ -355,6 +342,8 @@ public:
private:
QueryContext* _ctx = nullptr;
+ // Shared by local group states across input Blocks and result extraction.
+ std::shared_ptr<AIRequestExecutor> _request_executor =
std::make_shared<AIRequestExecutor>();
};
} // namespace doris
diff --git a/be/src/exprs/function/ai/ai_functions.h
b/be/src/exprs/function/ai/ai_functions.h
index f5bd10852eb..1ded3e94917 100644
--- a/be/src/exprs/function/ai/ai_functions.h
+++ b/be/src/exprs/function/ai/ai_functions.h
@@ -42,11 +42,10 @@
#include "core/data_type/primitive_type.h"
#include "exec/common/util.hpp"
#include "exprs/function/ai/ai_adapter.h"
+#include "exprs/function/ai/ai_request_executor.h"
#include "exprs/function/function.h"
#include "runtime/query_context.h"
#include "runtime/runtime_state.h"
-#include "service/http/http_client.h"
-#include "util/security.h"
#include "util/string_util.h"
#include "util/threadpool.h"
@@ -67,6 +66,14 @@ public:
bool is_blockable() const override { return true; }
+ Status close(FunctionContext* context, FunctionContext::FunctionStateScope
scope) override {
+ if (scope == FunctionContext::THREAD_LOCAL) {
+ // Release this execution context's cached clients.
+ context->set_function_state(scope, nullptr);
+ }
+ return Status::OK();
+ }
+
bool use_default_implementation_for_nulls() const final { return false; }
DataTypePtr get_return_type_impl(const DataTypes& arguments) const final {
@@ -169,53 +176,19 @@ protected:
}
}
- // Executes one HTTP POST request and validates transport-level success.
- Status do_send_request(HttpClient* client, const std::string& request_body,
- std::string& response, const AIResource& config,
- std::shared_ptr<AIAdapter>& adapter,
FunctionContext* context) const {
- RETURN_IF_ERROR(client->init(config.endpoint, false));
-
- QueryContext* query_ctx = context->state()->get_query_ctx();
- int64_t remaining_query_time =
query_ctx->get_remaining_query_time_seconds();
- if (remaining_query_time <= 0) {
- return Status::TimedOut("Query timeout exceeded before AI
request");
- }
-
- client->set_timeout_ms(remaining_query_time * 1000);
-
- if (!config.api_key.empty()) {
- RETURN_IF_ERROR(adapter->set_authentication(client));
- }
-
- Status st = client->execute_post_request(request_body, &response);
- long http_status = client->get_http_status();
-
- if (!st.ok()) {
- LOG(INFO) << "AI HTTP request failed before status validation,
provider="
- << config.provider_type << ", model=" <<
config.model_name
- << ", endpoint=" << mask_token(config.endpoint)
- << ", exec_status=" << st.to_string() << ",
response_body=" << response;
- return st;
- }
- if (http_status != 200) {
- return Status::HttpError(
- "http status code is not 200, code={}, url={},
response_body={}", http_status,
- mask_token(config.endpoint), response);
- }
- return Status::OK();
- }
-
- // Sends the request with retry mechanism for handling transient failures
Status send_request_to_llm(const std::string& request_body, std::string&
response,
const AIResource& config,
std::shared_ptr<AIAdapter>& adapter,
FunctionContext* context) const {
- return HttpClient::execute_with_retry(config.max_retries,
config.retry_delay_second,
- [this, &request_body, &response,
&config, &adapter,
- context](HttpClient* client) ->
Status {
- return
this->do_send_request(client, request_body,
-
response, config,
-
adapter, context);
- });
+ // Retain the cache across batches/Blocks.
+ auto* executor = static_cast<AIRequestExecutor*>(
+ context->get_function_state(FunctionContext::THREAD_LOCAL));
+ if (executor == nullptr) {
+ auto state = std::make_shared<AIRequestExecutor>();
+ executor = state.get();
+ context->set_function_state(FunctionContext::THREAD_LOCAL,
std::move(state));
+ }
+ return executor->execute(request_body, response, config, *adapter,
+ context->state()->get_query_ctx());
}
// Provider-reusable helper for string-returning functions.
diff --git a/be/src/exprs/function/ai/ai_request_executor.cpp
b/be/src/exprs/function/ai/ai_request_executor.cpp
new file mode 100644
index 00000000000..e2b77deee37
--- /dev/null
+++ b/be/src/exprs/function/ai/ai_request_executor.cpp
@@ -0,0 +1,301 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "exprs/function/ai/ai_request_executor.h"
+
+#include <algorithm>
+#include <charconv>
+#include <chrono>
+#include <limits>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+#include "exprs/function/ai/ai_adapter.h"
+#include "runtime/query_context.h"
+#include "service/http/http_client.h"
+#include "util/defer_op.h"
+#include "util/security.h"
+
+namespace doris {
+// Borrowing removes a client from idle, ensuring exclusive use of its easy
handle.
+struct AIRequestExecutor::ClientCache {
+ std::unique_ptr<HttpClient> borrow() {
+ {
+ std::lock_guard lock(mutex);
+ if (!idle.empty()) {
+ auto client = std::move(idle.back());
+ idle.pop_back();
+ return client;
+ }
+ }
+ return std::make_unique<HttpClient>();
+ }
+
+ void release(std::unique_ptr<HttpClient> client) {
+ (*client).reset();
+ // Bound idle clients, not TCP connections or in-flight requests.
+ // Excess clients are destroyed after releasing the lock.
+ constexpr size_t max_idle_clients = 8;
+ std::lock_guard lock(mutex);
+ if (idle.size() < max_idle_clients) {
+ idle.push_back(std::move(client));
+ }
+ }
+
+ std::mutex mutex;
+ std::vector<std::unique_ptr<HttpClient>> idle;
+};
+
+AIRequestExecutor::AIRequestExecutor()
+ : _now([] { return Clock::now(); }),
+ _sleep([](Clock::duration delay) {
std::this_thread::sleep_for(delay); }),
+ _clients(std::make_unique<ClientCache>()) {}
+AIRequestExecutor::~AIRequestExecutor() = default;
+
+AIErrorType AIRequestExecutor::classify_http_status(long http_status) {
+ switch (http_status) {
+ case 200:
+ return AIErrorType::NO_ERROR;
+ case 429:
+ case 408:
+ case 500:
+ case 502:
+ case 503:
+ case 504:
+ case 529: // Claude reports temporary overload with this status.
+ return AIErrorType::RETRYABLE;
+ default:
+ return AIErrorType::NON_RETRYABLE;
+ }
+}
+
+AIErrorType AIRequestExecutor::classify_curl_error(int curl_code) {
+ switch (curl_code) {
+ case CURLE_COULDNT_RESOLVE_PROXY:
+ case CURLE_COULDNT_RESOLVE_HOST:
+ case CURLE_COULDNT_CONNECT:
+ case CURLE_SEND_ERROR:
+ case CURLE_RECV_ERROR:
+ case CURLE_GOT_NOTHING:
+ case CURLE_PARTIAL_FILE:
+ case CURLE_HTTP2:
+ case CURLE_HTTP2_STREAM:
+ return AIErrorType::RETRYABLE;
+ default:
+ // Cancellation, timeout, TLS/configuration and local failures are
terminal.
+ return AIErrorType::NON_RETRYABLE;
+ }
+}
+
+int64_t AIRequestExecutor::calculate_retry_delay_ms(int32_t base_seconds) {
+ return static_cast<int64_t>(base_seconds) * 1000;
+}
+
+int64_t AIRequestExecutor::parse_retry_after_ms(const std::string& value,
int64_t unix_seconds) {
+ const auto first = value.find_first_not_of(" \t");
+ if (first == std::string::npos) {
+ return 0;
+ }
+ const auto last = value.find_last_not_of(" \t");
+ const std::string trimmed = value.substr(first, last - first + 1);
+ int64_t seconds = 0;
+ if (trimmed.find_first_not_of("0123456789") == std::string::npos) {
+ const auto parsed =
+ std::from_chars(trimmed.data(), trimmed.data() +
trimmed.size(), seconds);
+ if (parsed.ec == std::errc::result_out_of_range) {
+ return std::numeric_limits<int64_t>::max();
+ }
+ } else {
+ if (trimmed.find(':') == std::string::npos) {
+ return 0;
+ }
+ const auto date = curl_getdate(trimmed.c_str(), nullptr);
+ if (date < 0 || date <= unix_seconds) {
+ return 0;
+ }
+ seconds = date - unix_seconds;
+ }
+ return seconds > std::numeric_limits<int64_t>::max() / 1000
+ ? std::numeric_limits<int64_t>::max()
+ : seconds * 1000;
+}
+
+Status AIRequestExecutor::check_query_status(QueryContext* query_context,
+ Clock::time_point deadline) const
{
+ if (query_context->is_cancelled()) {
+ return query_context->exec_status();
+ }
+ if (_now() >= deadline ||
query_context->get_remaining_query_time_seconds() <= 0) {
+ return Status::TimedOut("Query timeout exceeded during AI request");
+ }
+ return Status::OK();
+}
+
+Status AIRequestExecutor::wait_for_retry(int64_t delay_ms, QueryContext*
query_context,
+ Clock::time_point deadline,
AIRequestStats& stats) const {
+ RETURN_IF_ERROR(check_query_status(query_context, deadline));
+ const auto started = _now();
+ const auto remaining_ms =
+ std::chrono::duration_cast<std::chrono::milliseconds>(deadline -
started).count();
+ // Never truncate Retry-After to the remaining budget and send earlier
than requested.
+ if (delay_ms >= remaining_ms) {
+ return Status::TimedOut("AI retry delay {} ms exceeds remaining query
budget {} ms",
+ delay_ms, remaining_ms);
+ }
+ Defer record_wait([&] {
+ stats.retry_wait_ms +=
+ std::chrono::duration_cast<std::chrono::milliseconds>(_now() -
started).count();
+ });
+ const auto retry_at = started + std::chrono::milliseconds(delay_ms);
+ while (_now() < retry_at) {
+ RETURN_IF_ERROR(check_query_status(query_context, deadline));
+ _sleep(std::min<Clock::duration>(retry_at - _now(),
std::chrono::milliseconds(100)));
+ }
+ return check_query_status(query_context, deadline);
+}
+
+AIErrorType AIRequestExecutor::classify_response(int curl_code, long
http_status,
+ const AIResource& config,
+ const std::string& response,
Status& status) {
+ if (!status.ok()) {
+ if (curl_code == CURLE_OPERATION_TIMEDOUT) {
+ status = Status::TimedOut("AI HTTP attempt timed out, endpoint={}",
+ mask_token(config.endpoint));
+ return AIErrorType::NON_RETRYABLE;
+ }
+ const auto error = classify_curl_error(curl_code);
+ // Received HTTP error headers still govern partial/disconnected
responses.
+ return error == AIErrorType::RETRYABLE && http_status >= 400
+ ? classify_http_status(http_status)
+ : error;
+ }
+ const auto error = classify_http_status(http_status);
+ if (error != AIErrorType::NO_ERROR) {
+ status = Status::HttpError("http status code is not 200, code={},
url={}, response_body={}",
+ http_status, mask_token(config.endpoint),
response);
+ }
+ return error;
+}
+
+Status AIRequestExecutor::execute_single_request(
+ HttpClient* client, const std::string& request_body, std::string&
response,
+ const AIResource& config, const AIAdapter& adapter, QueryContext*
query_context,
+ Clock::time_point deadline, AIRequestStats& stats) {
+ // Read HTTP error bodies for diagnostics, using the same mode for all AI
functions.
+ RETURN_IF_ERROR(client->init(config.endpoint, false));
+ RETURN_IF_ERROR(adapter.set_authentication(client));
+
+ // This attempt shares the logical request's deadline.
+ const auto remaining_ms =
+ std::chrono::duration_cast<std::chrono::milliseconds>(deadline -
_now()).count();
+ RETURN_IF_ERROR(check_query_status(query_context, deadline));
+ client->set_timeout_ms(std::max<int64_t>(1, remaining_ms));
+
+ // Check cancellation during the transfer, then clear the callback before
reuse.
+ client->set_abort_callback([this, query_context, deadline] {
+ return query_context->is_cancelled() || _now() >= deadline;
+ });
+ ++stats.attempts;
+ Status status = client->execute_post_request(request_body, &response);
+ client->set_abort_callback({});
+ return status;
+}
+
+Status AIRequestExecutor::prepare_retry(HttpClient& client, const AIResource&
config,
+ QueryContext* query_context,
Clock::time_point deadline,
+ AIErrorType error, const Status&
request_status,
+ AIRequestStats& stats) const {
+ // Stop on non-retryable errors or exhausted retry budgets.
+ if (error == AIErrorType::NON_RETRYABLE || stats.attempts - 1 >=
config.max_retries) {
+ return request_status;
+ }
+
+ // Keep the configured interval fixed; Retry-After may require a longer
wait.
+ int64_t delay_ms = calculate_retry_delay_ms(config.retry_delay_second);
+ std::string retry_after;
+ RETURN_IF_ERROR(client.get_response_header("Retry-After", &retry_after));
+ const auto unix_seconds = std::chrono::duration_cast<std::chrono::seconds>(
+
std::chrono::system_clock::now().time_since_epoch())
+ .count();
+ delay_ms = std::max(delay_ms, parse_retry_after_ms(retry_after,
unix_seconds));
+
+ // Wait within the query deadline, polling for cancellation before the
next send.
+ LOG(INFO) << "AI retry scheduled, next_attempt=" << stats.attempts + 1
+ << ", delay_ms=" << delay_ms;
+ return wait_for_retry(delay_ms, query_context, deadline, stats);
+}
+
+Status AIRequestExecutor::execute(const std::string& request_body,
std::string& response,
+ const AIResource& config, const AIAdapter&
adapter,
+ QueryContext* query_context, AIRequestStats*
output_stats) {
+ // Preserve statistics on success, failure, cancellation and timeout.
+ AIRequestStats stats;
+ Defer publish_stats([&] {
+ if (output_stats != nullptr) {
+ *output_stats = stats;
+ }
+ });
+
+ // Establish one deadline for all attempts and waits.
+ if (query_context == nullptr) {
+ return Status::InternalError("Query context is null");
+ }
+ const auto deadline =
+ _now() +
std::chrono::seconds(query_context->get_remaining_query_time_seconds());
+ RETURN_IF_ERROR(check_query_status(query_context, deadline));
+
+ // Keep the borrowed client across retries, then return it for subsequent
batches.
+ auto client = _clients->borrow();
+ Defer return_client([&] { _clients->release(std::move(client)); });
+
+ while (true) {
+ // Send one request; setup failures return without inspecting an HTTP
response.
+ RETURN_IF_ERROR(check_query_status(query_context, deadline));
+ response.clear();
+ const auto attempts_before = stats.attempts;
+ Status status = execute_single_request(client.get(), request_body,
response, config,
+ adapter, query_context,
deadline, stats);
+ if (stats.attempts == attempts_before) {
+ return status;
+ }
+
+ // Cancellation and timeout take precedence over HTTP/curl error
classification.
+ RETURN_IF_ERROR(check_query_status(query_context, deadline));
+ const auto curl_code = client->get_last_curl_code();
+ const long http_status = client->get_http_status();
+ const auto error = classify_response(curl_code, http_status, config,
response, status);
+ if (error == AIErrorType::NO_ERROR) {
+ return Status::OK();
+ }
+
+ // Record the failed transfer and provider response for diagnosis.
+ LOG(INFO) << "AI request failed, provider=" << config.provider_type
+ << ", model=" << config.model_name << ", endpoint=" <<
mask_token(config.endpoint)
+ << ", attempt=" << stats.attempts << ", http_status=" <<
http_status
+ << ", curl_code=" << curl_code
+ << ", curl_error=" <<
curl_easy_strerror(static_cast<CURLcode>(curl_code))
+ << ", retryable=" << (error != AIErrorType::NON_RETRYABLE)
+ << ", response_body=" << response;
+
+ // Complete retry checks and waiting before starting another request.
+ RETURN_IF_ERROR(
+ prepare_retry(*client, config, query_context, deadline, error,
status, stats));
+ }
+}
+
+} // namespace doris
diff --git a/be/src/exprs/function/ai/ai_request_executor.h
b/be/src/exprs/function/ai/ai_request_executor.h
new file mode 100644
index 00000000000..ec2648c4794
--- /dev/null
+++ b/be/src/exprs/function/ai/ai_request_executor.h
@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <chrono>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "common/status.h"
+
+namespace doris {
+
+struct AIResource;
+class AIAdapter;
+class QueryContext;
+class HttpClient;
+enum class AIErrorType {
+ NO_ERROR, // HTTP succeeded; the caller still needs to parse the
response.
+ NON_RETRYABLE, // Stop retrying this request.
+ RETRYABLE, // Retry within the configured count and query deadline.
+};
+
+struct AIRequestStats {
+ int64_t attempts = 0;
+ int64_t retry_wait_ms = 0;
+};
+
+// Shared synchronous AI transport with a client cache spanning batches and
Blocks.
+// Functions own payloads and result parsing; concurrent calls borrow
exclusive clients.
+class AIRequestExecutor {
+public:
+ AIRequestExecutor();
+ ~AIRequestExecutor();
+ AIRequestExecutor(const AIRequestExecutor&) = delete;
+ AIRequestExecutor& operator=(const AIRequestExecutor&) = delete;
+
+ // max_retries counts extra attempts for all retryable HTTP/curl errors.
+ Status execute(const std::string& request_body, std::string& response,
const AIResource& config,
+ const AIAdapter& adapter, QueryContext* query_context,
+ AIRequestStats* stats = nullptr);
+
+private:
+ using Clock = std::chrono::steady_clock;
+ static int64_t calculate_retry_delay_ms(int32_t base_seconds);
+ static int64_t parse_retry_after_ms(const std::string& value, int64_t
unix_seconds);
+ static AIErrorType classify_http_status(long http_status);
+ static AIErrorType classify_curl_error(int curl_code);
+ static AIErrorType classify_response(int curl_code, long http_status,
const AIResource& config,
+ const std::string& response, Status&
status);
+ Status check_query_status(QueryContext* query_context, Clock::time_point
deadline) const;
+ Status wait_for_retry(int64_t delay_ms, QueryContext* query_context,
Clock::time_point deadline,
+ AIRequestStats& stats) const;
+
+ // Return the original error when retries stop; otherwise wait before the
next send.
+ Status prepare_retry(HttpClient& client, const AIResource& config,
QueryContext* query_context,
+ Clock::time_point deadline, AIErrorType error,
+ const Status& request_status, AIRequestStats& stats)
const;
+
+ Status execute_single_request(HttpClient* client, const std::string&
request_body,
+ std::string& response, const AIResource&
config,
+ const AIAdapter& adapter, QueryContext*
query_context,
+ Clock::time_point deadline, AIRequestStats&
stats);
+
+ // Clock hooks keep waiting tests deterministic without real sleeps.
+ std::function<Clock::time_point()> _now;
+ std::function<void(Clock::duration)> _sleep;
+ struct ClientCache;
+ std::unique_ptr<ClientCache> _clients;
+};
+
+} // namespace doris
diff --git a/be/src/service/http/http_client.cpp
b/be/src/service/http/http_client.cpp
index 7a82549f78e..8f1e3ee548c 100644
--- a/be/src/service/http/http_client.cpp
+++ b/be/src/service/http/http_client.cpp
@@ -18,11 +18,17 @@
#include "service/http/http_client.h"
#include <absl/strings/str_split.h>
+#include <fcntl.h>
#include <glog/logging.h>
#include <unistd.h>
+#include <chrono>
+#include <filesystem>
+#include <fstream>
#include <memory>
+#include <mutex>
#include <ostream>
+#include <vector>
#include "common/cast_set.h"
#include "common/config.h"
@@ -281,29 +287,90 @@ HttpClient::~HttpClient() {
}
}
-Status HttpClient::init(const std::string& url, bool set_fail_on_error) {
- if (_curl == nullptr) {
- _curl = curl_easy_init();
- if (_curl == nullptr) {
- return Status::InternalError("fail to initialize curl");
+std::string HttpClient::_get_ca_certificate_path(const std::string&
configured_paths) {
+ struct CertificatePathCache {
+ std::mutex mutex;
+ std::string configured_paths;
+ std::string selected_path;
+ std::chrono::steady_clock::time_point refresh_after;
+ };
+ static CertificatePathCache cache;
+ std::scoped_lock lock(cache.mutex);
+ const auto now = std::chrono::steady_clock::now();
+ if (cache.configured_paths != configured_paths || now >=
cache.refresh_after) {
+ cache.selected_path.clear();
+ const std::vector<std::string> candidates =
+ absl::StrSplit(configured_paths, ';', absl::SkipEmpty());
+ for (const auto& path : candidates) {
+ std::error_code error;
+ if (!std::filesystem::is_regular_file(path, error) || error) {
+ continue;
+ }
+ std::ifstream file(path);
+ if (file.good()) {
+ cache.selected_path = path;
+ break;
+ }
}
- } else {
- curl_easy_reset(_curl);
+ cache.configured_paths = configured_paths;
+ // Configuration changes refresh immediately; also refresh filesystem
changes periodically.
+ cache.refresh_after = now + std::chrono::minutes(1);
+ }
+ return cache.selected_path;
+}
+
+Status HttpClient::_set_ca_certificate_from_config() {
+ std::string configured_paths;
+ {
+ std::scoped_lock lock(*config::get_mutable_string_config_lock());
+ configured_paths = config::ca_cert_file_paths;
+ }
+ const auto path = _get_ca_certificate_path(configured_paths);
+ if (path.empty()) {
+ // No usable candidate: retain libcurl's default CA settings and
verification.
+ return Status::OK();
}
+ const auto code = curl_easy_setopt(_curl, CURLOPT_CAINFO, path.c_str());
+ if (code != CURLE_OK) {
+ LOG(WARNING) << "failed to set CURLOPT_CAINFO, path=" << path
+ << ", errmsg=" << _to_errmsg(code);
+ return Status::InternalError("failed to set CA certificate {}: {}",
path, _to_errmsg(code));
+ }
+ return Status::OK();
+}
+
+void HttpClient::reset() {
+ if (_curl != nullptr) {
+ curl_easy_reset(_curl);
+ }
if (_header_list != nullptr) {
curl_slist_free_all(_header_list);
_header_list = nullptr;
}
- // curl_easy_reset() dropped the progress callback options, so drop the
functor too.
_abort_callback = nullptr;
- // set error_buf
+ _callback = nullptr;
_error_buf[0] = 0;
+ _last_curl_code = CURLE_OK;
+ _method = GET;
+}
+
+Status HttpClient::init(const std::string& url, bool set_fail_on_error) {
+ reset();
+ if (_curl == nullptr) {
+ _curl = curl_easy_init();
+ if (_curl == nullptr) {
+ return Status::InternalError("fail to initialize curl");
+ }
+ }
auto code = curl_easy_setopt(_curl, CURLOPT_ERRORBUFFER, _error_buf);
if (code != CURLE_OK) {
LOG(WARNING) << "fail to set CURLOPT_ERRORBUFFER, msg=" <<
_to_errmsg(code);
return Status::InternalError("fail to set error buffer");
}
+ // Apply runtime CA paths after reset, including HTTP requests redirected
to HTTPS.
+ RETURN_IF_ERROR(_set_ca_certificate_from_config());
+
// forbid signals
code = curl_easy_setopt(_curl, CURLOPT_NOSIGNAL, 1L);
if (code != CURLE_OK) {
@@ -450,6 +517,7 @@ Status HttpClient::execute(const std::function<bool(const
void* data, size_t len
}
_callback = &callback;
auto code = curl_easy_perform(_curl);
+ _last_curl_code = code;
if (code != CURLE_OK) {
std::string url = mask_token(_get_url());
LOG(WARNING) << "fail to execute HTTP client, errmsg=" <<
_to_errmsg(code)
@@ -464,24 +532,28 @@ Status HttpClient::execute(const std::function<bool(const
void* data, size_t len
return Status::OK();
}
-Status HttpClient::get_content_md5(std::string* md5) const {
- struct curl_header* header_ptr;
- auto code = curl_easy_header(_curl, HttpHeaders::CONTENT_MD5, 0,
CURLH_HEADER, 0, &header_ptr);
+Status HttpClient::get_response_header(const char* name, std::string* value)
const {
+ struct curl_header* header = nullptr;
+ // -1 selects the final response after redirects, rather than the first
request.
+ const auto code = curl_easy_header(_curl, name, 0, CURLH_HEADER, -1,
&header);
+ value->clear();
if (code == CURLHE_MISSING || code == CURLHE_NOHEADERS) {
- // no such headers exists
- md5->clear();
return Status::OK();
- } else if (code != CURLHE_OK) {
- auto msg = fmt::format("failed to get http header {}: {} ({})",
HttpHeaders::CONTENT_MD5,
- header_error_msg(code), code);
- LOG(WARNING) << msg << ", trace=" << get_stack_trace();
- return Status::HttpError(std::move(msg));
}
-
- *md5 = header_ptr->value;
+ if (code != CURLHE_OK) {
+ auto message = fmt::format("failed to get http header {}: {} ({})",
name,
+ header_error_msg(code), code);
+ LOG(WARNING) << message << ", trace=" << get_stack_trace();
+ return Status::HttpError(std::move(message));
+ }
+ *value = header->value;
return Status::OK();
}
+Status HttpClient::get_content_md5(std::string* md5) const {
+ return get_response_header(HttpHeaders::CONTENT_MD5, md5);
+}
+
Status HttpClient::download(const std::string& local_path) {
set_method(GET);
set_speed_limit();
diff --git a/be/src/service/http/http_client.h
b/be/src/service/http/http_client.h
index cdeaf9a4d34..ddbcea67c2f 100644
--- a/be/src/service/http/http_client.h
+++ b/be/src/service/http/http_client.h
@@ -49,6 +49,9 @@ public:
// you can call this multiple times to reuse this object
Status init(const std::string& url, bool set_fail_on_error = true);
+ // Release request-owned options and callbacks while retaining reusable
connections.
+ void reset();
+
void set_method(HttpMethod method);
void set_basic_auth(const std::string& user, const std::string& passwd) {
@@ -149,6 +152,11 @@ public:
// Get the value of the header CONTENT-MD5. The output is empty if no such
header exists.
Status get_content_md5(std::string* md5) const;
+ CURLcode get_last_curl_code() const { return _last_curl_code; }
+
+ // Read the final response header; clear value when the header is absent.
+ Status get_response_header(const char* name, std::string* value) const;
+
long get_http_status() const {
long code;
curl_easy_getinfo(_curl, CURLINFO_RESPONSE_CODE, &code);
@@ -202,11 +210,14 @@ public:
}
private:
+ static std::string _get_ca_certificate_path(const std::string&
configured_paths);
+ Status _set_ca_certificate_from_config();
const char* _to_errmsg(CURLcode code) const;
const char* _get_url() const;
private:
CURL* _curl = nullptr;
+ CURLcode _last_curl_code = CURLE_OK;
using HttpCallback = std::function<bool(const void* data, size_t length)>;
const HttpCallback* _callback = nullptr;
std::function<bool()> _abort_callback;
diff --git a/be/src/service/http/http_status.cpp
b/be/src/service/http/http_status.cpp
index 2804f15e91d..37ca7151a50 100644
--- a/be/src/service/http/http_status.cpp
+++ b/be/src/service/http/http_status.cpp
@@ -58,12 +58,14 @@ static std::map<HttpStatus, std::string> s_reason_map = {
{HttpStatus::UNSUPPORTED_MEDIA_TYPE, "Unsupported Media Type"},
{HttpStatus::REQUESTED_RANGE_NOT_SATISFIED, "Requested range not
satisfiable"},
{HttpStatus::EXPECTATION_FAILED, "Expectation Failed"},
+ {HttpStatus::TOO_MANY_REQUESTS, "Too Many Requests"},
{HttpStatus::INTERNAL_SERVER_ERROR, "Internal Server Error"},
{HttpStatus::NOT_IMPLEMENTED, "Not Implemented"},
{HttpStatus::BAD_GATEWAY, "Bad Gateway"},
{HttpStatus::SERVICE_UNAVAILABLE, "Service Unavailable"},
{HttpStatus::GATEWAY_TIMEOUT, "Gateway Time-out"},
- {HttpStatus::HTTP_VERSION_NOT_SUPPORTED, "HTTP Version not
supported"}};
+ {HttpStatus::HTTP_VERSION_NOT_SUPPORTED, "HTTP Version not supported"},
+ {HttpStatus::OVERLOADED, "Overloaded"}};
std::string default_reason(const HttpStatus& status) {
auto iter = s_reason_map.find(status);
diff --git a/be/src/service/http/http_status.h
b/be/src/service/http/http_status.h
index c6d8e21712e..c8b7a83fb73 100644
--- a/be/src/service/http/http_status.h
+++ b/be/src/service/http/http_status.h
@@ -56,12 +56,14 @@ enum HttpStatus {
UNSUPPORTED_MEDIA_TYPE = 415,
REQUESTED_RANGE_NOT_SATISFIED = 416,
EXPECTATION_FAILED = 417,
+ TOO_MANY_REQUESTS = 429,
INTERNAL_SERVER_ERROR = 500,
NOT_IMPLEMENTED = 501,
BAD_GATEWAY = 502,
SERVICE_UNAVAILABLE = 503,
GATEWAY_TIMEOUT = 504,
- HTTP_VERSION_NOT_SUPPORTED = 505
+ HTTP_VERSION_NOT_SUPPORTED = 505,
+ OVERLOADED = 529
};
std::string to_code(const HttpStatus& status);
diff --git a/be/test/ai/ai_function_test.cpp b/be/test/ai/ai_function_test.cpp
index a3d012c49f5..c2f01400d25 100644
--- a/be/test/ai/ai_function_test.cpp
+++ b/be/test/ai/ai_function_test.cpp
@@ -52,7 +52,7 @@ namespace doris {
class FunctionAITransportTestHelper : public FunctionAISentiment {
public:
- using FunctionAISentiment::do_send_request;
+ using FunctionAISentiment::send_request_to_llm;
};
class FunctionAIFilterBatchTestHelper : public
AIFunction<FunctionAIFilterBatchTestHelper> {
@@ -1646,7 +1646,7 @@ TEST(AIFunctionTest, ExecuteBatchRequestSuccess) {
config.provider_type = "OPENAI";
config.model_name = "test-model";
config.api_key = "secret";
- config.max_retries = 1;
+ config.max_retries = 0;
std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
adapter->init(config);
@@ -1685,7 +1685,7 @@ TEST(AIFunctionTest,
ExecuteBatchRequestResultSizeMismatch) {
config.provider_type = "OPENAI";
config.model_name = "test-model";
config.api_key = "secret";
- config.max_retries = 1;
+ config.max_retries = 0;
std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
adapter->init(config);
@@ -1715,14 +1715,14 @@ TEST(AIFunctionTest, DoSendRequestTransportError) {
config.provider_type = "OPENAI";
config.model_name = "test-model";
config.api_key = "secret";
+ config.max_retries = 0;
std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
adapter->init(config);
- HttpClient client;
std::string response;
FunctionAITransportTestHelper helper;
- Status st = helper.do_send_request(&client, "{}", response, config,
adapter, ctx.get());
+ Status st = helper.send_request_to_llm("{}", response, config, adapter,
ctx.get());
ASSERT_FALSE(st.ok());
ASSERT_EQ(response, "");
@@ -1744,15 +1744,15 @@ TEST(AIFunctionTest, DoSendRequestNon200) {
config.provider_type = "OPENAI";
config.model_name = "test-model";
config.api_key = "secret";
+ config.max_retries = 0;
std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
adapter->init(config);
- HttpClient client;
std::string response;
FunctionAITransportTestHelper helper;
- Status st = helper.do_send_request(&client, R"({"message":"hello"})",
response, config, adapter,
- ctx.get());
+ Status st = helper.send_request_to_llm(R"({"message":"hello"})", response,
config, adapter,
+ ctx.get());
ASSERT_FALSE(st.ok());
ASSERT_NE(st.to_string().find("http status code is not 200"),
std::string::npos);
@@ -1779,15 +1779,15 @@ TEST(AIFunctionTest, DoSendRequestSuccess) {
config.provider_type = "OPENAI";
config.model_name = "test-model";
config.api_key = "secret";
+ config.max_retries = 0;
std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
adapter->init(config);
- HttpClient client;
std::string response;
FunctionAITransportTestHelper helper;
- Status st = helper.do_send_request(&client, R"({"message":"hello"})",
response, config, adapter,
- ctx.get());
+ Status st = helper.send_request_to_llm(R"({"message":"hello"})", response,
config, adapter,
+ ctx.get());
ASSERT_TRUE(st.ok()) << st.to_string();
ASSERT_EQ(response, R"({"ok":true})");
diff --git a/be/test/ai/ai_request_executor_test.cpp
b/be/test/ai/ai_request_executor_test.cpp
new file mode 100644
index 00000000000..0439221a982
--- /dev/null
+++ b/be/test/ai/ai_request_executor_test.cpp
@@ -0,0 +1,1087 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "exprs/function/ai/ai_request_executor.h"
+
+#include <event2/http.h>
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <chrono>
+#include <condition_variable>
+#include <ctime>
+#include <functional>
+#include <future>
+#include <limits>
+#include <mutex>
+#include <set>
+#include <string>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include "core/data_type/data_type_jsonb.h"
+#include "core/value/jsonb_value.h"
+#include "exprs/aggregate/aggregate_function_ai_agg.h"
+#include "exprs/function/ai/ai_adapter.h"
+#include "exprs/function/ai/ai_summarize.h"
+#include "exprs/function/ai/embed.h"
+#include "runtime/exec_env.h"
+#include "runtime/query_context.h"
+#include "service/http/ev_http_server.h"
+#include "service/http/http_channel.h"
+#include "service/http/http_client.h"
+#include "service/http/http_handler.h"
+#include "service/http/http_request.h"
+#include "testutil/column_helper.h"
+#include "testutil/mock/mock_runtime_state.h"
+
+namespace doris {
+namespace {
+
+class AIResponseHandler : public HttpHandler {
+public:
+ void set_responses(std::vector<std::pair<HttpStatus, std::string>> values)
{
+ std::lock_guard lock(mutex);
+ responses = std::move(values);
+ response_index = 0;
+ }
+
+ void handle(HttpRequest* request) override {
+ const auto callback = on_request;
+ if (callback) {
+ callback();
+ }
+ std::lock_guard lock(mutex);
+ char* address = nullptr;
+ uint16_t port = 0;
+ // Peer ports identify actual TCP reuse, rather than C++ object reuse.
+
evhttp_connection_get_peer(evhttp_request_get_connection(request->get_evhttp_request()),
+ &address, &port);
+ peer_ports.push_back(port);
+ bodies.push_back(request->get_request_body());
+ authorizations.push_back(request->header("Authorization"));
+ content_types.push_back(request->header("Content-Type"));
+ paths.push_back(request->uri());
+ if (close_first_connection && bodies.size() == 1) {
+
evhttp_add_header(evhttp_request_get_output_headers(request->get_evhttp_request()),
+ "Connection", "close");
+ }
+ if (!retry_after.empty()) {
+
evhttp_add_header(evhttp_request_get_output_headers(request->get_evhttp_request()),
+ "Retry-After", retry_after.c_str());
+ }
+ const auto& [status, body] = responses[std::min(response_index++,
responses.size() - 1)];
+ if (truncate_response || (truncate_first_response && bodies.size() ==
1)) {
+
evhttp_add_header(evhttp_request_get_output_headers(request->get_evhttp_request()),
+ "Content-Length", std::to_string(body.size() +
10).c_str());
+
evhttp_add_header(evhttp_request_get_output_headers(request->get_evhttp_request()),
+ "Connection", "close");
+ }
+ HttpChannel::send_reply(request, status, body);
+ }
+
+ std::function<void()> on_request;
+ bool truncate_response = false;
+ bool truncate_first_response = false;
+ std::string retry_after;
+ bool close_first_connection = false;
+ std::mutex mutex;
+ size_t response_index = 0;
+ std::vector<std::pair<HttpStatus, std::string>> responses;
+ std::vector<uint16_t> peer_ports;
+ std::vector<std::string> bodies;
+ std::vector<std::string> authorizations;
+ std::vector<std::string> content_types;
+ std::vector<std::string> paths;
+};
+
+class AIRequestExecutorTest : public testing::Test {
+protected:
+ void SetUp() override {
+ auto options = create_fake_query_options();
+ options.__set_query_timeout(10);
+ options.__set_execution_timeout(10);
+ options.__set_ai_context_window_size(1);
+ options.__set_embed_max_batch_size(1);
+ query_ctx = MockQueryContext::create(TUniqueId(),
ExecEnv::GetInstance(), options);
+ TQueryGlobals globals;
+ state = std::make_unique<RuntimeState>(TUniqueId(), 0, options,
globals, nullptr,
+ query_ctx.get());
+ context = FunctionContext::create_context(state.get(), {}, {});
+ ASSERT_TRUE(server.register_handler(POST, "/ai", &handler));
+ ASSERT_TRUE(server.register_handler(POST, "/other", &handler));
+ static_cast<void>(server.start());
+ ASSERT_NE(server.get_real_port(), 0);
+ config.endpoint = "http://127.0.0.1:"; +
std::to_string(server.get_real_port()) + "/ai";
+ config.provider_type = "OPENAI";
+ config.model_name = "test-model";
+ config.api_key = "test-key";
+ config.max_retries = 3;
+ config.retry_delay_second = 0;
+ adapter = std::make_shared<OpenAIAdapter>();
+ adapter->init(config);
+ }
+
+ void use_fake_clock() {
+ fake_now = std::chrono::steady_clock::now();
+ executor._now = [&] { return fake_now; };
+ executor._sleep = [&](auto duration) { fake_now += duration; };
+ }
+
+ void replay_legacy_authentication_retry_policy() {
+ HttpClient legacy_client;
+ for (int attempt = 0; attempt <= config.max_retries; ++attempt) {
+ ASSERT_TRUE(legacy_client.init(config.endpoint, false).ok());
+ ASSERT_TRUE(adapter->set_authentication(&legacy_client).ok());
+ std::string body;
+ ASSERT_TRUE(legacy_client.execute_post_request("{}", &body).ok());
+ ASSERT_EQ(legacy_client.get_http_status(), 401);
+ if (attempt != config.max_retries) {
+ std::this_thread::sleep_for(std::chrono::seconds(1));
+ }
+ }
+ }
+
+ std::chrono::steady_clock::time_point fake_now;
+
+ Status execute(const std::string& body, std::string& response) {
+ return executor.execute(body, response, config, *adapter,
query_ctx.get());
+ }
+
+ void set_resource() {
+ TAIResource resource;
+ resource.endpoint = config.endpoint;
+ resource.provider_type = config.provider_type;
+ resource.model_name = config.model_name;
+ resource.api_key = config.api_key;
+ resource.max_retries = config.max_retries;
+ resource.retry_delay_second = config.retry_delay_second;
+ query_ctx->set_ai_resources(std::map<std::string, TAIResource>
{{"resource", resource}});
+ }
+
+ size_t connection_count() {
+ std::lock_guard lock(handler.mutex);
+ return std::set<uint16_t>(handler.peer_ports.begin(),
handler.peer_ports.end()).size();
+ }
+
+ void verify_concurrent_requests(bool share_executor) {
+ handler.set_responses(
+ {{HttpStatus::OK,
R"({"data":[{"index":0,"embedding":[0.5,1.0]}]})"}});
+ // Separate server event loops make simultaneous arrival deterministic.
+ EvHttpServer second_server {0};
+ ASSERT_TRUE(second_server.register_handler(POST, "/ai", &handler));
+ second_server.start();
+ AIResource second_config = config;
+ second_config.endpoint =
+ "http://127.0.0.1:"; +
std::to_string(second_server.get_real_port()) + "/ai";
+ auto second_adapter = std::make_shared<OpenAIAdapter>();
+ second_adapter->init(second_config);
+ auto second_context = context->clone();
+ if (share_executor) {
+ auto shared_executor = std::make_shared<AIRequestExecutor>();
+ context->set_function_state(FunctionContext::THREAD_LOCAL,
shared_executor);
+ second_context->set_function_state(FunctionContext::THREAD_LOCAL,
shared_executor);
+ }
+ std::mutex gate_mutex;
+ std::condition_variable gate;
+ int arrivals = 0;
+ bool both_arrived = true;
+ handler.on_request = [&] {
+ std::unique_lock lock(gate_mutex);
+ ++arrivals;
+ gate.notify_all();
+ both_arrived &=
+ gate.wait_for(lock, std::chrono::seconds(2), [&] { return
arrivals == 2; });
+ };
+ auto send = [&](FunctionContext* ctx, const AIResource& resource,
+ std::shared_ptr<AIAdapter> provider) {
+ std::vector<std::vector<float>> result;
+ return FunctionEmbed()._execute_prebuilt_embedding_request(
+ R"({"input":["first"]})", result, 1, resource, provider,
ctx);
+ };
+ auto first = std::async(std::launch::async, send, context.get(),
config, adapter);
+ auto second = std::async(std::launch::async, send,
second_context.get(), second_config,
+ second_adapter);
+ EXPECT_TRUE(first.get().ok());
+ EXPECT_TRUE(second.get().ok());
+ EXPECT_TRUE(both_arrived);
+ EXPECT_EQ(context->get_function_state(FunctionContext::THREAD_LOCAL) ==
+
second_context->get_function_state(FunctionContext::THREAD_LOCAL),
+ share_executor);
+ EXPECT_EQ(connection_count(), 2);
+ handler.on_request = {};
+ }
+
+ // The server stops before the handler and its recorded requests are
destroyed.
+ AIRequestExecutor executor;
+ AIResponseHandler handler;
+ EvHttpServer server {0};
+ std::shared_ptr<QueryContext> query_ctx;
+ std::unique_ptr<RuntimeState> state;
+ std::unique_ptr<FunctionContext> context;
+ AIResource config {TAIResource {}};
+ std::shared_ptr<AIAdapter> adapter;
+};
+
+TEST_F(AIRequestExecutorTest, ReusesConnectionAcrossTenRequests) {
+ handler.set_responses({{HttpStatus::OK, R"({"result":"ok"})"}});
+ for (int i = 0; i < 10; ++i) {
+ std::string response;
+ ASSERT_TRUE(execute(std::to_string(i), response).ok());
+ EXPECT_EQ(response, R"({"result":"ok"})");
+ }
+ std::lock_guard lock(handler.mutex);
+ ASSERT_EQ(handler.peer_ports.size(), 10);
+ EXPECT_EQ(std::set<uint16_t>(handler.peer_ports.begin(),
handler.peer_ports.end()).size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest,
CachedClientsDiscardRequestOptionsAfterSuccessAndFailure) {
+ class InspectingAdapter : public OpenAIAdapter {
+ public:
+ Status set_authentication(HttpClient* client) const override {
+ saved_client = client;
+ RETURN_IF_ERROR(OpenAIAdapter::set_authentication(client));
+ curl_easy_setopt(client->_curl, CURLOPT_PRIVATE, "request state");
+ return Status::OK();
+ }
+ mutable HttpClient* saved_client = nullptr;
+ } inspecting_adapter;
+ inspecting_adapter.init(config);
+ config.max_retries = 1;
+ for (const auto code :
+ {HttpStatus::OK, HttpStatus::UNAUTHORIZED,
HttpStatus::SERVICE_UNAVAILABLE}) {
+ handler.set_responses({{code, "response"}});
+ std::string response;
+ const auto status = executor.execute("request body", response, config,
inspecting_adapter,
+ query_ctx.get());
+ EXPECT_EQ(status.ok(), code == HttpStatus::OK);
+ ASSERT_NE(inspecting_adapter.saved_client, nullptr);
+ char* request_state = nullptr;
+ ASSERT_EQ(curl_easy_getinfo(inspecting_adapter.saved_client->_curl,
CURLINFO_PRIVATE,
+ &request_state),
+ CURLE_OK);
+ EXPECT_EQ(request_state, nullptr);
+ EXPECT_EQ(inspecting_adapter.saved_client->_header_list, nullptr);
+ }
+ EXPECT_EQ(connection_count(), 1);
+ for (const auto& body : handler.bodies) {
+ EXPECT_EQ(body, "request body");
+ }
+}
+
+TEST_F(AIRequestExecutorTest,
ScalarReusesConnectionAcrossBatchesAndBlocksWithNulls) {
+ handler.set_responses(
+ {{HttpStatus::OK,
R"({"choices":[{"message":{"content":"[\"summary\"]"}}]})"}});
+ FunctionAISummarize function;
+ for (int block_index = 0; block_index < 5; ++block_index) {
+ Block block;
+
block.insert({ColumnHelper::create_column<DataTypeString>({"resource"}),
+ std::make_shared<DataTypeString>(), "resource"});
+ block.insert({ColumnHelper::create_nullable_column<DataTypeString>(
+ {"first", "ignored", "second"}, {0, 1, 0}),
+ make_nullable(std::make_shared<DataTypeString>()),
"input"});
+ block.insert({nullptr,
make_nullable(std::make_shared<DataTypeString>()), "result"});
+ ASSERT_TRUE(function.execute(context.get(), block, {0, 1}, 2, 3,
config, adapter).ok());
+ const auto& result = assert_cast<const
ColumnNullable&>(*block.get_by_position(2).column);
+ ASSERT_EQ(result.size(), 3);
+ EXPECT_EQ(result.get_null_map_data()[1], 1);
+ EXPECT_EQ(result.get_data_at(0).to_string(), "summary");
+ EXPECT_EQ(result.get_data_at(2).to_string(), "summary");
+ }
+ ASSERT_EQ(handler.bodies.size(), 10);
+ EXPECT_EQ(connection_count(), 1);
+ EXPECT_NE(handler.bodies[0].find("first"), std::string::npos);
+ EXPECT_NE(handler.bodies[1].find("second"), std::string::npos);
+ std::weak_ptr<void> owner = context->_thread_local_fn_state;
+ ASSERT_TRUE(function.close(context.get(),
FunctionContext::THREAD_LOCAL).ok());
+ EXPECT_TRUE(owner.expired());
+}
+
+TEST_F(AIRequestExecutorTest,
EmbedReusesConnectionAcrossBatchesAndBlocksWithNulls) {
+ handler.set_responses({{HttpStatus::OK,
R"({"data":[{"index":0,"embedding":[0.5,1.0]}]})"}});
+ FunctionEmbed function;
+ for (int block_index = 0; block_index < 5; ++block_index) {
+ Block block;
+
block.insert({ColumnHelper::create_column<DataTypeString>({"resource"}),
+ std::make_shared<DataTypeString>(), "resource"});
+ block.insert({ColumnHelper::create_nullable_column<DataTypeString>(
+ {"first", "ignored", "second"}, {0, 1, 0}),
+ make_nullable(std::make_shared<DataTypeString>()),
"input"});
+ block.insert({nullptr,
+ make_nullable(std::make_shared<DataTypeArray>(
+
make_nullable(std::make_shared<DataTypeFloat32>()))),
+ "result"});
+ ASSERT_TRUE(function.execute(context.get(), block, {0, 1}, 2, 3,
config, adapter).ok());
+ const auto& result = assert_cast<const
ColumnNullable&>(*block.get_by_position(2).column);
+ ASSERT_EQ(result.size(), 3);
+ EXPECT_EQ(result.get_null_map_data()[1], 1);
+ const auto& arrays = assert_cast<const
ColumnArray&>(result.get_nested_column());
+ EXPECT_EQ(arrays.get_offsets()[0], 2);
+ EXPECT_EQ(arrays.get_offsets()[1], 2);
+ EXPECT_EQ(arrays.get_offsets()[2], 4);
+ }
+ ASSERT_EQ(handler.bodies.size(), 10);
+ EXPECT_EQ(connection_count(), 1);
+}
+
+TEST_F(AIRequestExecutorTest,
MultimodalEmbedReusesConnectionAcrossBatchesAndBlocks) {
+ handler.set_responses({{HttpStatus::OK,
R"({"data":[{"index":0,"embedding":[0.5,1.0]}]})"}});
+ config.provider_type = "VOYAGE";
+ config.model_name = "voyage-multimodal-3";
+ adapter = std::make_shared<VoyageAIAdapter>();
+ adapter->init(config);
+ FunctionEmbed function;
+ for (int block_index = 0; block_index < 5; ++block_index) {
+ auto input = ColumnString::create();
+ for (const auto& file :
+
{R"({"content_type":"image/png","uri":"https://example.com/a.png"})",
+
R"({"content_type":"video/mp4","uri":"https://example.com/b.mp4"})"}) {
+ JsonBinaryValue value;
+ ASSERT_TRUE(value.from_json_string(file).ok());
+ input->insert_data(value.value(), value.size());
+ }
+ Block block;
+
block.insert({ColumnHelper::create_column<DataTypeString>({"resource"}),
+ std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(input), std::make_shared<DataTypeJsonb>(),
"input"});
+ block.insert({nullptr,
+ std::make_shared<DataTypeArray>(
+
make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+ ASSERT_TRUE(function.execute(context.get(), block, {0, 1}, 2, 2,
config, adapter).ok());
+ const auto& arrays = assert_cast<const
ColumnArray&>(*block.get_by_position(2).column);
+ ASSERT_EQ(arrays.size(), 2);
+ EXPECT_EQ(arrays.get_offsets()[0], 2);
+ EXPECT_EQ(arrays.get_offsets()[1], 4);
+ }
+ ASSERT_EQ(handler.bodies.size(), 10);
+ EXPECT_EQ(connection_count(), 1);
+ EXPECT_NE(handler.bodies[0].find("image_url"), std::string::npos);
+ EXPECT_NE(handler.bodies[1].find("video_url"), std::string::npos);
+}
+
+TEST_F(AIRequestExecutorTest, ColdAndWarmConnectionsForFastAndDelayedMock) {
+ handler.set_responses({{HttpStatus::OK, "ok"}});
+ for (int delay_ms : {0, 5}) {
+ handler.on_request = [delay_ms] {
+ std::this_thread::sleep_for(std::chrono::milliseconds(delay_ms));
+ };
+ auto measure = [&](bool reuse) {
+ const size_t first_request = handler.peer_ports.size();
+ const auto start = std::chrono::steady_clock::now();
+ AIRequestExecutor warm;
+ for (int i = 0; i < 10; ++i) {
+ AIRequestExecutor cold;
+ std::string response;
+ ASSERT_TRUE((reuse ? warm : cold)
+ .execute("{}", response, config, *adapter,
query_ctx.get())
+ .ok());
+ EXPECT_EQ(response, "ok");
+ }
+ const auto elapsed =
std::chrono::duration_cast<std::chrono::microseconds>(
+ std::chrono::steady_clock::now() -
start)
+ .count();
+ const size_t connections =
+ std::set<uint16_t>(handler.peer_ports.begin() +
first_request,
+ handler.peer_ports.end())
+ .size();
+ EXPECT_EQ(connections, reuse ? 1 : 10);
+ // Diagnostic timings from an ASAN local mock, not a production
speedup claim.
+ const std::string label =
+ std::string(reuse ? "warm_" : "cold_") +
std::to_string(delay_ms) + "ms";
+ RecordProperty(label + "_connections",
std::to_string(connections));
+ RecordProperty(label + "_total_us", std::to_string(elapsed));
+ };
+ measure(false);
+ measure(true);
+ }
+}
+
+TEST_F(AIRequestExecutorTest, AggregateReusesConnectionAcrossFlushesAndResets)
{
+ handler.set_responses({{HttpStatus::OK,
R"({"choices":[{"message":{"content":"summary"}}]})"}});
+ set_resource();
+ AggregateFunctionAIAggData data;
+ data.set_query_context(query_ctx.get(),
std::make_shared<AIRequestExecutor>());
+ // A tiny context window forces add() to summarize the preceding context.
+ for (int block_index = 0; block_index < 5; ++block_index) {
+ data.prepare(StringRef("resource", 8), StringRef("summarize", 9));
+ data.add(StringRef("first", 5));
+ data.add(StringRef("second", 6));
+ EXPECT_EQ(data._execute_task(), "summary");
+ data.reset();
+ }
+ ASSERT_EQ(handler.bodies.size(), 10);
+ EXPECT_EQ(connection_count(), 1);
+}
+
+TEST_F(AIRequestExecutorTest,
AggregateGroupsShareConnectionsRatherThanRetainingOnePerGroup) {
+ handler.set_responses({{HttpStatus::OK,
R"({"choices":[{"message":{"content":"summary"}}]})"}});
+ set_resource();
+ AggregateFunctionAIAgg function({std::make_shared<DataTypeString>(),
+ std::make_shared<DataTypeString>(),
+ std::make_shared<DataTypeString>()});
+ function.set_query_context(query_ctx.get());
+ std::vector<std::unique_ptr<char[]>> groups;
+ for (int i = 0; i < 10; ++i) {
+ auto group = std::make_unique<char[]>(function.size_of_data());
+ function.create(group.get());
+ auto& data =
*reinterpret_cast<AggregateFunctionAIAggData*>(group.get());
+ data.prepare(StringRef("resource", 8), StringRef("summarize", 9));
+ data.add(StringRef("hello", 5));
+ EXPECT_EQ(data._execute_task(), "summary");
+ groups.push_back(std::move(group));
+ }
+ EXPECT_EQ(handler.bodies.size(), 10);
+ EXPECT_EQ(connection_count(), 1);
+ for (auto& group : groups) {
+ function.destroy(group.get());
+ }
+}
+
+TEST_F(AIRequestExecutorTest, RebuildsConnectionAfterServerClose) {
+ handler.close_first_connection = true;
+ handler.set_responses({{HttpStatus::OK, "ok"}});
+ for (int i = 0; i < 3; ++i) {
+ std::string response;
+ ASSERT_TRUE(execute("{}", response).ok());
+ EXPECT_EQ(response, "ok");
+ }
+ EXPECT_EQ(connection_count(), 2);
+}
+
+TEST_F(AIRequestExecutorTest, ResetsEndpointHeadersAndPayloadBetweenRequests) {
+ handler.set_responses({{HttpStatus::OK, "ok"}});
+ std::string response;
+ ASSERT_TRUE(execute("first", response).ok());
+ config.endpoint.replace(config.endpoint.size() - 3, 3, "/other");
+ config.provider_type = "LOCAL";
+ config.api_key.clear();
+ adapter = std::make_shared<LocalAdapter>();
+ adapter->init(config);
+ ASSERT_TRUE(execute("second", response).ok());
+ EXPECT_EQ(handler.bodies, (std::vector<std::string> {"first", "second"}));
+ EXPECT_EQ(handler.authorizations, (std::vector<std::string> {"Bearer
test-key", ""}));
+ EXPECT_EQ(handler.paths, (std::vector<std::string> {"/ai", "/other"}));
+ EXPECT_EQ(connection_count(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, FailedRequestDoesNotPoisonTheNextRequest) {
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"},
{HttpStatus::OK, "ok"}});
+ config.max_retries = 0;
+ std::string response;
+ EXPECT_FALSE(execute("first", response).ok());
+ ASSERT_TRUE(execute("second", response).ok());
+ EXPECT_EQ(response, "ok");
+ EXPECT_EQ(connection_count(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, IndependentClonedContextsCanSendConcurrently) {
+ verify_concurrent_requests(false);
+}
+
+TEST_F(AIRequestExecutorTest,
SharedExecutorLoansSeparateClientsForConcurrentRequests) {
+ verify_concurrent_requests(true);
+}
+
+TEST_F(AIRequestExecutorTest,
CancellationStopsInFlightWithoutRetriesAndDetachesQuery) {
+ // Match HttpClientTest.abort_in_flight_request: stack-trace symbolization
in
+ // BE_TEST can take seconds after curl has already aborted the transfer.
+ query_ctx->_timeout_second = 60;
+ handler.set_responses({{HttpStatus::OK, "ok"}});
+ std::promise<void> arrived;
+ std::promise<void> release;
+ auto released = release.get_future().share();
+ handler.on_request = [&] {
+ arrived.set_value();
+ released.wait();
+ };
+ auto pending = std::async(std::launch::async, [&] {
+ std::string response;
+ return execute("cancelled", response);
+ });
+ arrived.get_future().wait();
+ query_ctx->_exec_status.update(Status::Cancelled("test cancellation"));
+ // The server cannot reply until after we observe cancellation. This
checks the
+ // in-flight abort callback rather than merely checking status after a
reply.
+ const auto completion = pending.wait_for(std::chrono::seconds(15));
+ release.set_value();
+ EXPECT_EQ(completion, std::future_status::ready);
+ EXPECT_TRUE(pending.get().is<ErrorCode::CANCELLED>());
+ // Wait for the mock handler to complete before changing its callback.
+ server.stop();
+ server.join();
+ EXPECT_EQ(handler.bodies.size(), 1);
+ // A reused lease must bind a healthy query rather than retain
cancellation.
+ handler.on_request = {};
+ EvHttpServer next_server {0};
+ ASSERT_TRUE(next_server.register_handler(POST, "/ai", &handler));
+ next_server.start();
+ config.endpoint = "http://127.0.0.1:"; +
std::to_string(next_server.get_real_port()) + "/ai";
+ auto next_options = create_fake_query_options();
+ next_options.__set_query_timeout(10);
+ auto next_query = MockQueryContext::create(TUniqueId(),
ExecEnv::GetInstance(), next_options);
+ std::string response;
+ ASSERT_TRUE(executor.execute("healthy", response, config, *adapter,
next_query.get()).ok());
+ EXPECT_EQ(response, "ok");
+}
+
+TEST_F(AIRequestExecutorTest, TimeoutDoesNotPoisonClientForTheNextQuery) {
+ handler.set_responses({{HttpStatus::OK, "ok"}});
+ handler.on_request = [] {
std::this_thread::sleep_for(std::chrono::milliseconds(1500)); };
+ query_ctx->_timeout_second = 1;
+ config.max_retries = 0;
+ std::string response;
+ EXPECT_FALSE(execute("timeout", response).ok());
+ // The first handler will have finished by the time the second is
dispatched.
+ handler.on_request = {};
+ auto next_options = create_fake_query_options();
+ next_options.__set_query_timeout(10);
+ auto next_query = MockQueryContext::create(TUniqueId(),
ExecEnv::GetInstance(), next_options);
+ ASSERT_TRUE(executor.execute("healthy", response, config, *adapter,
next_query.get()).ok());
+ EXPECT_EQ(response, "ok");
+ EXPECT_EQ(handler.bodies.size(), 2);
+ EXPECT_EQ(connection_count(), 2);
+}
+
+TEST_F(AIRequestExecutorTest, SendsPayloadAndProviderHeaders) {
+ handler.set_responses({{HttpStatus::OK, R"({"result":"ok"})"}});
+ std::string response = "previous response";
+ ASSERT_TRUE(execute(R"({"input":"hello"})", response).ok());
+ EXPECT_EQ(response, R"({"result":"ok"})");
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies,
std::vector<std::string>({R"({"input":"hello"})"}));
+ EXPECT_EQ(handler.authorizations, std::vector<std::string>({"Bearer
test-key"}));
+ EXPECT_EQ(handler.content_types,
std::vector<std::string>({"application/json"}));
+}
+
+TEST_F(AIRequestExecutorTest, RetryReturnsOnlyTheSuccessfulAttemptBody) {
+ config.max_retries = 1;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE,
R"({"error":"busy"})"},
+ {HttpStatus::OK, R"({"result":"ok"})"}});
+ std::string response;
+ ASSERT_TRUE(execute("{}", response).ok());
+ EXPECT_EQ(response, R"({"result":"ok"})");
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies, std::vector<std::string>({"{}", "{}"}));
+ EXPECT_EQ(handler.peer_ports[0], handler.peer_ports[1]);
+}
+
+TEST_F(AIRequestExecutorTest, ZeroRetriesStillSendsSuccessfulInitialRequest) {
+ config.max_retries = 0;
+ handler.set_responses({{HttpStatus::OK, R"({"result":"ok"})"}});
+ std::string response;
+ ASSERT_TRUE(execute("{}", response).ok());
+ EXPECT_EQ(response, R"({"result":"ok"})");
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, ZeroRetriesReturnsInitialRequestFailure) {
+ config.max_retries = 0;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE,
R"({"error":"busy"})"}});
+ std::string response;
+ const auto status = execute("{}", response);
+ EXPECT_FALSE(status.ok());
+ EXPECT_NE(status.to_string().find("503"), std::string::npos);
+ EXPECT_EQ(response, R"({"error":"busy"})");
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, ZeroRetriesDoesNotWaitAfterFailure) {
+ config.max_retries = 0;
+ config.retry_delay_second = 3;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"}});
+ const auto started = std::chrono::steady_clock::now();
+ std::string response;
+ const auto status = execute("{}", response);
+ EXPECT_LT(std::chrono::steady_clock::now() - started,
std::chrono::seconds(2));
+ EXPECT_TRUE(status.is<ErrorCode::HTTP_ERROR>()) << status;
+ EXPECT_EQ(response, "busy");
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, FinalFailureOnlyWaitsBeforeTheRetry) {
+ config.max_retries = 1;
+ config.retry_delay_second = 2;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"}});
+ const auto started = std::chrono::steady_clock::now();
+ std::string response;
+ const auto status = execute("{}", response);
+ const auto elapsed = std::chrono::steady_clock::now() - started;
+ EXPECT_GE(elapsed, std::chrono::seconds(2));
+ EXPECT_LT(elapsed, std::chrono::milliseconds(3500));
+ EXPECT_TRUE(status.is<ErrorCode::HTTP_ERROR>()) << status;
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 2);
+ ASSERT_EQ(handler.peer_ports.size(), 2);
+ EXPECT_EQ(handler.peer_ports[0], handler.peer_ports[1]);
+}
+
+TEST_F(AIRequestExecutorTest, CancellationInterruptsRetryWait) {
+ query_ctx->_timeout_second = 60;
+ config.max_retries = 1;
+ config.retry_delay_second = 3;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"}});
+ std::promise<void> arrived;
+ std::once_flag first_request;
+ handler.on_request = [&] { std::call_once(first_request, [&] {
arrived.set_value(); }); };
+ auto pending = std::async(std::launch::async, [&] {
+ std::string response;
+ return execute("{}", response);
+ });
+ arrived.get_future().wait();
+ // Let the immediate 503 complete and leave the request waiting to retry.
+ EXPECT_EQ(pending.wait_for(std::chrono::milliseconds(300)),
std::future_status::timeout);
+ query_ctx->_exec_status.update(Status::Cancelled("cancel during AI retry
wait"));
+ EXPECT_EQ(pending.wait_for(std::chrono::seconds(1)),
std::future_status::ready);
+ EXPECT_TRUE(pending.get().is<ErrorCode::CANCELLED>());
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, QueryTimeoutInterruptsRetryWait) {
+ query_ctx->_timeout_second = 1;
+ config.max_retries = 1;
+ config.retry_delay_second = 3;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"}});
+ const auto started = std::chrono::steady_clock::now();
+ std::string response;
+ const auto status = execute("{}", response);
+ EXPECT_LT(std::chrono::steady_clock::now() - started,
std::chrono::seconds(2));
+ EXPECT_TRUE(status.is<ErrorCode::TIMEOUT>()) << status;
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, MaxRetriesCountsRetriesAfterInitialAttempt) {
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE,
R"({"error":"busy"})"}});
+ std::string response;
+ const auto status = execute("{}", response);
+ ASSERT_FALSE(status.ok());
+ EXPECT_NE(status.to_string().find("503"), std::string::npos);
+ EXPECT_NE(status.to_string().find("busy"), std::string::npos);
+ EXPECT_EQ(response, R"({"error":"busy"})");
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 4);
+}
+
+TEST_F(AIRequestExecutorTest, AuthenticationErrorsFailWithoutRetry) {
+ handler.set_responses({{HttpStatus::UNAUTHORIZED, R"({"error":"invalid
key"})"}});
+ std::string response;
+ EXPECT_FALSE(execute("{}", response).ok());
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest,
AuthenticationSetupFailureDoesNotReadResponseOrRetry) {
+ class FailedAuthenticationAdapter : public OpenAIAdapter {
+ public:
+ Status set_authentication(HttpClient*) const override {
+ return Status::InternalError("authentication setup failed");
+ }
+ } failed_adapter;
+ std::string response;
+ AIRequestStats stats;
+ const auto status =
+ executor.execute("{}", response, config, failed_adapter,
query_ctx.get(), &stats);
+ EXPECT_FALSE(status.ok());
+ EXPECT_NE(status.to_string().find("authentication setup failed"),
std::string::npos);
+ EXPECT_EQ(stats.attempts, 0);
+ EXPECT_TRUE(handler.bodies.empty());
+}
+
+TEST_F(AIRequestExecutorTest, PermanentHttpErrorsFailWithoutRetry) {
+ for (auto code : {400, 401, 403, 404, 422}) {
+ handler.set_responses({{static_cast<HttpStatus>(code), "invalid
request"}});
+ const auto before = handler.bodies.size();
+ std::string response;
+ EXPECT_FALSE(execute("{}", response).ok());
+ EXPECT_EQ(handler.bodies.size(), before + 1) << code;
+ }
+}
+
+TEST_F(AIRequestExecutorTest, RetryAfterExceedingQueryBudgetDoesNotSendAgain) {
+ handler.retry_after = "60";
+ handler.set_responses({{HttpStatus::TOO_MANY_REQUESTS, "rate limited"}});
+ const auto started = std::chrono::steady_clock::now();
+ std::string response;
+ const auto status = execute("{}", response);
+ EXPECT_TRUE(status.is<ErrorCode::TIMEOUT>()) << status;
+ EXPECT_LT(std::chrono::steady_clock::now() - started,
std::chrono::seconds(2));
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, TruncatedAuthenticationResponseDoesNotRetry) {
+ handler.truncate_response = true;
+ handler.set_responses({{HttpStatus::UNAUTHORIZED, "invalid key"}});
+ std::string response;
+ AIRequestStats stats;
+ EXPECT_FALSE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 1);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest,
TruncatedRateLimitResponsesUseConfiguredRetryBudget) {
+ query_ctx->_timeout_second = 60;
+ handler.truncate_response = true;
+ config.max_retries = 5;
+ handler.set_responses({{HttpStatus::TOO_MANY_REQUESTS, "throttled"}});
+ std::string response;
+ AIRequestStats stats;
+ EXPECT_FALSE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 6);
+ EXPECT_EQ(handler.bodies.size(), 6);
+}
+
+TEST_F(AIRequestExecutorTest, TruncatedSuccessfulResponseCanRecover) {
+ query_ctx->_timeout_second = 60;
+ handler.truncate_first_response = true;
+ handler.set_responses({{HttpStatus::OK, "partial"}, {HttpStatus::OK,
"success"}});
+ std::string response;
+ AIRequestStats stats;
+ EXPECT_TRUE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 2);
+ EXPECT_EQ(response, "success");
+ EXPECT_EQ(handler.bodies.size(), 2);
+}
+
+TEST_F(AIRequestExecutorTest, TransientHttpStatusPolicy) {
+ for (auto code : {408, 500, 502, 503, 504, 529, 429}) {
+ const auto before = handler.bodies.size();
+ config.max_retries = 1;
+ handler.set_responses({{static_cast<HttpStatus>(code), "busy"}});
+ AIRequestStats stats;
+ std::string response;
+ EXPECT_FALSE(
+ executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(handler.bodies.size(), before + 2) << code;
+ EXPECT_EQ(stats.attempts, 2);
+ }
+}
+
+TEST_F(AIRequestExecutorTest, StructuredProviderErrorsDoNotOverrideHttpStatus)
{
+ const std::vector<std::pair<std::string, std::string>> cases {
+ {"OPENAI", R"({"error":{"code":"rate_limit_exceeded"}})"},
+ {"OPENAI", R"({"error":{"type":"server_error"}})"},
+ {"ANTHROPIC", R"({"error":{"type":"overloaded_error"}})"},
+ {"GEMINI", R"({"error":{"status":"UNAVAILABLE"}})"},
+ {"GEMINI", R"({"error":{"status":"RESOURCE_EXHAUSTED"}})"}};
+ for (const auto& [provider, body] : cases) {
+ config.provider_type = provider;
+ adapter = AIAdapterFactory::create_adapter(provider);
+ adapter->init(config);
+ config.max_retries = 1;
+ handler.set_responses({{HttpStatus::BAD_REQUEST, body}});
+ const auto before = handler.bodies.size();
+ std::string response;
+ EXPECT_FALSE(execute("{}", response).ok());
+ EXPECT_EQ(handler.bodies.size(), before + 1) << provider << body;
+ }
+}
+
+TEST_F(AIRequestExecutorTest, AccountQuota429UsesRateLimitBudget) {
+ handler.set_responses(
+ {{HttpStatus::TOO_MANY_REQUESTS,
+
R"({"error":{"code":"insufficient_quota","type":"rate_limit_error"}})"}});
+ std::string response;
+ EXPECT_FALSE(execute("{}", response).ok());
+ EXPECT_EQ(handler.bodies.size(), 4);
+}
+
+TEST_F(AIRequestExecutorTest, ErrorMessageTextDoesNotControlRetryPolicy) {
+ handler.set_responses({{HttpStatus::BAD_REQUEST,
+ R"({"error":{"message":"rate_limit_exceeded
server_error"}})"}});
+ std::string response;
+ EXPECT_FALSE(execute("{}", response).ok());
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest,
PermanentFailureLatencyAndAttemptsComparedWithLegacyPolicy) {
+ handler.set_responses({{HttpStatus::UNAUTHORIZED, "invalid key"}});
+ // Replay the previous executor policy: retry every error, with a fixed
one-second wait.
+ const auto legacy_start = std::chrono::steady_clock::now();
+ replay_legacy_authentication_retry_policy();
+ const auto legacy_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(
+ std::chrono::steady_clock::now() -
legacy_start)
+ .count();
+ const auto legacy_attempts = handler.bodies.size();
+ config.retry_delay_second = 1;
+ const auto current_start = std::chrono::steady_clock::now();
+ std::string response;
+ AIRequestStats stats;
+ EXPECT_FALSE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ const auto current_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(
+ std::chrono::steady_clock::now() -
current_start)
+ .count();
+ EXPECT_EQ(legacy_attempts, 4);
+ EXPECT_EQ(stats.attempts, 1);
+ EXPECT_EQ(stats.retry_wait_ms, 0);
+ EXPECT_GE(legacy_ms, 3000);
+ EXPECT_LT(current_ms, legacy_ms);
+ RecordProperty("legacy_auth_attempts", static_cast<int>(legacy_attempts));
+ RecordProperty("current_auth_attempts", static_cast<int>(stats.attempts));
+ RecordProperty("legacy_auth_ms", static_cast<int>(legacy_ms));
+ RecordProperty("current_auth_ms", static_cast<int>(current_ms));
+}
+
+TEST_F(AIRequestExecutorTest, FixedRetryDelayDoesNotGrowAcrossAttempts) {
+ use_fake_clock();
+ config.retry_delay_second = 1;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"},
+ {HttpStatus::SERVICE_UNAVAILABLE, "busy"},
+ {HttpStatus::SERVICE_UNAVAILABLE, "busy"},
+ {HttpStatus::OK, "success"}});
+ AIRequestStats stats;
+ std::string response;
+ EXPECT_TRUE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 4);
+ EXPECT_EQ(stats.retry_wait_ms, 3000);
+ EXPECT_EQ(response, "success");
+ EXPECT_EQ(connection_count(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, MixedFailuresKeepTheConfiguredRetryInterval) {
+ use_fake_clock();
+ config.retry_delay_second = 1;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"},
+ {HttpStatus::TOO_MANY_REQUESTS, "throttled"},
+ {HttpStatus::SERVICE_UNAVAILABLE, "busy"},
+ {HttpStatus::TOO_MANY_REQUESTS, "throttled"}});
+ AIRequestStats stats;
+ std::string response;
+ EXPECT_FALSE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 4);
+ EXPECT_EQ(stats.retry_wait_ms, 3000);
+ EXPECT_EQ(handler.bodies.size(), 4);
+}
+
+TEST_F(AIRequestExecutorTest,
RateLimitAndServerErrorsShareConfiguredRetryBudget) {
+ config.max_retries = 5;
+ handler.set_responses({{HttpStatus::TOO_MANY_REQUESTS, "throttled"}});
+ AIRequestStats stats;
+ std::string response;
+ EXPECT_FALSE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 6);
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"}});
+ EXPECT_FALSE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 6);
+}
+
+TEST_F(AIRequestExecutorTest, ScalarEmbedAndAggregateShareTheErrorPolicy) {
+ set_resource();
+ AggregateFunctionAIAggData data;
+ data.set_query_context(query_ctx.get(),
std::make_shared<AIRequestExecutor>());
+ data.prepare(StringRef("resource", 8), StringRef("summarize", 9));
+ std::string response;
+ for (auto code : {400, 401, 403, 429, 503}) {
+ handler.set_responses({{static_cast<HttpStatus>(code), "error"}});
+ const auto expected = code == 429 || code == 503 ? 4 : 1;
+ auto before = handler.bodies.size();
+ EXPECT_FALSE(FunctionAISummarize()
+ .send_request_to_llm("{}", response, config,
adapter, context.get())
+ .ok());
+ EXPECT_EQ(handler.bodies.size(), before + expected) << "scalar " <<
code;
+ before = handler.bodies.size();
+ EXPECT_FALSE(FunctionEmbed()
+ .send_request_to_llm("{}", response, config,
adapter, context.get())
+ .ok());
+ EXPECT_EQ(handler.bodies.size(), before + expected) << "embed " <<
code;
+ before = handler.bodies.size();
+ EXPECT_FALSE(data.send_request_to_ai("{}", response).ok());
+ EXPECT_EQ(handler.bodies.size(), before + expected) << "aggregate " <<
code;
+ }
+}
+
+TEST_F(AIRequestExecutorTest, ParseFailuresDoNotResendSuccessfulHttpRequests) {
+ handler.set_responses({{HttpStatus::OK, "malformed json"}});
+ std::vector<std::vector<float>> results;
+ EXPECT_FALSE(FunctionEmbed()
+ ._execute_prebuilt_embedding_request("{}", results,
1, config, adapter,
+ context.get())
+ .ok());
+ EXPECT_EQ(handler.bodies.size(), 1);
+ set_resource();
+ AggregateFunctionAIAggData data;
+ data.set_query_context(query_ctx.get(),
std::make_shared<AIRequestExecutor>());
+ data.prepare(StringRef("resource", 8), StringRef("summarize", 9));
+ data.add(StringRef("hello", 5));
+ EXPECT_THROW(data._execute_task(), Exception);
+ EXPECT_EQ(handler.bodies.size(), 2);
+}
+
+TEST_F(AIRequestExecutorTest,
RetryAfterSetsWaitFloorEvenWithExplicitZeroDelay) {
+ use_fake_clock();
+ config.max_retries = 1;
+ handler.retry_after = "2";
+ handler.set_responses(
+ {{HttpStatus::TOO_MANY_REQUESTS, "throttled"}, {HttpStatus::OK,
"success"}});
+ AIRequestStats stats;
+ std::string response;
+ EXPECT_TRUE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_EQ(stats.attempts, 2);
+ EXPECT_EQ(stats.retry_wait_ms, 2000);
+}
+
+TEST_F(AIRequestExecutorTest, RetryAfterHttpDateSetsWaitFloor) {
+ use_fake_clock();
+ config.max_retries = 1;
+ const auto unix_now = std::time(nullptr);
+ const auto date = unix_now + 3;
+ std::tm utc {};
+ gmtime_r(&date, &utc);
+ char header[64];
+ std::strftime(header, sizeof(header), "%a, %d %b %Y %H:%M:%S GMT", &utc);
+ handler.retry_after = header;
+ handler.set_responses(
+ {{HttpStatus::TOO_MANY_REQUESTS, "throttled"}, {HttpStatus::OK,
"success"}});
+ AIRequestStats stats;
+ std::string response;
+ EXPECT_TRUE(executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats).ok());
+ EXPECT_GE(stats.retry_wait_ms, 2000);
+ EXPECT_LE(stats.retry_wait_ms, 3000);
+}
+
+TEST_F(AIRequestExecutorTest, OneDeadlineCoversEveryAttemptAndWait) {
+ use_fake_clock();
+ config.retry_delay_second = 4;
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"}});
+ AIRequestStats stats;
+ std::string response;
+ const auto status = executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats);
+ EXPECT_TRUE(status.is<ErrorCode::TIMEOUT>()) << status;
+ EXPECT_EQ(stats.attempts, 3);
+ EXPECT_EQ(stats.retry_wait_ms, 8000);
+ EXPECT_EQ(handler.bodies.size(), 3);
+}
+
+TEST_F(AIRequestExecutorTest, OversleptRetryWaitDoesNotSendAfterDeadline) {
+ use_fake_clock();
+ config.retry_delay_second = 1;
+ executor._sleep = [&](auto) { fake_now += std::chrono::seconds(20); };
+ handler.set_responses({{HttpStatus::SERVICE_UNAVAILABLE, "busy"}});
+ std::string response;
+ AIRequestStats stats;
+ const auto status = executor.execute("{}", response, config, *adapter,
query_ctx.get(), &stats);
+ EXPECT_TRUE(status.is<ErrorCode::TIMEOUT>()) << status;
+ EXPECT_EQ(stats.attempts, 1);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, RetryAfterParsingIsDeterministic) {
+ EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms(" 2 ", 0), 2000);
+ EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms("0", 0), 0);
+ EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms("-1", 0), 0);
+ EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms("1.5", 0), 0);
+ EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms("invalid", 0), 0);
+ EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms("Wed, 21 Oct 2015
07:28:00 GMT", 1445412479),
+ 1000);
+ EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms("Wed, 21 Oct 2015
07:28:00 GMT", 1445412481),
+ 0);
+
EXPECT_EQ(AIRequestExecutor::parse_retry_after_ms("9999999999999999999999999",
0),
+ std::numeric_limits<int64_t>::max());
+}
+
+TEST_F(AIRequestExecutorTest,
TransportFailuresUseCodesRatherThanStatusStrings) {
+ for (auto code :
+ {CURLE_COULDNT_RESOLVE_HOST, CURLE_COULDNT_CONNECT, CURLE_SEND_ERROR,
CURLE_RECV_ERROR,
+ CURLE_GOT_NOTHING, CURLE_PARTIAL_FILE, CURLE_HTTP2_STREAM}) {
+ EXPECT_EQ(AIRequestExecutor::classify_curl_error(code),
AIErrorType::RETRYABLE);
+ }
+ for (auto code : {CURLE_OPERATION_TIMEDOUT, CURLE_ABORTED_BY_CALLBACK,
CURLE_OUT_OF_MEMORY,
+ CURLE_WRITE_ERROR, CURLE_PEER_FAILED_VERIFICATION,
CURLE_URL_MALFORMAT}) {
+ EXPECT_EQ(AIRequestExecutor::classify_curl_error(code),
AIErrorType::NON_RETRYABLE);
+ }
+}
+
+TEST_F(AIRequestExecutorTest, ExpiredQueryDoesNotSend) {
+ query_ctx->_timeout_second = 0;
+ std::string response;
+ const auto status = execute("{}", response);
+ EXPECT_TRUE(status.is<ErrorCode::TIMEOUT>()) << status;
+ std::lock_guard lock(handler.mutex);
+ EXPECT_TRUE(handler.bodies.empty());
+}
+
+TEST_F(AIRequestExecutorTest, LocalProviderSetsJsonHeadersWithoutAnApiKey) {
+ handler.set_responses({{HttpStatus::OK, R"({"result":"ok"})"}});
+ config.provider_type = "LOCAL";
+ config.api_key.clear();
+ adapter = std::make_shared<LocalAdapter>();
+ adapter->init(config);
+ std::string response;
+ ASSERT_TRUE(execute("{}", response).ok());
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.content_types,
std::vector<std::string>({"application/json"}));
+ EXPECT_EQ(handler.authorizations, std::vector<std::string>({""}));
+}
+
+TEST_F(AIRequestExecutorTest,
AggregateHttpModeRejectsNon200WithoutAddingRetries) {
+ handler.set_responses({{HttpStatus::CREATED, R"({"result":"created"})"}});
+ std::string response;
+ const auto status = executor.execute("{}", response, config, *adapter,
query_ctx.get());
+ EXPECT_FALSE(status.ok());
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies.size(), 1);
+}
+
+TEST_F(AIRequestExecutorTest, MissingQueryContextDoesNotSend) {
+ std::string response;
+ const auto status = executor.execute("{}", response, config, *adapter,
nullptr);
+ EXPECT_FALSE(status.ok());
+ EXPECT_NE(status.to_string().find("Query context is null"),
std::string::npos);
+ std::lock_guard lock(handler.mutex);
+ EXPECT_TRUE(handler.bodies.empty());
+}
+
+TEST_F(AIRequestExecutorTest, EmbedUsesTheSharedSenderAndParsesVectors) {
+ handler.set_responses(
+ {{HttpStatus::OK,
+
R"({"data":[{"index":0,"embedding":[0.5,1.0]},{"index":1,"embedding":[2.0,3.0]}]})"}});
+ std::vector<std::vector<float>> results;
+ const auto status = FunctionEmbed()._execute_prebuilt_embedding_request(
+ R"({"input":["first","second"]})", results, 2, config, adapter,
context.get());
+ ASSERT_TRUE(status.ok()) << status;
+ EXPECT_EQ(results, (std::vector<std::vector<float>> {{0.5F, 1.0F}, {2.0F,
3.0F}}));
+ std::lock_guard lock(handler.mutex);
+ EXPECT_EQ(handler.bodies,
std::vector<std::string>({R"({"input":["first","second"]})"}));
+}
+
+TEST_F(AIRequestExecutorTest, AggregateUsesTheSharedSenderAndParsesOneResult) {
+ handler.set_responses({{HttpStatus::OK,
R"({"choices":[{"message":{"content":"summary"}}]})"}});
+ TAIResource resource;
+ resource.endpoint = config.endpoint;
+ resource.provider_type = config.provider_type;
+ resource.model_name = config.model_name;
+ resource.api_key = config.api_key;
+ resource.max_retries = config.max_retries;
+ resource.retry_delay_second = config.retry_delay_second;
+ query_ctx->set_ai_resources(std::map<std::string, TAIResource>
{{"resource", resource}});
+ AggregateFunctionAIAggData data;
+ data.set_query_context(query_ctx.get(),
std::make_shared<AIRequestExecutor>());
+ data.prepare(StringRef("resource", 8), StringRef("summarize", 9));
+ data.add(StringRef("hello", 5));
+ EXPECT_EQ(data._execute_task(), "summary");
+ std::lock_guard lock(handler.mutex);
+ ASSERT_EQ(handler.bodies.size(), 1);
+ EXPECT_NE(handler.bodies[0].find("hello"), std::string::npos);
+}
+
+} // namespace
+} // namespace doris
diff --git a/be/test/service/http/http_client_test.cpp
b/be/test/service/http/http_client_test.cpp
index 06d23a64185..35c7f3ad40a 100644
--- a/be/test/service/http/http_client_test.cpp
+++ b/be/test/service/http/http_client_test.cpp
@@ -22,6 +22,16 @@
#include <gtest/gtest-message.h>
#include <gtest/gtest-test-part.h>
#include <netinet/in.h>
+#include <openssl/buffer.h>
+#include <openssl/pem.h>
+#include <openssl/rsa.h>
+#include <openssl/ssl.h>
+#include <openssl/x509v3.h>
+#include <poll.h>
+#include <pthread.h>
+#ifdef ADDRESS_SANITIZER
+#include <sanitizer/allocator_interface.h>
+#endif
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/stat.h>
@@ -29,11 +39,16 @@
#include <atomic>
#include <boost/algorithm/string/predicate.hpp>
+#include <cerrno>
#include <chrono>
+#include <csignal>
#include <cstring>
#include <filesystem>
+#include <memory>
+#include <mutex>
#include <thread>
+#include "common/config.h"
#include "gtest/gtest_pred_impl.h"
#include "io/fs/local_file_system.h"
#include "runtime/exec_env.h"
@@ -46,10 +61,377 @@
#include "service/http/utils.h"
#include "service/http_service.h"
#include "testutil/http_utils.h"
+#include "util/defer_op.h"
#include "util/md5.h"
namespace doris {
+TEST(HttpClientResetTest, ReleasesCopiedPayloadWithoutDestroyingHandle) {
+ HttpClient client;
+ client.reset();
+ ASSERT_TRUE(client.init("http://127.0.0.1/";).ok());
+ auto* handle = client._curl;
+ const std::string payload(4 * 1024 * 1024, 'x');
+ client.set_payload(payload);
+ client.set_content_type("application/json");
+ client.set_abort_callback([] { return false; });
+#ifdef ADDRESS_SANITIZER
+ const auto retained_bytes = __sanitizer_get_current_allocated_bytes();
+#endif
+ client.reset();
+ EXPECT_EQ(client._curl, handle);
+ EXPECT_EQ(client._header_list, nullptr);
+ EXPECT_FALSE(client._abort_callback);
+ EXPECT_EQ(client._callback, nullptr);
+#ifdef ADDRESS_SANITIZER
+ // The original payload remains alive; only libcurl's request copy is
released.
+ EXPECT_LE(__sanitizer_get_current_allocated_bytes(), retained_bytes -
payload.size());
+#endif
+ ASSERT_TRUE(client.init("http://127.0.0.1/";).ok());
+ EXPECT_EQ(client._curl, handle);
+}
+
+class HttpClientTlsTest : public testing::Test {
+protected:
+ void SetUp() override {
+ {
+ std::scoped_lock lock(*config::get_mutable_string_config_lock());
+ _original_ca_paths = config::ca_cert_file_paths;
+ }
+ EVP_PKEY* raw_key = nullptr;
+ generate_key(&raw_key);
+ ASSERT_FALSE(HasFatalFailure());
+ std::unique_ptr<EVP_PKEY, decltype(&EVP_PKEY_free)> key(raw_key,
EVP_PKEY_free);
+ std::unique_ptr<X509, decltype(&X509_free)> certificate(X509_new(),
X509_free);
+ ASSERT_NE(certificate, nullptr);
+ generate_certificate(certificate.get(), key.get());
+ ASSERT_FALSE(HasFatalFailure());
+ start_server_and_write_ca(certificate.get(), key.get());
+ ASSERT_FALSE(HasFatalFailure());
+ }
+
+ static void generate_key(EVP_PKEY** key) {
+ std::unique_ptr<EVP_PKEY_CTX, decltype(&EVP_PKEY_CTX_free)>
key_context(
+ EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, nullptr), EVP_PKEY_CTX_free);
+ ASSERT_NE(key_context, nullptr);
+ ASSERT_EQ(EVP_PKEY_keygen_init(key_context.get()), 1);
+ ASSERT_EQ(EVP_PKEY_CTX_set_rsa_keygen_bits(key_context.get(), 2048),
1);
+ ASSERT_EQ(EVP_PKEY_keygen(key_context.get(), key), 1);
+ }
+
+ static void generate_certificate(X509* certificate, EVP_PKEY* key) {
+ ASSERT_EQ(X509_set_version(certificate, 2), 1);
+ ASSERT_EQ(ASN1_INTEGER_set(X509_get_serialNumber(certificate), 1), 1);
+ ASSERT_NE(X509_gmtime_adj(X509_getm_notBefore(certificate), -60),
nullptr);
+ ASSERT_NE(X509_gmtime_adj(X509_getm_notAfter(certificate), 3600),
nullptr);
+ ASSERT_EQ(X509_set_pubkey(certificate, key), 1);
+ auto* subject = X509_get_subject_name(certificate);
+ ASSERT_EQ(X509_NAME_add_entry_by_txt(subject, "CN", MBSTRING_ASC,
+ reinterpret_cast<const unsigned
char*>("127.0.0.1"),
+ -1, -1, 0),
+ 1);
+ ASSERT_EQ(X509_set_issuer_name(certificate, subject), 1);
+ add_extension(certificate, NID_basic_constraints, "critical,CA:TRUE");
+ add_extension(certificate, NID_subject_alt_name, "IP:127.0.0.1");
+ ASSERT_FALSE(HasFatalFailure());
+ ASSERT_GT(X509_sign(certificate, key, EVP_sha256()), 0);
+ }
+
+ void start_server_and_write_ca(X509* certificate, EVP_PKEY* key) {
+ _ssl_context.reset(SSL_CTX_new(TLS_server_method()));
+ ASSERT_NE(_ssl_context, nullptr);
+ ASSERT_EQ(SSL_CTX_use_certificate(_ssl_context.get(), certificate), 1);
+ ASSERT_EQ(SSL_CTX_use_PrivateKey(_ssl_context.get(), key), 1);
+ std::unique_ptr<BIO, decltype(&BIO_free)>
cert_bio(BIO_new(BIO_s_mem()), BIO_free);
+ ASSERT_NE(cert_bio, nullptr);
+ ASSERT_EQ(PEM_write_bio_X509(cert_bio.get(), certificate), 1);
+ BUF_MEM* cert_pem = nullptr;
+ BIO_get_mem_ptr(cert_bio.get(), &cert_pem);
+ char ca_path[] = "/tmp/doris-http-client-ca-XXXXXX";
+ const int fd = mkstemp(ca_path);
+ ASSERT_GE(fd, 0);
+ _ca_path = ca_path;
+ Defer close_file {[fd] { close(fd); }};
+ ASSERT_EQ(write(fd, cert_pem->data, cert_pem->length),
+ static_cast<ssize_t>(cert_pem->length));
+ start_listener();
+ ASSERT_FALSE(HasFatalFailure());
+ }
+
+ void start_listener() {
+ _listen_fd = socket(AF_INET, SOCK_STREAM, 0);
+ ASSERT_GE(_listen_fd, 0);
+ const int flags = fcntl(_listen_fd, F_GETFL);
+ ASSERT_GE(flags, 0);
+ ASSERT_EQ(fcntl(_listen_fd, F_SETFL, flags | O_NONBLOCK), 0);
+ sockaddr_in address {};
+ address.sin_family = AF_INET;
+ address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
+ ASSERT_EQ(bind(_listen_fd, reinterpret_cast<sockaddr*>(&address),
sizeof(address)), 0);
+ ASSERT_EQ(listen(_listen_fd, 8), 0);
+ socklen_t address_length = sizeof(address);
+ ASSERT_EQ(getsockname(_listen_fd,
reinterpret_cast<sockaddr*>(&address), &address_length),
+ 0);
+ _port = ntohs(address.sin_port);
+ _server_thread = std::thread([this] { serve_https(); });
+ }
+
+ void serve_https() {
+ sigset_t blocked_signals;
+ sigemptyset(&blocked_signals);
+ sigaddset(&blocked_signals, SIGPIPE);
+ ASSERT_EQ(pthread_sigmask(SIG_BLOCK, &blocked_signals, nullptr), 0);
+ Defer drain_sigpipe {[&blocked_signals] {
drain_pending_sigpipe(blocked_signals); }};
+ while (!_stopped) {
+ _listener_ready = true;
+ // Bound shutdown latency even when no client ever connects.
+ pollfd listener {.fd = _listen_fd, .events = POLLIN, .revents = 0};
+ const int ready = poll(&listener, 1, 100);
+ if (ready == 0 || (ready < 0 && errno == EINTR)) {
+ continue;
+ }
+ if (ready < 0 || _stopped) {
+ return;
+ }
+ const int fd = accept(_listen_fd, nullptr, nullptr);
+ if (fd < 0) {
+ if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
{
+ continue;
+ }
+ return;
+ }
+ {
+ std::scoped_lock lock(_connection_mutex);
+ _connection_fd = fd;
+ if (_stopped) {
+ shutdown(fd, SHUT_RDWR);
+ }
+ }
+ Defer close_connection {[this, fd] {
+ std::scoped_lock lock(_connection_mutex);
+ _connection_fd = -1;
+ close(fd);
+ }};
+ respond_https(fd);
+ }
+ }
+
+ static void drain_pending_sigpipe(const sigset_t& blocked_signals) {
+ // sigtimedwait() is unavailable on Darwin; consume only pending
SIGPIPE.
+ sigset_t pending;
+ while (sigpending(&pending) == 0 && sigismember(&pending, SIGPIPE) ==
1) {
+ int signal = 0;
+ if (sigwait(&blocked_signals, &signal) != 0) {
+ break;
+ }
+ }
+ }
+
+ void respond_https(int fd) {
+ // Accepted sockets may inherit O_NONBLOCK on Darwin; TLS uses
blocking I/O.
+ const int flags = fcntl(fd, F_GETFL);
+ ASSERT_GE(flags, 0);
+ ASSERT_EQ(fcntl(fd, F_SETFL, flags & ~O_NONBLOCK), 0);
+ const timeval timeout {.tv_sec = 5, .tv_usec = 0};
+ ASSERT_EQ(setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout,
sizeof(timeout)), 0);
+ ASSERT_EQ(setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &timeout,
sizeof(timeout)), 0);
+ std::unique_ptr<SSL, decltype(&SSL_free)>
connection(SSL_new(_ssl_context.get()), SSL_free);
+ ASSERT_NE(connection, nullptr);
+ ASSERT_EQ(SSL_set_fd(connection.get(), fd), 1);
+ if (SSL_accept(connection.get()) != 1) {
+ return;
+ }
+ char request[4096];
+ if (SSL_read(connection.get(), request, sizeof(request)) > 0) {
+ constexpr std::string_view response =
+ "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection:
close\r\n\r\nOK";
+ SSL_write(connection.get(), response.data(), response.size());
+ }
+ SSL_shutdown(connection.get());
+ }
+
+ void TearDown() override {
+ set_ca_paths(_original_ca_paths);
+ if (!_ca_path.empty()) {
+ EXPECT_EQ(unlink(_ca_path.c_str()), 0);
+ }
+ stop_server();
+ }
+
+ void stop_server() {
+ _stopped = true;
+ {
+ std::scoped_lock lock(_connection_mutex);
+ if (_connection_fd >= 0) {
+ shutdown(_connection_fd, SHUT_RDWR);
+ }
+ }
+ if (_listen_fd >= 0) {
+ if (_server_thread.joinable()) {
+ _server_thread.join();
+ }
+ close(_listen_fd);
+ _listen_fd = -1;
+ }
+ }
+
+ static void add_extension(X509* certificate, int nid, const char* value) {
+ std::unique_ptr<X509_EXTENSION, decltype(&X509_EXTENSION_free)>
extension(
+ X509V3_EXT_conf_nid(nullptr, nullptr, nid,
const_cast<char*>(value)),
+ X509_EXTENSION_free);
+ ASSERT_NE(extension, nullptr);
+ ASSERT_EQ(X509_add_ext(certificate, extension.get(), -1), 1);
+ }
+
+ static void set_ca_paths(const std::string& paths) {
+ std::scoped_lock lock(*config::get_mutable_string_config_lock());
+ config::ca_cert_file_paths = paths;
+ }
+
+ std::string https_url(const std::string& host = "127.0.0.1") const {
+ return "https://"; + host + ":" + std::to_string(_port) + "/health";
+ }
+
+ std::unique_ptr<SSL_CTX, decltype(&SSL_CTX_free)> _ssl_context {nullptr,
SSL_CTX_free};
+ int _listen_fd = -1;
+ uint16_t _port = 0;
+ std::atomic<bool> _stopped {false};
+ std::atomic<bool> _listener_ready {false};
+ std::mutex _connection_mutex;
+ int _connection_fd = -1;
+ std::thread _server_thread;
+ std::string _ca_path;
+ std::string _original_ca_paths;
+};
+
+TEST_F(HttpClientTlsTest, tls_listener_stops_without_connections) {
+ const auto deadline = std::chrono::steady_clock::now() +
std::chrono::seconds(1);
+ while (!_listener_ready && std::chrono::steady_clock::now() < deadline) {
+ std::this_thread::yield();
+ }
+ ASSERT_TRUE(_listener_ready);
+ const auto started = std::chrono::steady_clock::now();
+ stop_server();
+ EXPECT_LT(std::chrono::steady_clock::now() - started,
std::chrono::seconds(1));
+ EXPECT_FALSE(_server_thread.joinable());
+}
+
+TEST_F(HttpClientTlsTest, tls_sigpipe_drain_consumes_only_pending_signals) {
+ std::thread worker([] {
+ sigset_t blocked_signals;
+ sigemptyset(&blocked_signals);
+ sigaddset(&blocked_signals, SIGPIPE);
+ ASSERT_EQ(pthread_sigmask(SIG_BLOCK, &blocked_signals, nullptr), 0);
+ ASSERT_EQ(pthread_kill(pthread_self(), SIGPIPE), 0);
+ sigset_t pending;
+ ASSERT_EQ(sigpending(&pending), 0);
+ ASSERT_EQ(sigismember(&pending, SIGPIPE), 1);
+ drain_pending_sigpipe(blocked_signals);
+ ASSERT_EQ(sigpending(&pending), 0);
+ EXPECT_EQ(sigismember(&pending, SIGPIPE), 0);
+ drain_pending_sigpipe(blocked_signals);
+ });
+ worker.join();
+}
+
+TEST_F(HttpClientTlsTest, tls_configured_ca_survives_reinit) {
+ set_ca_paths(_ca_path + ".missing;/tmp;" + _ca_path);
+ HttpClient client;
+ for (int request = 0; request < 2; ++request) {
+ ASSERT_TRUE(client.init(https_url()).ok());
+ client.set_timeout_ms(5000);
+ std::string response;
+ auto status = client.execute(&response);
+ ASSERT_TRUE(status.ok()) << status;
+ EXPECT_EQ(client.get_http_status(), 200);
+ EXPECT_FALSE(response.empty());
+ }
+}
+
+TEST_F(HttpClientTlsTest,
tls_ca_path_probe_is_cached_and_configuration_invalidates_it) {
+ const auto paths = _ca_path + ".missing;/tmp;" + _ca_path;
+ EXPECT_EQ(HttpClient::_get_ca_certificate_path(paths), _ca_path);
+
+ const auto moved_path = _ca_path + ".moved";
+ ASSERT_EQ(rename(_ca_path.c_str(), moved_path.c_str()), 0);
+ Defer restore_ca([&] { EXPECT_EQ(rename(moved_path.c_str(),
_ca_path.c_str()), 0); });
+ // An unchanged configuration reuses the selected path without filesystem
probes.
+ EXPECT_EQ(HttpClient::_get_ca_certificate_path(paths), _ca_path);
+ EXPECT_TRUE(HttpClient::_get_ca_certificate_path(paths + ";").empty());
+}
+
+TEST_F(HttpClientTlsTest, tls_ca_configuration_updates_on_reused_client) {
+ HttpClient client;
+ set_ca_paths(_ca_path);
+ ASSERT_TRUE(client.init(https_url()).ok());
+ client.set_timeout_ms(5000);
+ ASSERT_TRUE(client.execute().ok());
+
+ set_ca_paths(_ca_path + ".missing");
+ ASSERT_TRUE(client.init(https_url()).ok());
+ client.set_timeout_ms(5000);
+ EXPECT_FALSE(client.execute().ok());
+
+ set_ca_paths(_ca_path);
+ ASSERT_TRUE(client.init(https_url()).ok());
+ client.set_timeout_ms(5000);
+ EXPECT_TRUE(client.execute().ok());
+}
+
+TEST_F(HttpClientTlsTest, tls_configured_ca_follows_http_redirect) {
+ class RedirectHandler : public HttpHandler {
+ public:
+ explicit RedirectHandler(std::string location) :
_location(std::move(location)) {}
+ void handle(HttpRequest* req) override {
+ req->add_output_header("Location", _location.c_str());
+ HttpChannel::send_reply(req, HttpStatus::FOUND, "");
+ }
+
+ private:
+ std::string _location;
+ } handler(https_url());
+ EvHttpServer server {0};
+ ASSERT_TRUE(server.register_handler(GET, "/redirect", &handler));
+ server.start();
+ ASSERT_GT(server.get_real_port(), 0);
+ set_ca_paths(_ca_path);
+ HttpClient client;
+ ASSERT_TRUE(
+ client.init("http://127.0.0.1:"; +
std::to_string(server.get_real_port()) + "/redirect")
+ .ok());
+ client.set_timeout_ms(5000);
+ auto status = client.execute();
+ ASSERT_TRUE(status.ok()) << status;
+ EXPECT_EQ(client.get_http_status(), 200);
+ ASSERT_TRUE(client.init(https_url()).ok());
+ client.set_timeout_ms(5000);
+ status = client.execute();
+ ASSERT_TRUE(status.ok()) << status;
+ EXPECT_EQ(client.get_http_status(), 200);
+}
+
+TEST_F(HttpClientTlsTest, tls_rejects_untrusted_certificate) {
+ for (const auto& paths : {std::string(), _ca_path + ".missing"}) {
+ set_ca_paths(paths);
+ HttpClient client;
+ ASSERT_TRUE(client.init(https_url()).ok());
+ client.set_timeout_ms(5000);
+ EXPECT_FALSE(client.execute().ok());
+ const auto code = client.get_last_curl_code();
+ EXPECT_TRUE(code == CURLE_PEER_FAILED_VERIFICATION || code ==
CURLE_SSL_CACERT_BADFILE)
+ << code;
+ }
+}
+
+TEST_F(HttpClientTlsTest, tls_rejects_hostname_mismatch) {
+ set_ca_paths(_ca_path);
+ HttpClient client;
+ ASSERT_TRUE(client.init(https_url("localhost")).ok());
+ client.set_timeout_ms(5000);
+ EXPECT_FALSE(client.execute().ok());
+ EXPECT_EQ(client.get_last_curl_code(), CURLE_PEER_FAILED_VERIFICATION);
+}
+
class HttpClientTestSimpleGetHandler : public HttpHandler {
public:
void handle(HttpRequest* req) override {
@@ -130,7 +512,6 @@ static EvHttpServer* s_server = nullptr;
static int real_port = 0;
static std::string hostname = "";
static std::string address = "";
-constexpr std::string_view TMP_DIR = "./http_test_tmp";
static HttpClientTestSimpleGetHandler s_simple_get_handler;
static HttpClientTestSimplePostHandler s_simple_post_handler;
@@ -143,6 +524,40 @@ public:
HttpClientTest() {}
~HttpClientTest() override {}
+ static void prepare_batch_download_files(const std::string& remote_dir) {
+ std::string large_file = remote_dir + "/a_large_file";
+ int fd = open(large_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
+ ASSERT_TRUE(fd >= 0);
+ std::string buf = "0123456789";
+ for (int i = 0; i < 10; i++) {
+ buf += buf;
+ }
+ for (int i = 0; i < 1024; i++) {
+ ASSERT_TRUE(write(fd, buf.c_str(), buf.size()) > 0);
+ }
+ close(fd);
+
+ // create some small files.
+ for (int i = 0; i < 32; i++) {
+ std::string small_file = remote_dir + "/small_file_" +
std::to_string(i);
+ fd = open(small_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
+ ASSERT_TRUE(fd >= 0);
+ ASSERT_TRUE(write(fd, buf.c_str(), buf.size()) > 0);
+ close(fd);
+ }
+
+ // create a empty file
+ std::string empty_file = remote_dir + "/empty_file";
+ fd = open(empty_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
+ ASSERT_TRUE(fd >= 0);
+ close(fd);
+
+ empty_file = remote_dir + "/zzzz";
+ fd = open(empty_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
+ ASSERT_TRUE(fd >= 0);
+ close(fd);
+ }
+
static void SetUpTestCase() {
s_server = new EvHttpServer(0);
s_server->register_handler(GET, "/simple_get", &s_simple_get_handler);
@@ -307,6 +722,40 @@ TEST_F(HttpClientTest, header_content_md5) {
}
}
+TEST_F(HttpClientTest, content_md5_uses_final_response_headers) {
+ class RedirectHandler : public HttpHandler {
+ public:
+ void handle(HttpRequest* req) override {
+ if (req->raw_path() == "/redirect") {
+ req->add_output_header("Location", "/final");
+ req->add_output_header(HttpHeaders::CONTENT_MD5,
"redirect-md5");
+ HttpChannel::send_reply(req, HttpStatus::FOUND, "");
+ } else {
+ req->add_output_header(HttpHeaders::CONTENT_MD5, "final-md5");
+ HttpChannel::send_reply(req, "OK");
+ }
+ }
+ } handler;
+ EvHttpServer server {0};
+ ASSERT_TRUE(server.register_handler(GET, "/redirect", &handler));
+ ASSERT_TRUE(server.register_handler(GET, "/final", &handler));
+ server.start();
+ HttpClient client;
+ ASSERT_TRUE(
+ client.init("http://127.0.0.1:"; +
std::to_string(server.get_real_port()) + "/redirect")
+ .ok());
+ ASSERT_TRUE(client.execute().ok());
+ std::string header;
+ ASSERT_TRUE(client.get_response_header(HttpHeaders::CONTENT_MD5,
&header).ok());
+ EXPECT_EQ(header, "final-md5");
+ std::string md5;
+ ASSERT_TRUE(client.get_content_md5(&md5).ok());
+ EXPECT_EQ(md5, header);
+ header = "previous value";
+ ASSERT_TRUE(client.get_response_header("X-Missing-Header", &header).ok());
+ EXPECT_TRUE(header.empty());
+}
+
TEST_F(HttpClientTest, download_file_md5) {
std::string url = hostname + "/download_file";
HttpClient client;
@@ -606,10 +1055,12 @@ TEST_F(HttpClientTest, enable_http_auth) {
}
TEST_F(HttpClientTest, batch_download) {
- EXPECT_TRUE(io::global_local_filesystem()->delete_directory(TMP_DIR).ok());
- EXPECT_TRUE(io::global_local_filesystem()->create_directory(TMP_DIR).ok());
-
- std::string root_dir(TMP_DIR);
+ char temporary_directory[] = "/tmp/doris-http-batch-XXXXXX";
+ ASSERT_NE(mkdtemp(temporary_directory), nullptr);
+ std::string root_dir(temporary_directory);
+ Defer cleanup([&root_dir] {
+
EXPECT_TRUE(io::global_local_filesystem()->delete_directory(root_dir).ok());
+ });
std::string remote_related_dir = root_dir + "/source";
std::string local_dir = root_dir + "/target";
EXPECT_TRUE(io::global_local_filesystem()->create_directory(remote_related_dir).ok());
@@ -618,40 +1069,8 @@ TEST_F(HttpClientTest, batch_download) {
std::string remote_dir;
EXPECT_TRUE(io::global_local_filesystem()->canonicalize(remote_related_dir,
&remote_dir).ok());
- // 0. create dir source and prepare a large file exceeds 1MB
- {
- std::string large_file = remote_dir + "/a_large_file";
- int fd = open(large_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
- ASSERT_TRUE(fd >= 0);
- std::string buf = "0123456789";
- for (int i = 0; i < 10; i++) {
- buf += buf;
- }
- for (int i = 0; i < 1024; i++) {
- ASSERT_TRUE(write(fd, buf.c_str(), buf.size()) > 0);
- }
- close(fd);
-
- // create some small files.
- for (int i = 0; i < 32; i++) {
- std::string small_file = remote_dir + "/small_file_" +
std::to_string(i);
- fd = open(small_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
- ASSERT_TRUE(fd >= 0);
- ASSERT_TRUE(write(fd, buf.c_str(), buf.size()) > 0);
- close(fd);
- }
-
- // create a empty file
- std::string empty_file = remote_dir + "/empty_file";
- fd = open(empty_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
- ASSERT_TRUE(fd >= 0);
- close(fd);
-
- empty_file = remote_dir + "/zzzz";
- fd = open(empty_file.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
- ASSERT_TRUE(fd >= 0);
- close(fd);
- }
+ prepare_batch_download_files(remote_dir);
+ ASSERT_FALSE(HasFatalFailure());
// 1. check remote support batch download
Status st = is_support_batch_download(address);
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java
index 43890d5110a..57ce106d8f5 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java
@@ -55,7 +55,7 @@ public class AIProperties extends BaseProperties {
public static final String DEFAULT_TEMPERATURE = "-1";
public static final String DEFAULT_MAX_TOKEN = "-1";
public static final String DEFAULT_MAX_RETRIES = "3";
- public static final String DEFAULT_RETRY_DELAY_SECOND = "0";
+ public static final String DEFAULT_RETRY_DELAY_SECOND = "1";
public static final String DEFAULT_ANTHROPIC_VERSION = "2023-06-01";
public static final String DEFAULT_DIMENSIONS = "-1";
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java
b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java
index 1fbafb7473d..a978efe06f9 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java
@@ -144,7 +144,7 @@ public class AIResourceTest {
aiResource.getProperty(AIProperties.MAX_TOKEN));
Assertions.assertEquals(AIProperties.DEFAULT_MAX_RETRIES,
aiResource.getProperty(AIProperties.MAX_RETRIES));
- Assertions.assertEquals(AIProperties.DEFAULT_RETRY_DELAY_SECOND,
+ Assertions.assertEquals("1",
aiResource.getProperty(AIProperties.RETRY_DELAY_SECOND));
// with no default settings
@@ -171,6 +171,38 @@ public class AIResourceTest {
}
}
+ @Test
+ public void testRetryDelayDefault() throws DdlException {
+ AIResource aiResource = new AIResource(name);
+ aiResource.setProperties(ImmutableMap.copyOf(aiProperties));
+
+ Assertions.assertEquals("1",
aiResource.getProperty(AIProperties.RETRY_DELAY_SECOND));
+ Assertions.assertEquals(1,
aiResource.toThrift().getRetryDelaySecond());
+ }
+
+ @Test
+ public void testExplicitZeroRetryDelay() throws DdlException {
+ aiProperties.put(AIProperties.RETRY_DELAY_SECOND, "0");
+ AIResource aiResource = new AIResource(name);
+ aiResource.setProperties(ImmutableMap.copyOf(aiProperties));
+
+ Assertions.assertEquals("0",
aiResource.getProperty(AIProperties.RETRY_DELAY_SECOND));
+ Assertions.assertEquals(0,
aiResource.toThrift().getRetryDelaySecond());
+ }
+
+ @Test
+ public void testExplicitRetryDelay() throws DdlException {
+ aiProperties.put(AIProperties.RETRY_DELAY_SECOND, retryDelaySecond);
+ aiProperties.put(AIProperties.MAX_RETRIES, maxRetries);
+ AIResource aiResource = new AIResource(name);
+ aiResource.setProperties(ImmutableMap.copyOf(aiProperties));
+
+ Assertions.assertEquals("2",
aiResource.getProperty(AIProperties.RETRY_DELAY_SECOND));
+ Assertions.assertEquals(2,
aiResource.toThrift().getRetryDelaySecond());
+ Assertions.assertEquals("5",
aiResource.getProperty(AIProperties.MAX_RETRIES));
+ Assertions.assertEquals(5, aiResource.toThrift().getMaxRetries());
+ }
+
@Test
public void testAnthropic() throws UserException {
try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
diff --git a/regression-test/suites/ai_p0/test_ai_functions.groovy
b/regression-test/suites/ai_p0/test_ai_functions.groovy
index 2404573e398..7309ed84fae 100644
--- a/regression-test/suites/ai_p0/test_ai_functions.groovy
+++ b/regression-test/suites/ai_p0/test_ai_functions.groovy
@@ -80,21 +80,8 @@ suite("test_ai_functions") {
exception "AI_AGG must accept literal for the task."
}
- // Test for normal call
- // test the default resource
- try {
- sql """set query_timeout=1;"""
- sql """set default_ai_resource='${resourceName}';"""
- test {
- sql """${sentiment_query}"""
- exception "timeout"
- }
- } finally {
- sql """UNSET VARIABLE query_timeout;"""
- sql """UNSET VARIABLE default_ai_resource;"""
- }
-
- def test_query_timeout_exception = { sql_text ->
+ // Dummy API keys may produce an HTTP error before the query timeout.
+ def test_ai_request_exception = { sql_text ->
try {
sql """set query_timeout=1;"""
sql """${sql_text}"""
@@ -109,17 +96,25 @@ suite("test_ai_functions") {
}
}
- test_query_timeout_exception("SELECT AI_TRANSLATE('${resourceName}', text,
'zh-CN') FROM ${test_table_for_ai_functions};")
- test_query_timeout_exception("SELECT AI_CLASSIFY('${resourceName}', text,
label) FROM ${test_table_for_ai_functions};")
- test_query_timeout_exception("SELECT AI_EXTRACT('${resourceName}', 'this
is a test', ['task']) FROM ${test_table_for_ai_functions};")
- test_query_timeout_exception("SELECT AI_FIXGRAMMAR('${resourceName}',
text) FROM ${test_table_for_ai_functions};")
- test_query_timeout_exception("SELECT AI_GENERATE('${resourceName}',
'generate something');")
- test_query_timeout_exception("SELECT AI_SUMMARIZE('${resourceName}',
'test,test,test,test')")
- test_query_timeout_exception("SELECT AI_SENTIMENT('${resourceName}', 'this
is a test');")
- test_query_timeout_exception("SELECT AI_MASK('${resourceName}', 'this is a
test', label) FROM ${test_table_for_ai_functions};")
- test_query_timeout_exception("SELECT AI_FILTER('${resourceName}', text)
FROM ${test_table_for_ai_functions};")
- test_query_timeout_exception("SELECT AI_SIMILARITY('${resourceName}',
'this is a similarity test', text) FROM ${test_table_for_ai_functions};")
- test_query_timeout_exception("SELECT AI_AGG('${resourceName}', text, 'this
is a test') FROM ${test_table_for_ai_functions};")
+ // Exercise the default-resource form with the same error expectations.
+ try {
+ sql """set default_ai_resource='${resourceName}';"""
+ test_ai_request_exception(sentiment_query)
+ } finally {
+ sql """UNSET VARIABLE default_ai_resource;"""
+ }
+
+ test_ai_request_exception("SELECT AI_TRANSLATE('${resourceName}', text,
'zh-CN') FROM ${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT AI_CLASSIFY('${resourceName}', text,
label) FROM ${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT AI_EXTRACT('${resourceName}', 'this is a
test', ['task']) FROM ${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT AI_FIXGRAMMAR('${resourceName}', text)
FROM ${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT AI_GENERATE('${resourceName}', 'generate
something');")
+ test_ai_request_exception("SELECT AI_SUMMARIZE('${resourceName}',
'test,test,test,test')")
+ test_ai_request_exception("SELECT AI_SENTIMENT('${resourceName}', 'this is
a test');")
+ test_ai_request_exception("SELECT AI_MASK('${resourceName}', 'this is a
test', label) FROM ${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT AI_FILTER('${resourceName}', text) FROM
${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT AI_SIMILARITY('${resourceName}', 'this
is a similarity test', text) FROM ${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT AI_AGG('${resourceName}', text, 'this is
a test') FROM ${test_table_for_ai_functions};")
String embedResourceName = "embedResourceName"
try_sql("""DROP RESOURCE IF EXISTS '${embedResourceName}'""")
@@ -157,7 +152,7 @@ suite("test_ai_functions") {
sql """UNSET VARIABLE embed_max_batch_size;"""
sql """UNSET VARIABLE ai_context_window_size;"""
- test_query_timeout_exception("SELECT EMBED('${embedResourceName}', text)
FROM ${test_table_for_ai_functions};")
+ test_ai_request_exception("SELECT EMBED('${embedResourceName}', text) FROM
${test_table_for_ai_functions};")
try_sql("""DROP TABLE IF EXISTS ${test_table_for_ai_functions}""")
try_sql("""DROP RESOURCE IF EXISTS '${resourceName}'""")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
![https://example.com/a.png"}](https://example.com/a.png"}){kind=link}