/* * 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; } }