/*
* ReceiveResourceCleanup gets called if an error occurs during file receiving.
* The function closes the connection, and closes and deletes the local file.
*/
static void
ReceiveResourceCleanup(int32 connectionId, const char *filename, int32 fileDescriptor)
{
if (connectionId != INVALID_CONNECTION_ID)
{
MultiClientDisconnect(connectionId);
}
if (fileDescriptor != -1)
{
int closed = -1;
int deleted = -1;
closed = close(fileDescriptor);
if (closed < 0)
{
ereport(WARNING, (errcode_for_file_access(),
errmsg("could not close file \"%s\": %m", filename)));
}
deleted = unlink(filename);
if (deleted != 0)
{
ereport(WARNING, (errcode_for_file_access(),
errmsg("could not delete file \"%s\": %m", filename)));
}
}
}
/*
* CleanupTaskExecution iterates over all connections and file descriptors for
* the given task execution. The function first closes all open connections and
* file descriptors, and then frees memory allocated for the task execution.
*/
void
CleanupTaskExecution(TaskExecution *taskExecution)
{
uint32 nodeIndex = 0;
for (nodeIndex = 0; nodeIndex < taskExecution->nodeCount; nodeIndex++)
{
int32 connectionId = taskExecution->connectionIdArray[nodeIndex];
int32 fileDescriptor = taskExecution->fileDescriptorArray[nodeIndex];
/* close open connection */
if (connectionId != INVALID_CONNECTION_ID)
{
MultiClientDisconnect(connectionId);
taskExecution->connectionIdArray[nodeIndex] = INVALID_CONNECTION_ID;
}
/* close open file */
if (fileDescriptor >= 0)
{
int closed = close(fileDescriptor);
taskExecution->fileDescriptorArray[nodeIndex] = -1;
if (closed < 0)
{
ereport(WARNING, (errcode_for_file_access(),
errmsg("could not close copy file: %m")));
}
}
}
/* deallocate memory and reset all fields */
pfree(taskExecution->taskStatusArray);
pfree(taskExecution->connectionIdArray);
pfree(taskExecution->fileDescriptorArray);
memset(taskExecution, 0, sizeof(TaskExecution));
}
/*
* ReceiveRegularFile creates a local file at the given file path, and connects
* to remote database that has the given node name and port number. The function
* then issues the given transmit command using client-side logic (libpq), reads
* the remote file's contents, and appends these contents to the local file. On
* success, the function returns success; on failure, it cleans up all resources
* and returns false.
*/
static bool
ReceiveRegularFile(const char *nodeName, uint32 nodePort,
StringInfo transmitCommand, StringInfo filePath)
{
int32 fileDescriptor = -1;
char filename[MAXPGPATH];
int closed = -1;
const int fileFlags = (O_APPEND | O_CREAT | O_RDWR | O_TRUNC | PG_BINARY);
const int fileMode = (S_IRUSR | S_IWUSR);
QueryStatus queryStatus = CLIENT_INVALID_QUERY;
int32 connectionId = INVALID_CONNECTION_ID;
char *nodeDatabase = NULL;
bool querySent = false;
bool queryReady = false;
bool copyDone = false;
/* create local file to append remote data to */
snprintf(filename, MAXPGPATH, "%s", filePath->data);
fileDescriptor = BasicOpenFile(filename, fileFlags, fileMode);
if (fileDescriptor < 0)
{
ereport(WARNING, (errcode_for_file_access(),
errmsg("could not open file \"%s\": %m", filePath->data)));
return false;
}
/* we use the same database name on the master and worker nodes */
nodeDatabase = get_database_name(MyDatabaseId);
/* connect to remote node */
connectionId = MultiClientConnect(nodeName, nodePort, nodeDatabase, NULL);
if (connectionId == INVALID_CONNECTION_ID)
{
ReceiveResourceCleanup(connectionId, filename, fileDescriptor);
return false;
}
/* send request to remote node to start transmitting data */
querySent = MultiClientSendQuery(connectionId, transmitCommand->data);
if (!querySent)
{
ReceiveResourceCleanup(connectionId, filename, fileDescriptor);
return false;
}
/* loop until the remote node acknowledges our transmit request */
while (!queryReady)
{
ResultStatus resultStatus = MultiClientResultStatus(connectionId);
if (resultStatus == CLIENT_RESULT_READY)
{
queryReady = true;
}
else if (resultStatus == CLIENT_RESULT_BUSY)
{
/* remote node did not respond; wait for longer */
long sleepIntervalPerCycle = RemoteTaskCheckInterval * 1000L;
pg_usleep(sleepIntervalPerCycle);
}
else
{
ReceiveResourceCleanup(connectionId, filename, fileDescriptor);
return false;
}
}
/* check query response is as expected */
queryStatus = MultiClientQueryStatus(connectionId);
if (queryStatus != CLIENT_QUERY_COPY)
{
ReceiveResourceCleanup(connectionId, filename, fileDescriptor);
return false;
}
/* loop until we receive and append all the data from remote node */
while (!copyDone)
{
CopyStatus copyStatus = MultiClientCopyData(connectionId, fileDescriptor);
if (copyStatus == CLIENT_COPY_DONE)
{
copyDone = true;
}
else if (copyStatus == CLIENT_COPY_MORE)
{
/* remote node will continue to send more data */
}
else
{
ReceiveResourceCleanup(connectionId, filename, fileDescriptor);
return false;
}
}
/* we are done executing; release the connection and the file handle */
MultiClientDisconnect(connectionId);
closed = close(fileDescriptor);
if (closed < 0)
{
ereport(WARNING, (errcode_for_file_access(),
errmsg("could not close file \"%s\": %m", filename)));
/* if we failed to close file, try to delete it before erroring out */
DeleteFile(filename);
return false;
}
/* we successfully received the remote file */
ereport(DEBUG2, (errmsg("received remote file \"%s\"", filename)));
return true;
}
