xadammr/znc-mssql
1
0
Fork 0
Code Releases Activity

improve threading, add modcommand, batch posts, default workers =4

Browse Source
This commit is contained in:
xadammr
2026-10-19 14:37:59 +10:00
parent 8531294ca5
commit eef1ce9cc2
1 changed files with 582 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mssql.py
+432 -83
View File
@@ -2,30 +2,57 @@
import znc
import re
import multiprocessing
import queue
import threading
import warnings
import traceback
import inspect
import os
import pprint
import json
import uuid
from datetime import datetime
from time import sleep
from contextlib import closing
class mssql(znc.Module):
module_types = [znc.CModInfo.GlobalModule]
description = "Log all channels to a Microsoft SQL Server database"
has_args = False
has_args = True
log_queue = multiprocessing.SimpleQueue()
log_queue = None
worker_threads = None
internal_log = None
hook_debugging = True
hook_debugging = False
debug_log_saves = False
default_workers = 4
num_workers = 4
max_workers = 32
batch_size = 50
db_config = None
next_worker_id = 0
def put_log(self, code, channel, host, user, ident, user_mode, date, target_user, message, network):
"""
Adds the log line to database write queue.
"""
regex = re.compile("\x1f|\x02|\x03(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE)
message = regex.sub("", message)
message = self.colourstrip(message)
decolouriser = re.compile(r"\x1f|\x02|\x1d|\x16|\x11|\x1e|\x0f|\x03(?:\d{1,2}(?:,\d{1,2})?)?|\x04(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE)
if message is not None:
message.encode('utf8', 'replace').decode('utf8')
message = decolouriser.sub("", message)
if self.log_queue is None:
return
queue_len = self.log_queue.qsize() if self.log_queue is not None else 0
if queue_len > 10:
with self.internal_log.error() as target:
target.write('Queue length high: {} entries at {}\n'.format(queue_len, datetime.utcnow()))
try:
self.log_queue.put({
'znc_user': self.GetUser().GetUserName() if self.GetUser() is not None else None,
'code': code,
@@ -34,35 +61,135 @@ class mssql(znc.Module):
'user': user,
'ident': ident,
'user_mode': user_mode,
'date': date,
'date': str(date),
'target_user': target_user,
'message': message,
'network': network
'network': network,
'_id': uuid.uuid4()
})
except Exception as e:
with self.internal_log.error() as target:
target.write('Failed to enqueue log event: {0} {1}\n'.format(type(e), str(e)))
target.write('Stack trace: ' + traceback.format_exc())
target.write('\n')
def debug_hook(self):
"""
Dumps parent calling method name and its arguments to debug logfile.
"""
if self.hook_debugging is not True:
return
frameinfo = inspect.stack()[1]
argvals = frameinfo.frame.f_locals
with self.internal_log.debug() as target:
target.write('Called method: ' + frameinfo.function + '()\n')
for argname in argvals:
if argname == 'self':
continue
target.write(' ' + argname + ' -> ' + pprint.pformat(argvals[argname]) + '\n')
target.write('\n')
def _start_workers(self, count):
"""Start `count` more worker threads on the shared queue."""
for _ in range(count):
idx = self.next_worker_id
self.next_worker_id += 1
db = MSSQLDatabase(self.db_config)
worker_thread = threading.Thread(
target=DatabaseThread.worker_safe,
args=(db, self.log_queue, self.internal_log, self.debug_log_saves, self.batch_size),
daemon=True,
name=f'DatabaseWorker-{idx}'
)
worker_thread.start()
self.worker_threads.append(worker_thread)
def OnModCommand(self, command):
"""
Commands: Status, AddWorkers <n>
"""
parts = str(command).split()
if not parts:
return
cmd = parts[0].lower()
if cmd == 'status':
alive = len([t for t in self.worker_threads if t.is_alive()]) if self.worker_threads else 0
qsize = self.log_queue.qsize() if self.log_queue is not None else 0
self.PutModule('Workers alive: {}, queue length: {}, batch size: {}'.format(alive, qsize, self.batch_size))
elif cmd == 'addworkers':
try:
n = int(parts[1])
except (IndexError, ValueError):
self.PutModule('Usage: AddWorkers <n>')
return
current = len(self.worker_threads)
n = max(0, min(n, self.max_workers - current))
if n == 0:
self.PutModule('No workers added (limit is {}, currently {})'.format(self.max_workers, current))
return
self._start_workers(n)
self.PutModule('Started {} workers, now {}'.format(n, len(self.worker_threads)))
else:
self.PutModule('Commands: Status, AddWorkers <n>')
def OnLoad(self, args, message):
"""
This module hook is called when a module is loaded.
:type args: const CString &
:type args: CString &
:rtype: bool
:param args: The arguments for the modules.
:param message: A message that may be displayed to the user after loading the module.
:return: True if the module loaded successfully, else False.
"""
# Initialize internal log first for debugging
self.internal_log = InternalLog(self.GetSavePath())
# Parse num_workers from args (default 4)
self.num_workers = self.default_workers
args_str = ""
if args:
try:
args_str = str(args) if not hasattr(args, 's') else args.s
if args_str:
self.num_workers = int(args_str.strip())
if self.num_workers < 1:
self.num_workers = 1
if self.num_workers > self.max_workers:
self.num_workers = self.max_workers
except (ValueError, AttributeError, TypeError) as e:
with self.internal_log.debug() as target:
target.write('Failed to parse args: {} (args_str: "{}", type: {})\n'.format(e, args_str, type(args)))
with self.internal_log.debug() as target:
target.write('mssql module loaded at: {}\n'.format(datetime.utcnow()))
target.write('save_path: {}\n'.format(self.GetSavePath()))
target.write('raw args: {} (type: {})\n'.format(repr(args), type(args).__name__))
target.write('parsed args_str: "{}"\n'.format(args_str))
target.write('num_workers: {}\n'.format(self.num_workers))
target.write('\n')
with self.internal_log.error() as target:
target.write('mssql module loaded at: {}\n'.format(datetime.utcnow()))
target.write('hook_debugging: {}\n'.format(self.hook_debugging))
target.write('num_workers: {}\n'.format(self.num_workers))
target.write('\n')
self.debug_hook()
try:
db = MSSQLDatabase({
'host': "db.spd.internal",
self.log_queue = queue.Queue()
self.db_config = {
'host': "dbx.spd.internal",
'user': "_srv_znc",
'passwd': "zncznc",
'db': "znc"
})
'password': "znc",
'database': "znc",
'port': '1433',
'appname': 'ZNC Logger'
}
self.worker_threads = []
self.next_worker_id = 0
self._start_workers(self.num_workers)
multiprocessing.Process(target=DatabaseThread.worker_safe,
args=(db, self.log_queue, self.internal_log)).start()
return True
except Exception as e:
message.s = str(e)
@@ -74,27 +201,108 @@ class mssql(znc.Module):
return False
def OnUnload(self):
if self.log_queue is None:
return
with self.internal_log.debug() as target:
target.write('Unloading module, draining queue length {}\n'.format(self.log_queue.qsize()))
# Send one sentinel value per worker thread
if self.worker_threads:
for _ in self.worker_threads:
self.log_queue.put(None)
self.log_queue.join()
if self.worker_threads is not None:
for worker_thread in self.worker_threads:
worker_thread.join(timeout=30)
with self.internal_log.debug() as target:
target.write('Unload complete, queue drained\n')
def __del__(self):
# Terminate worker threads cleanly if possible.
try:
if self.log_queue is not None:
# Send sentinel values for all workers
if self.worker_threads:
for _ in self.worker_threads:
self.log_queue.put_nowait(None)
except Exception:
pass
def OnIRCConnected(self):
"""
This module hook is called after a successful login to IRC.
:rtype: None
"""
self.debug_hook()
self.put_log("CONNECT", None, None, None, None, None, datetime.now(), None, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnIRCDisconnected(self):
"""
This module hook is called when a user gets disconnected from IRC.
:rtype: None
"""
self.debug_hook()
self.put_log("DISCONNECT", None, None, None, None, None, datetime.now(), None, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnOp(self, user, target_user, channel, noChange):
self.insert("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnChanNotice(self, nick, channel, message):
"""
Called when we receive a channel NOTICE message from IRC.
:type nick: CNick &
:type channel: CChan &
:type message: CString &
:rtype: EModRet
"""
self.debug_hook()
self.put_log("NOTICE", channel.GetName(), nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnPrivNotice(self, user, message):
self.insert("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnUserNotice(self, target, message):
"""
This module hook is called when a user sends a NOTICE message.
:type target: CString &
:type message: CString &
:rtype: EModRet
"""
self.debug_hook()
network = self.GetNetwork()
if network:
self.put_log("NOTICE", None, target.GetHost(), target.GetNick(), target.GetIdent(), None, datetime.now(), target.GetNick(), message, str(self.GetNetwork()))
def OnDeop(self, user, target_user, channel, noChange):
self.insert("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnVoice(self, user, target_user, channel, noChange):
self.insert("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnDevoice(self, user, target_user, channel, noChange):
self.insert("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnMode(self, user, channel, mode, argument, added, noChange):
self.debug_hook()
mode = chr(mode)
code = "MODE"
if mode == "b":
@@ -103,58 +311,149 @@ class mssql(znc.Module):
sign = "+" if added else "-"
argument = "{0}{1} {2}".format(sign, mode, argument)
self.insert(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork()))
return True
self.put_log(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork()))
return znc.CONTINUE
def OnQuit(self, user, message, channels):
self.debug_hook()
for channel in channels:
self.insert("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return True
self.put_log("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return znc.CONTINUE
def OnNick(self, user, new_nick, channels):
self.debug_hook()
for channel in channels:
self.insert("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork()))
return True
self.put_log("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnKick(self, user, target_nick, channel, message):
self.insert("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork()))
return znc.CONTINUE
def OnJoin(self, user, channel):
self.insert("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnPart(self, user, channel, message):
self.insert("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return znc.CONTINUE
def OnUserAction (self, target, message):
self.debug_hook()
(channel, user) = self.resolveTarget(target)
if channel == None:
return True
self.put_log("ME", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return znc.CONTINUE
return self.OnChanAction(user, channel, message)
def OnPrivAction(self, nick, message):
"""
Called when we receive a private CTCP ACTION ("/me" in query) from IRC.
:type nick: CNick &
:type message: CString &
:rtype: EModRet
"""
self.debug_hook()
self.put_log("ME", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnPrivMsg(self, nick, message):
self.insert("MSG", "", nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("MSG", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnPrivCTCP(self, nick, message):
self.debug_hook()
self.put_log("CTCP", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnChanCTCP(self, user, channel, message):
self.debug_hook()
self.put_log("CTCP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnUserCTCP(self, target, message):
self.debug_hook()
local_nick = self.GetNetwork().GetIRCNick()
if local_nick is not None:
self.put_log(
"CTCP",
None,
local_nick.GetHost(),
local_nick.GetNick(),
local_nick.GetIdent(),
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
else:
self.put_log(
"CTCP",
None,
None,
None,
None,
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
return znc.CONTINUE
def OnUserMsg (self, target, message):
self.debug_hook()
(channel, user) = self.resolveTarget(target)
if channel == None:
return True
local_nick = self.GetNetwork().GetIRCNick()
if local_nick is not None:
self.put_log(
"MSG",
None,
local_nick.GetHost(),
local_nick.GetNick(),
local_nick.GetIdent(),
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
else:
self.put_log(
"MSG",
None,
None,
None,
None,
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
return znc.CONTINUE
return self.OnChanMsg(user, channel, message)
def OnChanAction(self, user, channel, message):
self.insert("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
self.debug_hook()
self.put_log("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return True
def OnChanMsg(self, user, channel, message):
self.insert("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnTopic(self, user, channel, message):
self.insert("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return True
self.debug_hook()
self.put_log("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def findMode(self, channel, user):
realUser = channel.FindNick(user.GetNick())
@@ -203,6 +502,7 @@ class mssql(znc.Module):
except: pass
find = data.find('\x03')
data = data.replace('\x02','')
data = data.replace('\x1d','')
data = data.replace('\x1f','')
@@ -212,47 +512,74 @@ class mssql(znc.Module):
class DatabaseThread:
@staticmethod
def worker_safe(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None:
def worker_safe(db, log_queue, internal_log, debug_log_saves=False, batch_size=50) -> None:
while True:
try:
DatabaseThread.worker(db, log_queue, internal_log)
DatabaseThread.worker(db, log_queue, internal_log, debug_log_saves, batch_size)
break
except Exception as e:
with internal_log.error() as target:
target.write('Unrecoverable exception in worker thread: {0} {1}\n'.format(type(e), str(e)))
target.write('Worker thread crashed: {0} {1}\n'.format(type(e), str(e)))
target.write('Stack trace: ' + traceback.format_exc())
target.write('\n')
raise
target.write('Restarting worker thread in 5 seconds...\n\n')
sleep(5)
continue
@staticmethod
def worker(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None:
db.connect()
def worker(db, log_queue, internal_log, debug_log_saves=False, batch_size=50) -> None:
stop = False
while not stop:
first = log_queue.get()
if first is None:
log_queue.task_done() # count the sentinel too, or join() hangs
return
while True:
item = log_queue.get()
if item is None:
# Collect more rows without blocking, up to batch_size
taken = 1
items = [first]
while len(items) < batch_size:
try:
nxt = log_queue.get_nowait()
except queue.Empty:
break
taken += 1
if nxt is None:
stop = True
break
items.append(nxt)
try:
db.ensure_connected()
db.insert_into('logs', item)
DatabaseThread.save_batch(db, items, log_queue, internal_log, debug_log_saves)
finally:
# exactly one task_done() per get(), after the batch is finished
for _ in range(taken):
log_queue.task_done()
@staticmethod
def save_batch(db, items, log_queue, internal_log, debug_log_saves) -> None:
while True:
try:
db.ensure_connected(internal_log, str(items[0]['_id']))
db.insert_batch(items)
if debug_log_saves:
with internal_log.debug() as target:
target.write('Saved {} log items queue_size={}\n'.format(len(items), log_queue.qsize()))
return
except Exception as e:
sleep_for = 10
with internal_log.error() as target:
target.write('Could not save to database caused by: {0} {1}\n'.format(type(e), str(e)))
if 'open' in dir(db.conn):
target.write('Database handle state: {}\n'.format(db.conn.open))
target.write('Stack trace: ' + traceback.format_exc())
target.write('Current log: ')
json.dump(item, target)
target.write('Batch of {} rows, first log: '.format(len(items)))
item_copy = items[0].copy()
item_copy['_id'] = str(item_copy['_id'])
json.dump(item_copy, target)
target.write('\n\n')
target.write('Retry in {} s\n'.format(sleep_for))
sleep(sleep_for)
with internal_log.error() as target:
target.write('Retrying now.\n'.format(sleep_for))
log_queue.put(item)
class InternalLog:
def __init__(self, save_path: str):
self.save_path = save_path
@@ -265,9 +592,6 @@ class InternalLog:
def open(self, level: str):
target = open(os.path.join(self.save_path, level + '.log'), 'a')
line = 'Log opened at: {} UTC\n'.format(datetime.utcnow())
target.write(line)
target.write('=' * len(line) + '\n\n')
return target
@@ -282,13 +606,38 @@ class MSSQLDatabase(Database):
import pymssql
self.conn = pymssql.connect(**self.dsn)
def ensure_connected(self):
if self.conn.status == 0:
def ensure_connected(self, internal_log, log_id):
# No per-row SELECT 1. A dead connection is detected when an
# insert fails; insert_batch() then drops it and the retry reconnects.
if self.conn is None:
self.connect()
def insert_into(self, table, row):
# self.PutModule("Inserting row...")
cols = ', '.join('"{}"'.format(col) for col in row.keys())
vals = ', '.join('%({})s'.format(col) for col in row.keys())
self.conn.cursor().callproc("add_log", vals)
def insert_batch(self, rows):
"""Insert all rows, then commit once. All or nothing: on any error
the transaction is rolled back and the connection is dropped."""
try:
for row in rows:
params = (row['znc_user'], row['date'], row['code'], row['network'], row['channel'], row['host'], row['user'], row['ident'], row['user_mode'], row['target_user'], row['message'])
cur = self.conn.cursor()
try:
cur.callproc("add_log", params)
finally:
try:
cur.close()
except Exception:
pass
self.conn.commit()
except Exception:
try:
self.conn.rollback()
except Exception:
pass
try:
self.conn.close()
except Exception:
pass
self.conn = None
raise
def insert_into(self, table, row):
self.insert_batch([row])