diff --git a/src/backend/distributed/test/connection_cache.c b/src/backend/distributed/test/connection_cache.c index 3d641d837..8c46d0dd5 100644 --- a/src/backend/distributed/test/connection_cache.c +++ b/src/backend/distributed/test/connection_cache.c @@ -13,7 +13,6 @@ #include "postgres.h" #include "c.h" #include "fmgr.h" - #include "libpq-fe.h" #include @@ -22,7 +21,9 @@ #include "catalog/pg_type.h" #include "distributed/connection_cache.h" +#include "distributed/metadata_cache.h" #include "distributed/test_helper_functions.h" /* IWYU pragma: keep */ +#include "utils/elog.h" #include "utils/lsyscache.h" @@ -34,6 +35,7 @@ static Datum ExtractIntegerDatum(char *input); PG_FUNCTION_INFO_V1(initialize_remote_temp_table); PG_FUNCTION_INFO_V1(count_remote_temp_table_rows); PG_FUNCTION_INFO_V1(get_and_purge_connection); +PG_FUNCTION_INFO_V1(connect_and_purge_connection); PG_FUNCTION_INFO_V1(set_connection_status_bad); @@ -128,6 +130,30 @@ get_and_purge_connection(PG_FUNCTION_ARGS) } +/* + * get_and_purge_connection first gets a connection using the provided hostname + * and port before immediately passing that connection to PurgeConnection. This + * is to test PurgeConnection behvaior when circumventing the cache. + */ +Datum +connect_and_purge_connection(PG_FUNCTION_ARGS) +{ + char *nodeName = PG_GETARG_CSTRING(0); + int32 nodePort = PG_GETARG_INT32(1); + char *nodeUser = CurrentUserName(); + + PGconn *connection = ConnectToNode(nodeName, nodePort, nodeUser); + if (connection == NULL) + { + PG_RETURN_BOOL(false); + } + + PurgeConnection(connection); + + PG_RETURN_BOOL(true); +} + + /* * set_connection_status_bad does not remove the given connection from the connection hash. * It simply shuts down the underlying socket. On success, it returns true. diff --git a/src/backend/distributed/utils/connection_cache.c b/src/backend/distributed/utils/connection_cache.c index c6f062028..766a5ade2 100644 --- a/src/backend/distributed/utils/connection_cache.c +++ b/src/backend/distributed/utils/connection_cache.c @@ -14,16 +14,14 @@ #include "libpq-fe.h" #include "miscadmin.h" - #include +#include #include #include "commands/dbcommands.h" #include "distributed/connection_cache.h" #include "distributed/metadata_cache.h" -#include "lib/stringinfo.h" #include "mb/pg_wchar.h" -#include "nodes/pg_list.h" #include "utils/builtins.h" #include "utils/elog.h" #include "utils/errcodes.h" @@ -172,10 +170,6 @@ PurgeConnection(PGconn *connection) */ if (purgedConnection != connection) { - ereport(WARNING, (errmsg("hash entry for \"%s:%d\" contained different " - "connection than that provided by caller", - nodeConnectionKey.nodeName, - nodeConnectionKey.nodePort))); PQfinish(connection); } } @@ -186,22 +180,21 @@ PurgeConnectionByKey(NodeConnectionKey *nodeConnectionKey) { bool entryFound = false; NodeConnectionEntry *nodeConnectionEntry = NULL; + PGconn *connection = NULL; + + if (NodeConnectionHash != NULL) + { + nodeConnectionEntry = hash_search(NodeConnectionHash, nodeConnectionKey, + HASH_REMOVE, &entryFound); + } - nodeConnectionEntry = hash_search(NodeConnectionHash, nodeConnectionKey, HASH_REMOVE, - &entryFound); if (entryFound) { + connection = nodeConnectionEntry->connection; PQfinish(nodeConnectionEntry->connection); } - else - { - ereport(WARNING, (errcode(ERRCODE_NO_DATA), - errmsg("could not find hash entry for connection to \"%s:%d\"", - nodeConnectionKey->nodeName, - nodeConnectionKey->nodePort))); - } - return nodeConnectionEntry->connection; + return connection; } diff --git a/src/include/distributed/test_helper_functions.h b/src/include/distributed/test_helper_functions.h index 6ba52a1b0..22bcad541 100644 --- a/src/include/distributed/test_helper_functions.h +++ b/src/include/distributed/test_helper_functions.h @@ -39,6 +39,7 @@ extern Datum fake_fdw_handler(PG_FUNCTION_ARGS); extern Datum initialize_remote_temp_table(PG_FUNCTION_ARGS); extern Datum count_remote_temp_table_rows(PG_FUNCTION_ARGS); extern Datum get_and_purge_connection(PG_FUNCTION_ARGS); +extern Datum connect_and_purge_connection(PG_FUNCTION_ARGS); extern Datum set_connection_status_bad(PG_FUNCTION_ARGS); /* function declarations for exercising metadata functions */ diff --git a/src/test/regress/expected/multi_connection_cache.out b/src/test/regress/expected/multi_connection_cache.out index fd17f3845..79cf59dd2 100644 --- a/src/test/regress/expected/multi_connection_cache.out +++ b/src/test/regress/expected/multi_connection_cache.out @@ -15,6 +15,10 @@ CREATE FUNCTION get_and_purge_connection(cstring, integer) RETURNS bool AS 'citus' LANGUAGE C STRICT; +CREATE FUNCTION connect_and_purge_connection(cstring, integer) + RETURNS bool + AS 'citus' + LANGUAGE C STRICT; CREATE FUNCTION set_connection_status_bad(cstring, integer) RETURNS bool AS 'citus' @@ -104,3 +108,11 @@ SELECT get_and_purge_connection('localhost', :worker_port); (1 row) SET client_min_messages TO DEFAULT; +\c +-- purge existing connection to localhost +SELECT connect_and_purge_connection('localhost', :worker_port); + connect_and_purge_connection +------------------------------ + t +(1 row) + diff --git a/src/test/regress/sql/multi_connection_cache.sql b/src/test/regress/sql/multi_connection_cache.sql index 442545c73..996eb8975 100644 --- a/src/test/regress/sql/multi_connection_cache.sql +++ b/src/test/regress/sql/multi_connection_cache.sql @@ -22,6 +22,11 @@ CREATE FUNCTION get_and_purge_connection(cstring, integer) AS 'citus' LANGUAGE C STRICT; +CREATE FUNCTION connect_and_purge_connection(cstring, integer) + RETURNS bool + AS 'citus' + LANGUAGE C STRICT; + CREATE FUNCTION set_connection_status_bad(cstring, integer) RETURNS bool AS 'citus' @@ -76,3 +81,8 @@ SELECT count_remote_temp_table_rows('localhost', :worker_port); SELECT get_and_purge_connection('localhost', :worker_port); SET client_min_messages TO DEFAULT; + +\c + +-- purge existing connection to localhost +SELECT connect_and_purge_connection('localhost', :worker_port);