"""Cascaded worker. CascadedConsumer that also maintains node. """ import time import datetime from typing import Sequence, List, Dict, Optional, cast import skytools from skytools.basetypes import Cursor, Connection, DictRow from pgq.cascade.consumer import CascadedConsumer from pgq.baseconsumer import BatchInfo, EventList from pgq.event import Event from pgq.producer import bulk_insert_events __all__ = ['CascadedWorker'] class WorkerState(object): """Depending on node state decides on actions worker needs to do.""" # node_type, # node_name, provider_node, # global_watermark, local_watermark # combined_queue, combined_type process_batch = 0 # handled in CascadedConsumer copy_events = 0 # ok global_wm_event = 0 # ok local_wm_publish = 1 # ok process_events = 0 # ok send_tick_event = 0 # ok wait_behind = 0 # ok process_tick_event = 0 # ok target_queue = '' # ok keep_event_ids = 0 # ok create_tick = 0 # ok filtered_copy = 0 # ok process_global_wm = 0 # ok sync_watermark = 0 # ? wm_sync_nodes: Sequence[str] = [] def __init__(self, queue_name: str, nst: DictRow) -> None: self.node_type = nst['node_type'] self.node_name = nst['node_name'] self.local_watermark = nst['local_watermark'] self.global_watermark = nst['global_watermark'] self.node_attrs = {} attrs = nst.get('node_attrs', '') if attrs: self.node_attrs = skytools.db_urldecode(attrs) ntype = nst['node_type'] ctype = nst['combined_type'] if ntype == 'root': self.global_wm_event = 1 self.local_wm_publish = 0 elif ntype == 'branch': self.target_queue = queue_name self.process_batch = 1 self.process_events = 1 self.copy_events = 1 self.process_tick_event = 1 self.keep_event_ids = 1 self.create_tick = 1 if 'sync_watermark' in self.node_attrs: slist = self.node_attrs['sync_watermark'] self.sync_watermark = 1 self.wm_sync_nodes = slist.split(',') if slist else [] else: self.process_global_wm = 1 elif ntype == 'leaf' and not ctype: self.process_batch = 1 self.process_events = 1 elif ntype == 'leaf' and ctype: self.target_queue = nst['combined_queue'] if ctype == 'root': self.process_batch = 1 self.process_events = 1 self.copy_events = 1 self.filtered_copy = 1 self.send_tick_event = 1 elif ctype == 'branch': self.process_batch = 1 self.wait_behind = 1 else: raise Exception('invalid state 1') else: raise Exception('invalid state 2') if ctype and ntype != 'leaf': raise Exception('invalid state 3') class CascadedWorker(CascadedConsumer): """CascadedWorker base class. Config fragment:: ## Parameters for pgq.CascadedWorker ## # how often the root node should push wm downstream (seconds) #global_wm_publish_period = 300 # how often the nodes should report their wm upstream (seconds) #local_wm_publish_period = 300 """ global_wm_publish_time: float = 0 global_wm_publish_period: float = 5 * 60 local_wm_publish_time: float = 0 local_wm_publish_period: float = 5 * 60 max_evbuf: int = 500 cur_event_seq: int = 0 cur_max_id: int = 0 seq_buffer: int = 10000 main_worker: bool = True _worker_state: Optional[WorkerState] = None ev_buf: List[Event] = [] real_global_wm: Optional[int] = None was_wait_behind: bool = False def reload(self) -> None: super().reload() self.global_wm_publish_period = self.cf.getfloat('global_wm_publish_period', CascadedWorker.global_wm_publish_period) self.local_wm_publish_period = self.cf.getfloat('local_wm_publish_period', CascadedWorker.local_wm_publish_period) def _load_next_batch(self, curs: Cursor) -> Optional[int]: # handle wait_behind without blocking whole loop batch_id = super()._load_next_batch(curs) wst = self._worker_state cst = self._consumer_state info = self.batch_info if batch_id is not None and wst and cst and info and wst.wait_behind: tick_id = info['tick_id'] completed_tick_id = cst['completed_tick'] ahead = tick_id > completed_tick_id if ahead: # pretend no batch was available batch_id = None self.log.debug("Wait behind: tick_id=%r completed=%r", tick_id, completed_tick_id) return batch_id def process_batch(self, db: Connection, batch_id: int, event_list: EventList) -> None: # switch to alternative path for wait_behind wst = self._worker_state if wst and wst.wait_behind: dst_db = self.get_database(self.target_db) self.process_wait_behind(db, batch_id, event_list, dst_db) else: super().process_batch(db, batch_id, event_list) def process_wait_behind(self, src_db: Connection, batch_id: int, event_list: EventList, dst_db: Connection) -> None: # data events are already applied to target from main (target) queue, # need to process system events from events for leaf node. src_curs = src_db.cursor() dst_curs = dst_db.cursor() for ev in event_list: if ev.ev_type.split('.', 1)[0] in ("pgq", "londiste"): self.process_remote_event(src_curs, dst_curs, ev) st = self._worker_state assert st if st.local_wm_publish and self.main_worker: self.publish_local_wm(src_db, dst_db) def process_remote_batch(self, src_db: Connection, tick_id: int, event_list: EventList, dst_db: Connection) -> None: """Worker-specific event processing.""" self.ev_buf = [] max_id = 0 assert self._worker_state st = self._worker_state src_curs = src_db.cursor() dst_curs = dst_db.cursor() for ev in event_list: if st.copy_events: self.copy_event(dst_curs, ev, st.filtered_copy) if ev.ev_type.split('.', 1)[0] in ("pgq", "londiste"): # process cascade events even on waiting leaf node self.process_remote_event(src_curs, dst_curs, ev) else: if st.process_events: self.process_remote_event(src_curs, dst_curs, ev) if ev.ev_id > max_id: max_id = ev.ev_id if max_id > self.cur_max_id: self.cur_max_id = max_id def is_batch_done(self, state: DictRow, batch_info: BatchInfo, dst_db: Connection) -> bool: # handle combined_queue type change (branch->root) if self._worker_state and self.was_wait_behind: cur_tick = batch_info['tick_id'] dst_tick = state['completed_tick'] if cur_tick <= dst_tick: # current batch is already applied, skip it return True if not self._worker_state.wait_behind: # forget previous state self.was_wait_behind = False wst = self._worker_state assert wst # check if events have processed done = super().is_batch_done(state, batch_info, dst_db) if not wst.create_tick: return done if not done: return False # check if tick is done - it happens in separate tx # fetch last tick from target queue q = "select t.tick_id from pgq.tick t, pgq.queue q"\ " where t.tick_queue = q.queue_id and q.queue_name = %s"\ " order by t.tick_queue desc, t.tick_id desc"\ " limit 1" curs = dst_db.cursor() curs.execute(q, [self.queue_name]) last_tick = curs.fetchone()['tick_id'] dst_db.commit() # insert tick if missing cur_tick = batch_info['tick_id'] if last_tick != cur_tick: prev_tick = batch_info['prev_tick_id'] tick_time = batch_info['batch_end'] if last_tick != prev_tick: raise Exception('is_batch_done: last branch tick = %d, expected %d or %d' % ( last_tick, prev_tick, cur_tick)) self.create_branch_tick(dst_db, cur_tick, tick_time) return True def publish_local_wm(self, src_db: Connection, dst_db: Connection) -> None: """Send local watermark to provider. """ t = time.time() if t - self.local_wm_publish_time < self.local_wm_publish_period: return st = self._worker_state assert st assert self.batch_info wm = st.local_watermark if st.sync_watermark: # dont send local watermark upstream wm = self.batch_info['prev_tick_id'] elif wm > self.batch_info['cur_tick_id']: # in wait-behind-leaf case, the wm from target can be # ahead from source queue, use current batch then wm = self.batch_info['cur_tick_id'] self.log.debug("Publishing local watermark: %d", wm) src_curs = src_db.cursor() q = "select * from pgq_node.set_subscriber_watermark(%s, %s, %s)" src_curs.execute(q, [self.pgq_queue_name, st.node_name, wm]) src_db.commit() # if next part fails, dont repeat it immediately self.local_wm_publish_time = t if st.sync_watermark and self.real_global_wm is not None: # instead sync 'global-watermark' with specific nodes dst_curs = dst_db.cursor() nmap = self._get_node_map(dst_curs) dst_db.commit() # local lowest wm = st.local_watermark # the global-watermark in subtree can stay behind # upstream global-watermark, but must not go ahead if self.real_global_wm < wm: wm = self.real_global_wm for node in st.wm_sync_nodes: if node == st.node_name: continue if node not in nmap: # dont ignore missing nodes - cluster may be partially set up self.log.warning('Unknown node in sync_watermark list: %s', node) return n = nmap[node] if n['dead']: # ignore dead nodes continue wmdb = self.get_database('wmdb', connstr=n['node_location'], autocommit=1, profile='remote') wmcurs = wmdb.cursor() q = 'select local_watermark from pgq_node.get_node_info(%s)' wmcurs.execute(q, [self.queue_name]) row = wmcurs.fetchone() if not row: # partially set up node? self.log.warning('Node not working: %s', node) elif row['local_watermark'] < wm: # keep lowest wm wm = row['local_watermark'] self.close_database('wmdb') # now we have lowest wm, store it q = "select pgq_node.set_global_watermark(%s, %s)" dst_curs.execute(q, [self.queue_name, wm]) dst_db.commit() def _get_node_map(self, curs: Cursor) -> Dict[str, DictRow]: q = "select node_name, node_location, dead from pgq_node.get_queue_locations(%s)" curs.execute(q, [self.queue_name]) res = {} for row in curs.fetchall(): res[row['node_name']] = row return res def process_remote_event(self, src_curs: Cursor, dst_curs: Cursor, ev: Event) -> None: """Handle cascading events. """ if ev.retry: raise Exception('CascadedWorker must not get retry events') # non cascade events send to CascadedConsumer to error out if ev.ev_type[:4] != 'pgq.': super().process_remote_event(src_curs, dst_curs, ev) return # ignore cascade events if not main worker if not self.main_worker: return # check if for right queue t = ev.ev_type if ev.ev_extra1 != self.pgq_queue_name and t != "pgq.tick-id": raise Exception("bad event in queue: " + str(ev)) self.log.debug("got cascade event: %s(%s)", t, ev.ev_data) st = self._worker_state assert st if t == "pgq.location-info": node = ev.ev_data loc = ev.ev_extra2 dead = ev.ev_extra3 q = "select * from pgq_node.register_location(%s, %s, %s, %s)" dst_curs.execute(q, [self.pgq_queue_name, node, loc, dead]) elif t == "pgq.unregister-location": node = ev.ev_data q = "select * from pgq_node.unregister_location(%s, %s)" dst_curs.execute(q, [self.pgq_queue_name, node]) elif t == "pgq.global-watermark": if st.sync_watermark: tick_id = int(ev.ev_data) self.log.debug('Half-ignoring global watermark %d', tick_id) self.real_global_wm = tick_id elif st.process_global_wm: tick_id = int(ev.ev_data) q = "select * from pgq_node.set_global_watermark(%s, %s)" dst_curs.execute(q, [self.pgq_queue_name, tick_id]) elif t == "pgq.tick-id": tick_id = int(ev.ev_data) if ev.ev_extra1 == self.pgq_queue_name: raise Exception('tick-id event for own queue?') if st.process_tick_event: q = "select * from pgq_node.set_partition_watermark(%s, %s, %s)" dst_curs.execute(q, [self.pgq_queue_name, ev.ev_extra1, tick_id]) else: raise Exception("unknown cascade event: %s" % t) def finish_remote_batch(self, src_db: Connection, dst_db: Connection, tick_id: int) -> None: """Worker-specific cleanup on target node. """ # merge-leaf on branch should not update tick pos st = self._worker_state assert st if self.main_worker: dst_curs = dst_db.cursor() self.flush_events(dst_curs) # send tick event into queue if st.send_tick_event: q = "select pgq.insert_event(%s, 'pgq.tick-id', %s, %s, null, null, null)" dst_curs.execute(q, [st.target_queue, str(tick_id), self.pgq_queue_name]) super().finish_remote_batch(src_db, dst_db, tick_id) if self.main_worker: if st.create_tick and self.batch_info: # create actual tick tick_id = self.batch_info['tick_id'] tick_time = self.batch_info['batch_end'] self.create_branch_tick(dst_db, tick_id, tick_time) if st.local_wm_publish: self.publish_local_wm(src_db, dst_db) def create_branch_tick(self, dst_db: Connection, tick_id: int, tick_time: datetime.datetime) -> None: q = "select pgq.ticker(%s, %s, %s, %s)" # execute it in autocommit mode ilev = dst_db.isolation_level dst_db.set_isolation_level(0) dst_curs = dst_db.cursor() dst_curs.execute(q, [self.pgq_queue_name, tick_id, tick_time, self.cur_max_id]) dst_db.set_isolation_level(ilev) def copy_event(self, dst_curs: Cursor, ev: Event, filtered_copy: int) -> None: """Add event to copy buffer. """ if not self.main_worker: return if filtered_copy: if ev.type[:4] == "pgq.": return if len(self.ev_buf) >= self.max_evbuf: self.flush_events(dst_curs) if ev.type == 'pgq.global-watermark': st = self._worker_state if st and st.sync_watermark: # replace payload with synced global watermark row = dict(ev._event_row.items()) row['ev_data'] = str(st.global_watermark) ev = Event(self.queue_name, cast(DictRow, row)) self.ev_buf.append(ev) def flush_events(self, dst_curs: Cursor) -> None: """Send copy buffer to target queue. """ if len(self.ev_buf) == 0: return flds = ['ev_time', 'ev_type', 'ev_data', 'ev_extra1', 'ev_extra2', 'ev_extra3', 'ev_extra4'] st = self._worker_state assert st if st.keep_event_ids: flds.append('ev_id') bulk_insert_events(dst_curs, self.ev_buf, flds, st.target_queue) self.ev_buf = [] def refresh_state(self, dst_db: Connection, full_logic: bool = True) -> DictRow: """Load also node state from target node. """ queue_name = self.pgq_queue_name or '?' res = super().refresh_state(dst_db, full_logic) q = "select * from pgq_node.get_node_info(%s)" rows = self.exec_cmd(dst_db, q, [queue_name]) assert rows self._worker_state = WorkerState(queue_name, rows[0]) if self._worker_state.wait_behind: self.was_wait_behind = True return res def process_root_node(self, dst_db: Connection) -> None: """On root node send global watermark downstream. """ super().process_root_node(dst_db) t = time.time() if t - self.global_wm_publish_time < self.global_wm_publish_period: return self.log.debug("Publishing global watermark") dst_curs = dst_db.cursor() q = "select * from pgq_node.set_global_watermark(%s, NULL)" dst_curs.execute(q, [self.pgq_queue_name]) dst_db.commit() self.global_wm_publish_time = t