mirror of
https://github.com/outbackdingo/patroni.git
synced 2026-08-25 14:53:37 +00:00
420 lines
19 KiB
Python
420 lines
19 KiB
Python
import logging
|
|
import re
|
|
import time
|
|
|
|
from threading import Condition, Event, Thread
|
|
from urllib.parse import urlparse
|
|
from typing import Any, Dict, List, Optional, Union, Tuple, TYPE_CHECKING
|
|
|
|
from .connection import Connection
|
|
from ..dcs import CITUS_COORDINATOR_GROUP_ID, Cluster
|
|
from ..psycopg import connect, quote_ident
|
|
|
|
if TYPE_CHECKING: # pragma: no cover
|
|
from psycopg import Cursor
|
|
from psycopg2 import cursor
|
|
from . import Postgresql
|
|
|
|
CITUS_SLOT_NAME_RE = re.compile(r'^citus_shard_(move|split)_slot(_[1-9][0-9]*){2,3}$')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PgDistNode(object):
|
|
"""Represents a single row in the `pg_dist_node` table"""
|
|
|
|
def __init__(self, group: int, host: str, port: int, event: str, nodeid: Optional[int] = None,
|
|
timeout: Optional[float] = None, cooldown: Optional[float] = None) -> None:
|
|
self.group = group
|
|
# A weird way of pausing client connections by adding the `-demoted` suffix to the hostname
|
|
self.host = host + ('-demoted' if event == 'before_demote' else '')
|
|
self.port = port
|
|
# Event that is trying to change or changed the given row.
|
|
# Possible values: before_demote, before_promote, after_promote.
|
|
self.event = event
|
|
self.nodeid = nodeid
|
|
|
|
# If transaction was started, we need to COMMIT/ROLLBACK before the deadline
|
|
self.timeout = timeout
|
|
self.cooldown = cooldown or 10000 # 10s by default
|
|
self.deadline: float = 0
|
|
|
|
# All changes in the pg_dist_node are serialized on the Patroni
|
|
# side by performing them from a thread. The thread, that is
|
|
# requested a change, sometimes needs to wait for a result.
|
|
# For example, we want to pause client connections before demoting
|
|
# the worker, and once it is done notify the calling thread.
|
|
self._event = Event()
|
|
|
|
def wait(self) -> None:
|
|
self._event.wait()
|
|
|
|
def wakeup(self) -> None:
|
|
self._event.set()
|
|
|
|
def __eq__(self, other: Any) -> bool:
|
|
return isinstance(other, PgDistNode) and self.event == other.event\
|
|
and self.host == other.host and self.port == other.port
|
|
|
|
def __ne__(self, other: Any) -> bool:
|
|
return not self == other
|
|
|
|
def __str__(self) -> str:
|
|
return ('PgDistNode(nodeid={0},group={1},host={2},port={3},event={4})'
|
|
.format(self.nodeid, self.group, self.host, self.port, self.event))
|
|
|
|
def __repr__(self) -> str:
|
|
return str(self)
|
|
|
|
|
|
class CitusHandler(Thread):
|
|
|
|
def __init__(self, postgresql: 'Postgresql', config: Optional[Dict[str, Union[str, int]]]) -> None:
|
|
super(CitusHandler, self).__init__()
|
|
self.daemon = True
|
|
self._postgresql = postgresql
|
|
self._config = config
|
|
self._connection = Connection()
|
|
self._pg_dist_node: Dict[int, PgDistNode] = {} # Cache of pg_dist_node: {groupid: PgDistNode()}
|
|
self._tasks: List[PgDistNode] = [] # Requests to change pg_dist_node, every task is a `PgDistNode`
|
|
self._in_flight: Optional[PgDistNode] = None # Reference to the `PgDistNode` being changed in a transaction
|
|
self._schedule_load_pg_dist_node = True # Flag that "pg_dist_node" should be queried from the database
|
|
self._condition = Condition() # protects _pg_dist_node, _tasks, _in_flight, and _schedule_load_pg_dist_node
|
|
self.schedule_cache_rebuild()
|
|
|
|
def is_enabled(self) -> bool:
|
|
return isinstance(self._config, dict)
|
|
|
|
def group(self) -> Optional[int]:
|
|
return int(self._config['group']) if isinstance(self._config, dict) else None
|
|
|
|
def is_coordinator(self) -> bool:
|
|
return self.is_enabled() and self.group() == CITUS_COORDINATOR_GROUP_ID
|
|
|
|
def is_worker(self) -> bool:
|
|
return self.is_enabled() and not self.is_coordinator()
|
|
|
|
def set_conn_kwargs(self, kwargs: Dict[str, Any]) -> None:
|
|
if isinstance(self._config, dict): # self.is_enabled():
|
|
kwargs.update({'dbname': self._config['database'],
|
|
'options': '-c statement_timeout=0 -c idle_in_transaction_session_timeout=0'})
|
|
self._connection.set_conn_kwargs(kwargs)
|
|
|
|
def schedule_cache_rebuild(self) -> None:
|
|
with self._condition:
|
|
self._schedule_load_pg_dist_node = True
|
|
|
|
def on_demote(self) -> None:
|
|
with self._condition:
|
|
self._pg_dist_node.clear()
|
|
self._tasks[:] = []
|
|
self._in_flight = None
|
|
|
|
def query(self, sql: str, *params: Any) -> Union['Cursor[Any]', 'cursor']:
|
|
try:
|
|
logger.debug('query(%s, %s)', sql, params)
|
|
cursor = self._connection.cursor()
|
|
cursor.execute(sql.encode('utf-8'), params or None)
|
|
return cursor
|
|
except Exception as e:
|
|
logger.error('Exception when executing query "%s", (%s): %r', sql, params, e)
|
|
self._connection.close()
|
|
with self._condition:
|
|
self._in_flight = None
|
|
self.schedule_cache_rebuild()
|
|
raise e
|
|
|
|
def load_pg_dist_node(self) -> bool:
|
|
"""Read from the `pg_dist_node` table and put it into the local cache"""
|
|
|
|
with self._condition:
|
|
if not self._schedule_load_pg_dist_node:
|
|
return True
|
|
self._schedule_load_pg_dist_node = False
|
|
|
|
try:
|
|
cursor = self.query("SELECT nodeid, groupid, nodename, nodeport, noderole"
|
|
" FROM pg_catalog.pg_dist_node WHERE noderole = 'primary'")
|
|
except Exception:
|
|
return False
|
|
|
|
with self._condition:
|
|
self._pg_dist_node = {r[1]: PgDistNode(r[1], r[2], r[3], 'after_promote', r[0]) for r in cursor}
|
|
return True
|
|
|
|
def sync_pg_dist_node(self, cluster: Cluster) -> None:
|
|
"""Maintain the `pg_dist_node` from the coordinator leader every heartbeat loop.
|
|
|
|
We can't always rely on REST API calls from worker nodes in order
|
|
to maintain `pg_dist_node`, therefore at least once per heartbeat
|
|
loop we make sure that workes registered in `self._pg_dist_node`
|
|
cache are matching the cluster view from DCS by creating tasks
|
|
the same way as it is done from the REST API."""
|
|
|
|
if not self.is_coordinator():
|
|
return
|
|
|
|
with self._condition:
|
|
if not self.is_alive():
|
|
self.start()
|
|
|
|
self.add_task('after_promote', CITUS_COORDINATOR_GROUP_ID, self._postgresql.connection_string)
|
|
|
|
for group, worker in cluster.workers.items():
|
|
leader = worker.leader
|
|
if leader and leader.conn_url\
|
|
and leader.data.get('role') in ('master', 'primary') and leader.data.get('state') == 'running':
|
|
self.add_task('after_promote', group, leader.conn_url)
|
|
|
|
def find_task_by_group(self, group: int) -> Optional[int]:
|
|
for i, task in enumerate(self._tasks):
|
|
if task.group == group:
|
|
return i
|
|
|
|
def pick_task(self) -> Tuple[Optional[int], Optional[PgDistNode]]:
|
|
"""Returns the tuple(i, task), where `i` - is the task index in the self._tasks list
|
|
|
|
Tasks are picked by following priorities:
|
|
1. If there is already a transaction in progress, pick a task
|
|
that that will change already affected worker primary.
|
|
2. If the coordinator address should be changed - pick a task
|
|
with group=0 (coordinators are always in group 0).
|
|
3. Pick a task that is the oldest (first from the self._tasks)"""
|
|
|
|
with self._condition:
|
|
if self._in_flight:
|
|
i = self.find_task_by_group(self._in_flight.group)
|
|
else:
|
|
while True:
|
|
i = self.find_task_by_group(CITUS_COORDINATOR_GROUP_ID) # set_coordinator
|
|
if i is None and self._tasks:
|
|
i = 0
|
|
if i is None:
|
|
break
|
|
task = self._tasks[i]
|
|
if task == self._pg_dist_node.get(task.group):
|
|
self._tasks.pop(i) # nothing to do because cached version of pg_dist_node already matches
|
|
else:
|
|
break
|
|
task = self._tasks[i] if i is not None else None
|
|
|
|
# When tasks are added it could happen that self._pg_dist_node
|
|
# wasn't ready (self._schedule_load_pg_dist_node is False)
|
|
# and hence the nodeid wasn't filled.
|
|
if task and task.group in self._pg_dist_node:
|
|
task.nodeid = self._pg_dist_node[task.group].nodeid
|
|
return i, task
|
|
|
|
def update_node(self, task: PgDistNode) -> None:
|
|
if task.nodeid is not None:
|
|
self.query('SELECT pg_catalog.citus_update_node(%s, %s, %s, true, %s)',
|
|
task.nodeid, task.host, task.port, task.cooldown)
|
|
elif task.event != 'before_demote':
|
|
row = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, 'primary', 'default')",
|
|
task.host, task.port, task.group).fetchone()
|
|
if row is not None:
|
|
task.nodeid = row[0]
|
|
|
|
def process_task(self, task: PgDistNode) -> bool:
|
|
"""Updates a single row in `pg_dist_node` table, optionally in a transaction.
|
|
|
|
The transaction is started if we do a demote of the worker node or before promoting the other worker if
|
|
there is no transaction in progress. And, the transaction is committed when the switchover/failover completed.
|
|
|
|
.. note:
|
|
The maximum lifetime of the transaction in progress is controlled outside of this method.
|
|
|
|
.. note:
|
|
Read access to `self._in_flight` isn't protected because we know it can't be changed outside of our thread.
|
|
|
|
:param task: reference to a :class:`PgDistNode` object that represents a row to be updated/created.
|
|
:returns: `True` if the row was succesfully created/updated or transaction in progress
|
|
was committed as an indicator that the `self._pg_dist_node` cache should be updated,
|
|
or, if the new transaction was opened, this method returns `False`.
|
|
"""
|
|
|
|
if task.event == 'after_promote':
|
|
# The after_promote may happen without previous before_demote and/or
|
|
# before_promore. In this case we just call self.update_node() method.
|
|
# If there is a transaction in progress, it could be that it already did
|
|
# required changes and we can simply COMMIT.
|
|
if not self._in_flight or self._in_flight.host != task.host or self._in_flight.port != task.port:
|
|
self.update_node(task)
|
|
if self._in_flight:
|
|
self.query('COMMIT')
|
|
return True
|
|
else: # before_demote, before_promote
|
|
if task.timeout:
|
|
task.deadline = time.time() + task.timeout
|
|
if not self._in_flight:
|
|
self.query('BEGIN')
|
|
self.update_node(task)
|
|
return False
|
|
|
|
def process_tasks(self) -> None:
|
|
while True:
|
|
# Read access to `_in_flight` isn't protected because we know it can't be changed outside of our thread.
|
|
if not self._in_flight and not self.load_pg_dist_node():
|
|
break
|
|
|
|
i, task = self.pick_task()
|
|
if not task or i is None:
|
|
break
|
|
try:
|
|
update_cache = self.process_task(task)
|
|
except Exception as e:
|
|
logger.error('Exception when working with pg_dist_node: %r', e)
|
|
update_cache = None
|
|
with self._condition:
|
|
if self._tasks:
|
|
if update_cache:
|
|
self._pg_dist_node[task.group] = task
|
|
|
|
if update_cache is False: # an indicator that process_tasks has started a transaction
|
|
self._in_flight = task
|
|
else:
|
|
self._in_flight = None
|
|
|
|
if id(self._tasks[i]) == id(task):
|
|
self._tasks.pop(i)
|
|
task.wakeup()
|
|
|
|
def run(self) -> None:
|
|
while True:
|
|
try:
|
|
with self._condition:
|
|
if self._schedule_load_pg_dist_node:
|
|
timeout = -1
|
|
elif self._in_flight:
|
|
timeout = self._in_flight.deadline - time.time() if self._tasks else None
|
|
else:
|
|
timeout = -1 if self._tasks else None
|
|
|
|
if timeout is None or timeout > 0:
|
|
self._condition.wait(timeout)
|
|
elif self._in_flight:
|
|
logger.warning('Rolling back transaction. Last known status: %s', self._in_flight)
|
|
self.query('ROLLBACK')
|
|
self._in_flight = None
|
|
self.process_tasks()
|
|
except Exception:
|
|
logger.exception('run')
|
|
|
|
def _add_task(self, task: PgDistNode) -> bool:
|
|
with self._condition:
|
|
i = self.find_task_by_group(task.group)
|
|
|
|
# The `PgDistNode.timeout` == None is an indicator that it was scheduled from the sync_pg_dist_node().
|
|
if task.timeout is None:
|
|
# We don't want to override the already existing task created from REST API.
|
|
if i is not None and self._tasks[i].timeout is not None:
|
|
return False
|
|
|
|
# There is a little race condition with tasks created from REST API - the call made "before" the member
|
|
# key is updated in DCS. Therefore it is possible that :func:`sync_pg_dist_node` will try to create a
|
|
# task based on the outdated values of "state"/"role". To solve it we introduce an artificial timeout.
|
|
# Only when the timeout is reached new tasks could be scheduled from sync_pg_dist_node()
|
|
if self._in_flight and self._in_flight.group == task.group and self._in_flight.timeout is not None\
|
|
and self._in_flight.deadline > time.time():
|
|
return False
|
|
|
|
# Override already existing task for the same worker group
|
|
if i is not None:
|
|
if task != self._tasks[i]:
|
|
logger.debug('Overriding existing task: %s != %s', self._tasks[i], task)
|
|
self._tasks[i] = task
|
|
self._condition.notify()
|
|
return True
|
|
# Add the task to the list if Worker node state is different from the cached `pg_dist_node`
|
|
elif self._schedule_load_pg_dist_node or task != self._pg_dist_node.get(task.group)\
|
|
or self._in_flight and task.group == self._in_flight.group:
|
|
logger.debug('Adding the new task: %s', task)
|
|
self._tasks.append(task)
|
|
self._condition.notify()
|
|
return True
|
|
return False
|
|
|
|
def add_task(self, event: str, group: int, conn_url: str,
|
|
timeout: Optional[float] = None, cooldown: Optional[float] = None) -> Optional[PgDistNode]:
|
|
try:
|
|
r = urlparse(conn_url)
|
|
except Exception as e:
|
|
return logger.error('Failed to parse connection url %s: %r', conn_url, e)
|
|
host = r.hostname
|
|
if host:
|
|
port = r.port or 5432
|
|
task = PgDistNode(group, host, port, event, timeout=timeout, cooldown=cooldown)
|
|
return task if self._add_task(task) else None
|
|
|
|
def handle_event(self, cluster: Cluster, event: Dict[str, Any]) -> None:
|
|
if not self.is_alive():
|
|
return
|
|
|
|
worker = cluster.workers.get(event['group'])
|
|
if not (worker and worker.leader and worker.leader.name == event['leader'] and worker.leader.conn_url):
|
|
return
|
|
|
|
task = self.add_task(event['type'], event['group'],
|
|
worker.leader.conn_url,
|
|
event['timeout'], event['cooldown'] * 1000)
|
|
if task and event['type'] == 'before_demote':
|
|
task.wait()
|
|
|
|
def bootstrap(self) -> None:
|
|
if not isinstance(self._config, dict): # self.is_enabled()
|
|
return
|
|
|
|
conn_kwargs = self._postgresql.config.local_connect_kwargs
|
|
conn_kwargs['options'] = '-c synchronous_commit=local -c statement_timeout=0'
|
|
if self._config['database'] != self._postgresql.database:
|
|
conn = connect(**conn_kwargs)
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute('CREATE DATABASE {0}'.format(
|
|
quote_ident(self._config['database'], conn)).encode('utf-8'))
|
|
finally:
|
|
conn.close()
|
|
|
|
conn_kwargs['dbname'] = self._config['database']
|
|
conn = connect(**conn_kwargs)
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute('CREATE EXTENSION citus')
|
|
|
|
superuser = self._postgresql.config.superuser
|
|
params = {k: superuser[k] for k in ('password', 'sslcert', 'sslkey') if k in superuser}
|
|
if params:
|
|
cur.execute("INSERT INTO pg_catalog.pg_dist_authinfo VALUES"
|
|
"(0, pg_catalog.current_user(), %s)",
|
|
(self._postgresql.config.format_dsn(params),))
|
|
|
|
if self.is_coordinator():
|
|
r = urlparse(self._postgresql.connection_string)
|
|
cur.execute("SELECT pg_catalog.citus_set_coordinator_host(%s, %s, 'primary', 'default')",
|
|
(r.hostname, r.port or 5432))
|
|
finally:
|
|
conn.close()
|
|
|
|
def adjust_postgres_gucs(self, parameters: Dict[str, Any]) -> None:
|
|
if not self.is_enabled():
|
|
return
|
|
|
|
# citus extension must be on the first place in shared_preload_libraries
|
|
shared_preload_libraries = list(filter(
|
|
lambda el: el and el != 'citus',
|
|
[p.strip() for p in parameters.get('shared_preload_libraries', '').split(',')]))
|
|
parameters['shared_preload_libraries'] = ','.join(['citus'] + shared_preload_libraries)
|
|
|
|
# if not explicitly set Citus overrides max_prepared_transactions to max_connections*2
|
|
if parameters['max_prepared_transactions'] == 0:
|
|
parameters['max_prepared_transactions'] = parameters['max_connections'] * 2
|
|
|
|
# Resharding in Citus implemented using logical replication
|
|
parameters['wal_level'] = 'logical'
|
|
|
|
def ignore_replication_slot(self, slot: Dict[str, str]) -> bool:
|
|
if isinstance(self._config, dict) and self._postgresql.is_leader() and\
|
|
slot['type'] == 'logical' and slot['database'] == self._config['database']:
|
|
m = CITUS_SLOT_NAME_RE.match(slot['name'])
|
|
return bool(m and {'move': 'pgoutput', 'split': 'citus'}.get(m.group(1)) == slot['plugin'])
|
|
return False
|