#!/usr/bin/python3

import sys
import time
import argparse
import logging
import logging.handlers
import traceback
import pymysql
from GRMConst import MySQL_URL, DISCONNECT_TIMER
from dbutils.pooled_db import PooledDB

logging.basicConfig(
  level = logging.INFO,
  format = '%(asctime)s %(levelname)-8s %(message)s',
  datefmt = '%Y-%m-%d %H:%M:%S',
  handlers = [ logging.handlers.TimedRotatingFileHandler('/var/log/locker_cleanup.log', when='MIDNIGHT', backupCount=31, encoding='utf-8') ]
)

pymysql.install_as_MySQLdb()

MySQLdb = pymysql

pool = PooledDB(creator=MySQLdb, blocking=True, maxconnections=2, setsession=[ 'set session transaction isolation level read committed' ], **MySQL_URL)

logger = logging.getLogger(__name__)

def info(*args):
  logger.info(''.join(map(str,args)))

def error(*args):
  logger.error(''.join(map(str,args)))

BATCH = 100

class DB:
  def __enter__(self):
    self.db = pool.connection()
    self.cur = self.db.cursor()
    self.db.begin()
    return self.cur

  def __exit__(self, exc_type, exc_value, traceback):
    try:
      if exc_type is None:
        self.db.commit()
      else:
        self.db.rollback()
    finally:
      self.cur.close()
      self.db.close()

def cleanup_batch():
  with DB() as cur:
    cur.execute(
      "SELECT conn_id FROM conn "
      "WHERE disconnected < ADDTIME(NOW(),%s) "
      "ORDER BY conn_id LIMIT %s "
      "FOR UPDATE SKIP LOCKED",
      (DISCONNECT_TIMER, BATCH))
    rows = cur.fetchall()
    if not rows:
      return 0
    ids = [r[0] for r in rows]
    ph = ','.join(['%s'] * len(ids))
    cur.execute("DELETE FROM lock_file WHERE conn_id IN ({})".format(ph), ids)
    cur.execute("DELETE FROM conn      WHERE conn_id IN ({})".format(ph), ids)
    return len(ids)

def cleanup_once():
  total = 0
  while True:
    try:
      n = cleanup_batch()
    except:
      error(traceback.format_exc())
      return total
    total += n
    if n < BATCH:
      return total

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--interval', type=int, default=60,
                      help='実行間隔(秒)。0以下なら1回だけ実行して終了。')
  args = parser.parse_args()

  if args.interval <= 0:
    n = cleanup_once()
    if n:
      info('cleaned ', n, ' conn(s)')
    return

  while True:
    try:
      n = cleanup_once()
      if n:
        info('cleaned ', n, ' conn(s)')
    except:
      error(traceback.format_exc())
    time.sleep(args.interval)

main()