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