This is an automated email from the ASF dual-hosted git repository. djwang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/cloudberry.git
commit 7c0bb732728f2801472de47bd88c0514029c9366 Author: NJrslv <[email protected]> AuthorDate: Tue Jan 20 17:03:53 2026 +0300 [yagp_hooks_collector] Add UDS round-trip test and fix send() accounting Add regression test for UDS transport. Fix send() return value: do not add -1 to total_bytes_sent on error. General refactoring. --- expected/yagp_uds.out | 42 ++++++++++++ gpcontrib/yagp_hooks_collector/Makefile | 2 +- sql/yagp_uds.sql | 31 +++++++++ src/Config.cpp | 10 +-- src/Config.h | 8 +-- src/UDSConnector.cpp | 117 +++++++++++++++++--------------- src/hook_wrappers.cpp | 96 ++++++++++++++++++++++++-- src/hook_wrappers.h | 4 ++ src/yagp_hooks_collector.c | 64 +++++++++++++++-- yagp_hooks_collector--1.1.sql | 15 ++++ 10 files changed, 318 insertions(+), 71 deletions(-) diff --git a/expected/yagp_uds.out b/expected/yagp_uds.out new file mode 100644 index 00000000000..d04929ffb4a --- /dev/null +++ b/expected/yagp_uds.out @@ -0,0 +1,42 @@ +-- Test UDS socket +-- start_ignore +CREATE EXTENSION IF NOT EXISTS yagp_hooks_collector; +-- end_ignore +\set UDS_PATH '/tmp/yagpcc_test.sock' +-- Configure extension to send via UDS +SET yagpcc.uds_path TO :'UDS_PATH'; +SET yagpcc.ignored_users_list TO ''; +SET yagpcc.enable TO TRUE; +SET yagpcc.logging_mode TO 'UDS'; +-- Start receiver +SELECT yagpcc.__test_uds_start_server(:'UDS_PATH'); + __test_uds_start_server +------------------------- +(0 rows) + +-- Send +SELECT 1; + ?column? +---------- + 1 +(1 row) + +-- Receive +SELECT yagpcc.__test_uds_receive() > 0 as received; + received +---------- + t +(1 row) + +-- Stop receiver +SELECT yagpcc.__test_uds_stop_server(); + __test_uds_stop_server +------------------------ +(0 rows) + +-- Cleanup +DROP EXTENSION yagp_hooks_collector; +RESET yagpcc.uds_path; +RESET yagpcc.ignored_users_list; +RESET yagpcc.enable; +RESET yagpcc.logging_mode; diff --git a/gpcontrib/yagp_hooks_collector/Makefile b/gpcontrib/yagp_hooks_collector/Makefile index 79f5401c8d1..eb6541b7687 100644 --- a/gpcontrib/yagp_hooks_collector/Makefile +++ b/gpcontrib/yagp_hooks_collector/Makefile @@ -1,7 +1,7 @@ MODULE_big = yagp_hooks_collector EXTENSION = yagp_hooks_collector DATA = $(wildcard *--*.sql) -REGRESS = yagp_cursors yagp_dist yagp_select yagp_utf8_trim yagp_utility yagp_guc_cache +REGRESS = yagp_cursors yagp_dist yagp_select yagp_utf8_trim yagp_utility yagp_guc_cache yagp_uds PROTO_BASES = yagpcc_plan yagpcc_metrics yagpcc_set_service PROTO_OBJS = $(patsubst %,src/protos/%.pb.o,$(PROTO_BASES)) diff --git a/sql/yagp_uds.sql b/sql/yagp_uds.sql new file mode 100644 index 00000000000..3eef697a4e7 --- /dev/null +++ b/sql/yagp_uds.sql @@ -0,0 +1,31 @@ +-- Test UDS socket +-- start_ignore +CREATE EXTENSION IF NOT EXISTS yagp_hooks_collector; +-- end_ignore + +\set UDS_PATH '/tmp/yagpcc_test.sock' + +-- Configure extension to send via UDS +SET yagpcc.uds_path TO :'UDS_PATH'; +SET yagpcc.ignored_users_list TO ''; +SET yagpcc.enable TO TRUE; +SET yagpcc.logging_mode TO 'UDS'; + +-- Start receiver +SELECT yagpcc.__test_uds_start_server(:'UDS_PATH'); + +-- Send +SELECT 1; + +-- Receive +SELECT yagpcc.__test_uds_receive() > 0 as received; + +-- Stop receiver +SELECT yagpcc.__test_uds_stop_server(); + +-- Cleanup +DROP EXTENSION yagp_hooks_collector; +RESET yagpcc.uds_path; +RESET yagpcc.ignored_users_list; +RESET yagpcc.enable; +RESET yagpcc.logging_mode; diff --git a/src/Config.cpp b/src/Config.cpp index 4fb58677018..2c2032ebb03 100644 --- a/src/Config.cpp +++ b/src/Config.cpp @@ -16,9 +16,9 @@ static bool guc_enable_cdbstats = true; static bool guc_enable_collector = true; static bool guc_report_nested_queries = true; static char *guc_ignored_users = nullptr; -static int guc_max_text_size = 1 << 20; // in bytes (1MB) -static int guc_max_plan_size = 1024; // in KB -static int guc_min_analyze_time = 10000; // in ms +static int guc_max_text_size = 1 << 20; // in bytes (1MB) +static int guc_max_plan_size = 1024; // in KB +static int guc_min_analyze_time = 10000; // in ms static int guc_logging_mode = LOG_MODE_UDS; static bool guc_enable_utility = false; @@ -143,8 +143,8 @@ void Config::sync() { enable_collector_ = guc_enable_collector; enable_utility_ = guc_enable_utility; report_nested_queries_ = guc_report_nested_queries; - max_text_size_ = static_cast<size_t>(guc_max_text_size); - max_plan_size_ = static_cast<size_t>(guc_max_plan_size); + max_text_size_ = guc_max_text_size; + max_plan_size_ = guc_max_plan_size; min_analyze_time_ = guc_min_analyze_time; logging_mode_ = guc_logging_mode; } diff --git a/src/Config.h b/src/Config.h index b4a393b0383..aa6b5bdc0ba 100644 --- a/src/Config.h +++ b/src/Config.h @@ -21,8 +21,8 @@ public: bool enable_collector() const { return enable_collector_; } bool enable_utility() const { return enable_utility_; } bool report_nested_queries() const { return report_nested_queries_; } - size_t max_text_size() const { return max_text_size_; } - size_t max_plan_size() const { return max_plan_size_ * 1024; } + int max_text_size() const { return max_text_size_; } + int max_plan_size() const { return max_plan_size_ * 1024; } int min_analyze_time() const { return min_analyze_time_; } int logging_mode() const { return logging_mode_; } bool filter_user(const std::string &username) const; @@ -37,8 +37,8 @@ private: bool enable_collector_; bool enable_utility_; bool report_nested_queries_; - size_t max_text_size_; - size_t max_plan_size_; + int max_text_size_; + int max_plan_size_; int min_analyze_time_; int logging_mode_; }; diff --git a/src/UDSConnector.cpp b/src/UDSConnector.cpp index 74fd57a3ac0..ea118fca783 100644 --- a/src/UDSConnector.cpp +++ b/src/UDSConnector.cpp @@ -27,66 +27,77 @@ static void inline log_tracing_failure(const yagpcc::SetQueryReq &req, bool UDSConnector::report_query(const yagpcc::SetQueryReq &req, const std::string &event, const Config &config) { - sockaddr_un address; + sockaddr_un address{}; address.sun_family = AF_UNIX; - const std::string &uds_path = config.uds_path(); + const auto &uds_path = config.uds_path(); + if (uds_path.size() >= sizeof(address.sun_path)) { ereport(WARNING, (errmsg("UDS path is too long for socket buffer"))); YagpStat::report_error(); return false; } strcpy(address.sun_path, uds_path.c_str()); - bool success = true; - auto sockfd = socket(AF_UNIX, SOCK_STREAM, 0); - if (sockfd != -1) { - if (fcntl(sockfd, F_SETFL, O_NONBLOCK) != -1) { - if (connect(sockfd, (sockaddr *)&address, sizeof(address)) != -1) { - auto data_size = req.ByteSize(); - auto total_size = data_size + sizeof(uint32_t); - uint8_t *buf = (uint8_t *)ya_gpdb::palloc(total_size); - uint32_t *size_payload = (uint32_t *)buf; - *size_payload = data_size; - req.SerializeWithCachedSizesToArray(buf + sizeof(uint32_t)); - int64_t sent = 0, sent_total = 0; - do { - sent = send(sockfd, buf + sent_total, total_size - sent_total, - MSG_DONTWAIT); - sent_total += sent; - } while ( - sent > 0 && size_t(sent_total) != total_size && - // the line below is a small throttling hack: - // if a message does not fit a single packet, we take a nap - // before sending the next one. - // Otherwise, MSG_DONTWAIT send might overflow the UDS - (std::this_thread::sleep_for(std::chrono::milliseconds(1)), true)); - if (sent < 0) { - log_tracing_failure(req, event); - success = false; - YagpStat::report_bad_send(total_size); - } else { - YagpStat::report_send(total_size); - } - ya_gpdb::pfree(buf); - } else { - // log the error and go on - log_tracing_failure(req, event); - success = false; - YagpStat::report_bad_connection(); - } - } else { - // That's a very important error that should never happen, so make it - // visible to an end-user and admins. - ereport(WARNING, - (errmsg("Unable to create non-blocking socket connection %m"))); - success = false; - YagpStat::report_error(); - } - close(sockfd); - } else { - // log the error and go on + + const auto sockfd = socket(AF_UNIX, SOCK_STREAM, 0); + if (sockfd == -1) { log_tracing_failure(req, event); - success = false; YagpStat::report_error(); + return false; } - return success; -} \ No newline at end of file + + // Close socket automatically on error path. + struct SockGuard { + int fd; + ~SockGuard() { close(fd); } + } sock_guard{sockfd}; + + if (fcntl(sockfd, F_SETFL, O_NONBLOCK) == -1) { + // That's a very important error that should never happen, so make it + // visible to an end-user and admins. + ereport(WARNING, + (errmsg("Unable to create non-blocking socket connection %m"))); + YagpStat::report_error(); + return false; + } + + if (connect(sockfd, reinterpret_cast<sockaddr *>(&address), + sizeof(address)) == -1) { + log_tracing_failure(req, event); + YagpStat::report_bad_connection(); + return false; + } + + const auto data_size = req.ByteSize(); + const auto total_size = data_size + sizeof(uint32_t); + auto *buf = static_cast<uint8_t *>(ya_gpdb::palloc(total_size)); + // Free buf automatically on error path. + struct BufGuard { + void *p; + ~BufGuard() { ya_gpdb::pfree(p); } + } buf_guard{buf}; + + *reinterpret_cast<uint32_t *>(buf) = data_size; + req.SerializeWithCachedSizesToArray(buf + sizeof(uint32_t)); + + int64_t sent = 0, sent_total = 0; + do { + sent = + send(sockfd, buf + sent_total, total_size - sent_total, MSG_DONTWAIT); + if (sent > 0) + sent_total += sent; + } while (sent > 0 && size_t(sent_total) != total_size && + // the line below is a small throttling hack: + // if a message does not fit a single packet, we take a nap + // before sending the next one. + // Otherwise, MSG_DONTWAIT send might overflow the UDS + (std::this_thread::sleep_for(std::chrono::milliseconds(1)), true)); + + if (sent < 0) { + log_tracing_failure(req, event); + YagpStat::report_bad_send(total_size); + return false; + } + + YagpStat::report_send(total_size); + return true; +} diff --git a/src/hook_wrappers.cpp b/src/hook_wrappers.cpp index 8cf74641c29..602a2470805 100644 --- a/src/hook_wrappers.cpp +++ b/src/hook_wrappers.cpp @@ -11,6 +11,12 @@ extern "C" { #include "cdb/ml_ipc.h" #include "tcop/utility.h" #include "stat_statements_parser/pg_stat_statements_ya_parser.h" + +#include <sys/socket.h> +#include <sys/un.h> +#include <unistd.h> +#include <errno.h> +#include <poll.h> } #undef typeid @@ -52,6 +58,13 @@ static void ya_process_utility_hook(PlannedStmt *pstmt, const char *queryString, QueryEnvironment *queryEnv, DestReceiver *dest, QueryCompletion *qc); +#define TEST_MAX_CONNECTIONS 4 +#define TEST_RCV_BUF_SIZE 8192 +#define TEST_POLL_TIMEOUT_MS 200 + +static int test_server_fd = -1; +static char *test_sock_path = NULL; + static EventSender *sender = nullptr; static inline EventSender *get_sender() { @@ -226,8 +239,9 @@ static void ya_process_utility_hook(PlannedStmt *pstmt, const char *queryString, } get_sender()->decr_depth(); - cpp_call(get_sender(), &EventSender::query_metrics_collect, METRICS_QUERY_DONE, - (void *)query_desc, true /* utility */, (ErrorData *)NULL); + cpp_call(get_sender(), &EventSender::query_metrics_collect, + METRICS_QUERY_DONE, (void *)query_desc, true /* utility */, + (ErrorData *)NULL); pfree(query_desc); } @@ -242,8 +256,9 @@ static void ya_process_utility_hook(PlannedStmt *pstmt, const char *queryString, MemoryContextSwitchTo(oldctx); get_sender()->decr_depth(); - cpp_call(get_sender(), &EventSender::query_metrics_collect, METRICS_QUERY_ERROR, - (void *)query_desc, true /* utility */, edata); + cpp_call(get_sender(), &EventSender::query_metrics_collect, + METRICS_QUERY_ERROR, (void *)query_desc, true /* utility */, + edata); pfree(query_desc); ReThrowError(edata); @@ -294,4 +309,77 @@ Datum yagp_functions_get(FunctionCallInfo fcinfo) { HeapTuple tuple = ya_gpdb::heap_form_tuple(tupdesc, values, nulls); Datum result = HeapTupleGetDatum(tuple); PG_RETURN_DATUM(result); +} + +void test_uds_stop_server() { + if (test_server_fd >= 0) { + close(test_server_fd); + test_server_fd = -1; + } + if (test_sock_path) { + unlink(test_sock_path); + pfree(test_sock_path); + test_sock_path = NULL; + } +} + +void test_uds_start_server(const char *path) { + struct sockaddr_un addr = {.sun_family = AF_UNIX}; + + if (strlen(path) >= sizeof(addr.sun_path)) + ereport(ERROR, (errmsg("path too long"))); + + test_uds_stop_server(); + + strlcpy(addr.sun_path, path, sizeof(addr.sun_path)); + test_sock_path = MemoryContextStrdup(TopMemoryContext, path); + unlink(path); + + if ((test_server_fd = socket(AF_UNIX, SOCK_STREAM, 0)) < 0 || + bind(test_server_fd, (struct sockaddr *)&addr, sizeof(addr)) < 0 || + listen(test_server_fd, TEST_MAX_CONNECTIONS) < 0) { + test_uds_stop_server(); + ereport(ERROR, (errmsg("socket setup failed: %m"))); + } +} + +int64 test_uds_receive(int timeout_ms) { + char buf[TEST_RCV_BUF_SIZE]; + int rc; + struct pollfd pfd = {.fd = test_server_fd, .events = POLLIN}; + int64 total = 0; + + if (test_server_fd < 0) + ereport(ERROR, (errmsg("server not started"))); + + for (;;) { + CHECK_FOR_INTERRUPTS(); + rc = poll(&pfd, 1, Min(timeout_ms, TEST_POLL_TIMEOUT_MS)); + if (rc > 0) + break; + if (rc < 0 && errno != EINTR) + ereport(ERROR, (errmsg("poll: %m"))); + timeout_ms -= TEST_POLL_TIMEOUT_MS; + if (timeout_ms <= 0) + return total; + } + + if (pfd.revents & POLLIN) { + int client = accept(test_server_fd, NULL, NULL); + ssize_t n; + + if (client < 0) + ereport(ERROR, (errmsg("accept: %m"))); + + while ((n = recv(client, buf, sizeof(buf), 0)) != 0) { + if (n > 0) + total += n; + else if (errno != EINTR) + break; + } + + close(client); + } + + return total; } \ No newline at end of file diff --git a/src/hook_wrappers.h b/src/hook_wrappers.h index cfabf39485e..236c6eb9d79 100644 --- a/src/hook_wrappers.h +++ b/src/hook_wrappers.h @@ -12,6 +12,10 @@ extern Datum yagp_functions_get(FunctionCallInfo fcinfo); extern void init_log(); extern void truncate_log(); +extern void test_uds_start_server(const char *path); +extern int64_t test_uds_receive(int timeout_ms); +extern void test_uds_stop_server(); + #ifdef __cplusplus } #endif \ No newline at end of file diff --git a/src/yagp_hooks_collector.c b/src/yagp_hooks_collector.c index 27fd0e04b26..f7863a38921 100644 --- a/src/yagp_hooks_collector.c +++ b/src/yagp_hooks_collector.c @@ -14,16 +14,18 @@ PG_FUNCTION_INFO_V1(yagp_stat_messages); PG_FUNCTION_INFO_V1(yagp_init_log); PG_FUNCTION_INFO_V1(yagp_truncate_log); +PG_FUNCTION_INFO_V1(yagp_test_uds_start_server); +PG_FUNCTION_INFO_V1(yagp_test_uds_receive); +PG_FUNCTION_INFO_V1(yagp_test_uds_stop_server); + void _PG_init(void) { - if (Gp_role == GP_ROLE_DISPATCH || Gp_role == GP_ROLE_EXECUTE) { + if (Gp_role == GP_ROLE_DISPATCH || Gp_role == GP_ROLE_EXECUTE) hooks_init(); - } } void _PG_fini(void) { - if (Gp_role == GP_ROLE_DISPATCH || Gp_role == GP_ROLE_EXECUTE) { + if (Gp_role == GP_ROLE_DISPATCH || Gp_role == GP_ROLE_EXECUTE) hooks_deinit(); - } } Datum yagp_stat_messages_reset(PG_FUNCTION_ARGS) { @@ -65,3 +67,57 @@ Datum yagp_truncate_log(PG_FUNCTION_ARGS) { funcctx = SRF_PERCALL_SETUP(); SRF_RETURN_DONE(funcctx); } + +Datum yagp_test_uds_start_server(PG_FUNCTION_ARGS) { + FuncCallContext *funcctx; + + if (SRF_IS_FIRSTCALL()) { + funcctx = SRF_FIRSTCALL_INIT(); + char *path = text_to_cstring(PG_GETARG_TEXT_PP(0)); + test_uds_start_server(path); + pfree(path); + } + + funcctx = SRF_PERCALL_SETUP(); + SRF_RETURN_DONE(funcctx); +} + +Datum yagp_test_uds_receive(PG_FUNCTION_ARGS) { + FuncCallContext *funcctx; + int64 *result; + + if (SRF_IS_FIRSTCALL()) { + MemoryContext oldcontext; + + funcctx = SRF_FIRSTCALL_INIT(); + oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + result = (int64 *)palloc(sizeof(int64)); + funcctx->user_fctx = result; + funcctx->max_calls = 1; + MemoryContextSwitchTo(oldcontext); + + int timeout_ms = PG_GETARG_INT32(0); + *result = test_uds_receive(timeout_ms); + } + + funcctx = SRF_PERCALL_SETUP(); + + if (funcctx->call_cntr < funcctx->max_calls) { + result = (int64 *)funcctx->user_fctx; + SRF_RETURN_NEXT(funcctx, Int64GetDatum(*result)); + } + + SRF_RETURN_DONE(funcctx); +} + +Datum yagp_test_uds_stop_server(PG_FUNCTION_ARGS) { + FuncCallContext *funcctx; + + if (SRF_IS_FIRSTCALL()) { + funcctx = SRF_FIRSTCALL_INIT(); + test_uds_stop_server(); + } + + funcctx = SRF_PERCALL_SETUP(); + SRF_RETURN_DONE(funcctx); +} diff --git a/yagp_hooks_collector--1.1.sql b/yagp_hooks_collector--1.1.sql index e0e94b51493..83bfb553638 100644 --- a/yagp_hooks_collector--1.1.sql +++ b/yagp_hooks_collector--1.1.sql @@ -93,3 +93,18 @@ BEGIN PERFORM yagpcc.__truncate_log_on_segments(); END; $$ LANGUAGE plpgsql VOLATILE; + +CREATE FUNCTION yagpcc.__test_uds_start_server(path text) +RETURNS SETOF void +AS 'MODULE_PATHNAME', 'yagp_test_uds_start_server' +LANGUAGE C STRICT EXECUTE ON MASTER; + +CREATE FUNCTION yagpcc.__test_uds_receive(timeout_ms int DEFAULT 2000) +RETURNS SETOF bigint +AS 'MODULE_PATHNAME', 'yagp_test_uds_receive' +LANGUAGE C STRICT EXECUTE ON MASTER; + +CREATE FUNCTION yagpcc.__test_uds_stop_server() +RETURNS SETOF void +AS 'MODULE_PATHNAME', 'yagp_test_uds_stop_server' +LANGUAGE C EXECUTE ON MASTER; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
