/*
* CreateTask creates a new task in shared hash, initializes the task, and sets
* the task to assigned state. Note that this function expects the caller to
* hold an exclusive lock over the shared hash.
*/
static void
CreateTask(uint64 jobId, uint32 taskId, char *taskCallString)
{
WorkerTask *workerTask = NULL;
uint32 assignmentTime = 0;
char *databaseName = get_database_name(MyDatabaseId);
char *userName = CurrentUserName();
/* increase task priority for cleanup tasks */
assignmentTime = (uint32) time(NULL);
if (taskId == JOB_CLEANUP_TASK_ID)
{
assignmentTime = HIGH_PRIORITY_TASK_TIME;
}
/* enter the worker task into shared hash and initialize the task */
workerTask = WorkerTasksHashEnter(jobId, taskId);
workerTask->assignedAt = assignmentTime;
strlcpy(workerTask->taskCallString, taskCallString, TASK_CALL_STRING_SIZE);
workerTask->taskStatus = TASK_ASSIGNED;
workerTask->connectionId = INVALID_CONNECTION_ID;
workerTask->failureCount = 0;
strlcpy(workerTask->databaseName, databaseName, NAMEDATALEN);
strlcpy(workerTask->userName, userName, NAMEDATALEN);
}
/*
* GetOrEstablishConnection returns a PGconn which can be used to execute
* queries on a remote PostgreSQL server. If no suitable connection to the
* specified node on the specified port yet exists, the function establishes
* a new connection and adds it to the connection cache before returning it.
*
* Returned connections are guaranteed to be in the CONNECTION_OK state. If the
* requested connection cannot be established, or if it was previously created
* but is now in an unrecoverable bad state, this function returns NULL.
*
* This function throws an error if a hostname over 255 characters is provided.
*/
PGconn *
GetOrEstablishConnection(char *nodeName, int32 nodePort)
{
PGconn *connection = NULL;
NodeConnectionKey nodeConnectionKey;
NodeConnectionEntry *nodeConnectionEntry = NULL;
bool entryFound = false;
bool needNewConnection = true;
char *userName = CurrentUserName();
/* check input */
if (strnlen(nodeName, MAX_NODE_LENGTH + 1) > MAX_NODE_LENGTH)
{
ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("hostname exceeds the maximum length of %d",
MAX_NODE_LENGTH)));
}
/* if first call, initialize the connection hash */
if (NodeConnectionHash == NULL)
{
NodeConnectionHash = CreateNodeConnectionHash();
}
memset(&nodeConnectionKey, 0, sizeof(nodeConnectionKey));
strncpy(nodeConnectionKey.nodeName, nodeName, MAX_NODE_LENGTH);
nodeConnectionKey.nodePort = nodePort;
strncpy(nodeConnectionKey.nodeUser, userName, NAMEDATALEN);
nodeConnectionEntry = hash_search(NodeConnectionHash, &nodeConnectionKey,
HASH_FIND, &entryFound);
if (entryFound)
{
connection = nodeConnectionEntry->connection;
if (PQstatus(connection) == CONNECTION_OK)
{
needNewConnection = false;
}
else
{
PurgeConnection(connection);
}
}
if (needNewConnection)
{
connection = ConnectToNode(nodeName, nodePort, nodeConnectionKey.nodeUser);
if (connection != NULL)
{
nodeConnectionEntry = hash_search(NodeConnectionHash, &nodeConnectionKey,
HASH_ENTER, &entryFound);
nodeConnectionEntry->connection = connection;
}
}
return connection;
}
/*
* MultiClientConnectStart asynchronously tries to establish a connection. If it
* succeeds, it returns the connection id. Otherwise, it reports connection
* error and returns INVALID_CONNECTION_ID.
*/
int32
MultiClientConnectStart(const char *nodeName, uint32 nodePort, const char *nodeDatabase)
{
PGconn *connection = NULL;
char connInfoString[STRING_BUFFER_SIZE];
ConnStatusType connStatusType = CONNECTION_BAD;
char *userName = CurrentUserName();
int32 connectionId = AllocateConnectionId();
if (connectionId == INVALID_CONNECTION_ID)
{
ereport(WARNING, (errmsg("could not allocate connection in connection pool")));
return connectionId;
}
if (XactModificationLevel > XACT_MODIFICATION_NONE)
{
ereport(ERROR, (errcode(ERRCODE_ACTIVE_SQL_TRANSACTION),
errmsg("cannot open new connections after the first modification "
"command within a transaction")));
}
/* transcribe connection paremeters to string */
snprintf(connInfoString, STRING_BUFFER_SIZE, CONN_INFO_TEMPLATE,
nodeName, nodePort, nodeDatabase, userName, CLIENT_CONNECT_TIMEOUT);
/* prepare asynchronous request for worker node connection */
connection = PQconnectStart(connInfoString);
connStatusType = PQstatus(connection);
/*
* If prepared, we save the connection, and set its initial polling status
* to PGRES_POLLING_WRITING as specified in "Database Connection Control
* Functions" section of the PostgreSQL documentation.
*/
if (connStatusType != CONNECTION_BAD)
{
ClientConnectionArray[connectionId] = connection;
ClientPollingStatusArray[connectionId] = PGRES_POLLING_WRITING;
}
else
{
WarnRemoteError(connection, NULL);
PQfinish(connection);
connectionId = INVALID_CONNECTION_ID;
}
return connectionId;
}
/*
* 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);
}
/*
* MultiClientConnect synchronously tries to establish a connection. If it
* succeeds, it returns the connection id. Otherwise, it reports connection
* error and returns INVALID_CONNECTION_ID.
*
* nodeDatabase and userName can be NULL, in which case values from the
* current session are used.
*/
int32
MultiClientConnect(const char *nodeName, uint32 nodePort, const char *nodeDatabase,
const char *userName)
{
PGconn *connection = NULL;
char connInfoString[STRING_BUFFER_SIZE];
ConnStatusType connStatusType = CONNECTION_OK;
int32 connectionId = AllocateConnectionId();
char *effectiveDatabaseName = NULL;
char *effectiveUserName = NULL;
if (XactModificationLevel > XACT_MODIFICATION_NONE)
{
ereport(ERROR, (errcode(ERRCODE_ACTIVE_SQL_TRANSACTION),
errmsg("cannot open new connections after the first modification "
"command within a transaction")));
}
if (connectionId == INVALID_CONNECTION_ID)
{
ereport(WARNING, (errmsg("could not allocate connection in connection pool")));
return connectionId;
}
if (nodeDatabase == NULL)
{
effectiveDatabaseName = get_database_name(MyDatabaseId);
}
else
{
effectiveDatabaseName = pstrdup(nodeDatabase);
}
if (userName == NULL)
{
effectiveUserName = CurrentUserName();
}
else
{
effectiveUserName = pstrdup(userName);
}
/*
* FIXME: This code is bad on several levels. It completely forgoes any
* escaping, it misses setting a number of parameters, it works with a
* limited string size without erroring when it's too long. We shouldn't
* even build a query string this way, there's PQconnectdbParams()!
*/
/* transcribe connection paremeters to string */
snprintf(connInfoString, STRING_BUFFER_SIZE, CONN_INFO_TEMPLATE,
nodeName, nodePort,
effectiveDatabaseName, effectiveUserName,
CLIENT_CONNECT_TIMEOUT);
/* establish synchronous connection to worker node */
connection = PQconnectdb(connInfoString);
connStatusType = PQstatus(connection);
if (connStatusType == CONNECTION_OK)
{
ClientConnectionArray[connectionId] = connection;
}
else
{
WarnRemoteError(connection, NULL);
PQfinish(connection);
connectionId = INVALID_CONNECTION_ID;
}
pfree(effectiveDatabaseName);
pfree(effectiveUserName);
return connectionId;
}
/*
* StartPlacementListConnection returns a connection to a remote node suitable for
* a placement accesses (SELECT, DML, DDL) or throws an error if no suitable
* connection can be established if would cause a self-deadlock or consistency
* violation.
*/
MultiConnection *
StartPlacementListConnection(uint32 flags, List *placementAccessList,
const char *userName)
{
char *freeUserName = NULL;
ListCell *placementAccessCell = NULL;
List *placementEntryList = NIL;
ListCell *placementEntryCell = NULL;
MultiConnection *chosenConnection = NULL;
if (userName == NULL)
{
userName = freeUserName = CurrentUserName();
}
chosenConnection = FindPlacementListConnection(flags, placementAccessList, userName,
&placementEntryList);
if (chosenConnection == NULL)
{
/* use the first placement from the list to extract nodename and nodeport */
ShardPlacementAccess *placementAccess =
(ShardPlacementAccess *) linitial(placementAccessList);
ShardPlacement *placement = placementAccess->placement;
char *nodeName = placement->nodeName;
int nodePort = placement->nodePort;
/*
* No suitable connection in the placement->connection mapping, get one from
* the node->connection pool.
*/
chosenConnection = StartNodeUserDatabaseConnection(flags, nodeName, nodePort,
userName, NULL);
if (flags & CONNECTION_PER_PLACEMENT &&
ConnectionAccessedDifferentPlacement(chosenConnection, placement))
{
/*
* Cached connection accessed a non-co-located placement in the same
* table or co-location group, while the caller asked for a connection
* per placement. Open a new connection instead.
*
* We use this for situations in which we want to use a different
* connection for every placement, such as COPY. If we blindly returned
* a cached conection that already modified a different, non-co-located
* placement B in the same table or in a table with the same co-location
* ID as the current placement, then we'd no longer able to write to
* placement B later in the COPY.
*/
chosenConnection = StartNodeUserDatabaseConnection(flags |
FORCE_NEW_CONNECTION,
nodeName, nodePort,
userName, NULL);
Assert(!ConnectionAccessedDifferentPlacement(chosenConnection, placement));
}
}
/*
* Now that a connection has been chosen, initialise or update the connection
* references for all placements.
*/
forboth(placementAccessCell, placementAccessList,
placementEntryCell, placementEntryList)
{
ShardPlacementAccess *placementAccess =
(ShardPlacementAccess *) lfirst(placementAccessCell);
ShardPlacementAccessType accessType = placementAccess->accessType;
ConnectionPlacementHashEntry *placementEntry =
(ConnectionPlacementHashEntry *) lfirst(placementEntryCell);
ConnectionReference *placementConnection = placementEntry->primaryConnection;
if (placementConnection->connection == chosenConnection)
{
/* using the connection that was already assigned to the placement */
}
else if (placementConnection->connection == NULL)
{
/* placement does not have a connection assigned yet */
placementConnection->connection = chosenConnection;
placementConnection->hadDDL = false;
placementConnection->hadDML = false;
placementConnection->userName = MemoryContextStrdup(TopTransactionContext,
userName);
placementConnection->placementId = placementAccess->placement->placementId;
/* record association with connection */
dlist_push_tail(&chosenConnection->referencedPlacements,
&placementConnection->connectionNode);
}
else
{
/* using a different connection than the one assigned to the placement */
if (accessType != PLACEMENT_ACCESS_SELECT)
{
/*
* We previously read from the placement, but now we're writing to
* it (if we had written to the placement, we would have either chosen
* the same connection, or errored out). Update the connection reference
* to point to the connection used for writing. We don't need to remember
* the existing connection since we won't be able to reuse it for
* accessing the placement. However, we do register that it exists in
* hasSecondaryConnections.
*/
placementConnection->connection = chosenConnection;
placementConnection->userName = MemoryContextStrdup(TopTransactionContext,
userName);
Assert(!placementConnection->hadDDL);
Assert(!placementConnection->hadDML);
/* record association with connection */
dlist_push_tail(&chosenConnection->referencedPlacements,
&placementConnection->connectionNode);
}
/*
* There are now multiple connections that read from the placement
* and DDL commands are forbidden.
*/
placementEntry->hasSecondaryConnections = true;
if (placementEntry->colocatedEntry != NULL)
{
/* we also remember this for co-located placements */
placementEntry->colocatedEntry->hasSecondaryConnections = true;
}
}
/*
* Remember that we used the current connection for writes.
*/
if (accessType == PLACEMENT_ACCESS_DDL)
{
placementConnection->hadDDL = true;
}
if (accessType == PLACEMENT_ACCESS_DML)
{
placementConnection->hadDML = true;
}
}
