mirror of
https://github.com/outbackdingo/patroni.git
synced 2026-08-25 07:20:20 +00:00
Matt BakerandGitHub
b7ea511511
Expanding on the addition of docstrings in code, this adds python module API docs to sphinx documentation. A developer can preview what this might look like by running this locally: ``` tox -m docs ``` The option `-W` is added to the tox env so that warning messages are considered errors. Adds doc generation using the above method to the test GitHub workflow to catch documentation problems on PRs. Some docstrings have been reformatted and fixed to satisfy errors generated with the above setup.
422 lines
19 KiB
Python
422 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.get('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_primary() 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
|