mirror of https://github.com/citusdata/citus.git
141 lines
3.8 KiB
C
141 lines
3.8 KiB
C
/*-------------------------------------------------------------------------
|
|
*
|
|
* citus_global_signal.c
|
|
* Commands for Citus' overriden versions of pg_cancel_backend
|
|
* and pg_terminate_backend statements.
|
|
*
|
|
* Copyright (c) Citus Data, Inc.
|
|
*
|
|
*-------------------------------------------------------------------------
|
|
*/
|
|
|
|
#include "postgres.h"
|
|
|
|
#include "signal.h"
|
|
|
|
#include "lib/stringinfo.h"
|
|
|
|
#include "pg_version_constants.h"
|
|
|
|
#include "distributed/backend_data.h"
|
|
#include "distributed/metadata_cache.h"
|
|
#include "distributed/remote_commands.h"
|
|
#include "distributed/worker_manager.h"
|
|
|
|
static bool CitusSignalBackend(uint64 globalPID, uint64 timeout, int sig);
|
|
|
|
PG_FUNCTION_INFO_V1(citus_cancel_backend);
|
|
PG_FUNCTION_INFO_V1(citus_terminate_backend);
|
|
|
|
/*
|
|
* pg_cancel_backend overrides the Postgres' pg_cancel_backend to cancel
|
|
* a query with a global pid so a query can be cancelled from another node.
|
|
*
|
|
* To cancel a query that is on another node, a pg_cancel_backend command is sent
|
|
* to that node. This new command is sent with pid instead of global pid, so original
|
|
* pg_cancel_backend function is used.
|
|
*/
|
|
Datum
|
|
citus_cancel_backend(PG_FUNCTION_ARGS)
|
|
{
|
|
CheckCitusVersion(ERROR);
|
|
|
|
uint64 pid = PG_GETARG_INT64(0);
|
|
|
|
int sig = SIGINT;
|
|
uint64 timeout = 0;
|
|
bool success = CitusSignalBackend(pid, timeout, sig);
|
|
|
|
PG_RETURN_BOOL(success);
|
|
}
|
|
|
|
|
|
/*
|
|
* pg_terminate_backend overrides the Postgres' pg_terminate_backend to terminate
|
|
* a query with a global pid so a query can be terminated from another node.
|
|
*
|
|
* To terminate a query that is on another node, a pg_terminate_backend command is sent
|
|
* to that node. This new command is sent with pid instead of global pid, so original
|
|
* pg_terminate_backend function is used.
|
|
*/
|
|
Datum
|
|
citus_terminate_backend(PG_FUNCTION_ARGS)
|
|
{
|
|
CheckCitusVersion(ERROR);
|
|
|
|
uint64 pid = PG_GETARG_INT64(0);
|
|
uint64 timeout = PG_GETARG_INT64(1);
|
|
|
|
int sig = SIGTERM;
|
|
bool success = CitusSignalBackend(pid, timeout, sig);
|
|
|
|
PG_RETURN_BOOL(success);
|
|
}
|
|
|
|
|
|
/*
|
|
* CitusSignalBackend gets a global pid and and ends the original query with the global pid
|
|
* that might have started in another node by connecting to that node and running either
|
|
* pg_cancel_backend or pg_terminate_backend based on the withTerminate argument.
|
|
*/
|
|
static bool
|
|
CitusSignalBackend(uint64 globalPID, uint64 timeout, int sig)
|
|
{
|
|
Assert((sig == SIGINT) || (sig == SIGTERM));
|
|
|
|
bool missingOk = false;
|
|
int nodeId = ExtractNodeIdFromGlobalPID(globalPID, missingOk);
|
|
int processId = ExtractProcessIdFromGlobalPID(globalPID);
|
|
|
|
WorkerNode *workerNode = FindNodeWithNodeId(nodeId, missingOk);
|
|
|
|
StringInfo cancelQuery = makeStringInfo();
|
|
|
|
if (sig == SIGINT)
|
|
{
|
|
appendStringInfo(cancelQuery, "SELECT pg_cancel_backend(%d::integer)", processId);
|
|
}
|
|
else
|
|
{
|
|
appendStringInfo(cancelQuery,
|
|
"SELECT pg_terminate_backend(%d::integer, %lu::bigint)",
|
|
processId, timeout);
|
|
}
|
|
|
|
int connectionFlags = 0;
|
|
MultiConnection *connection = GetNodeConnection(connectionFlags,
|
|
workerNode->workerName,
|
|
workerNode->workerPort);
|
|
|
|
if (!SendRemoteCommand(connection, cancelQuery->data))
|
|
{
|
|
/* if we cannot connect, we warn and report false */
|
|
ReportConnectionError(connection, WARNING);
|
|
return false;
|
|
}
|
|
|
|
bool raiseInterrupts = true;
|
|
PGresult *queryResult = GetRemoteCommandResult(connection, raiseInterrupts);
|
|
|
|
/* if remote node throws an error, we also throw an error */
|
|
if (!IsResponseOK(queryResult))
|
|
{
|
|
ReportResultError(connection, queryResult, ERROR);
|
|
}
|
|
|
|
StringInfo queryResultString = makeStringInfo();
|
|
bool success = EvaluateSingleQueryResult(connection, queryResult, queryResultString);
|
|
if (success && strcmp(queryResultString->data, "f") == 0)
|
|
{
|
|
/* worker node returned "f" */
|
|
success = false;
|
|
}
|
|
|
|
PQclear(queryResult);
|
|
|
|
bool raiseErrors = false;
|
|
ClearResults(connection, raiseErrors);
|
|
|
|
return success;
|
|
}
|