diff --git a/src/DIRAC/AccountingSystem/DB/AccountingDB.py b/src/DIRAC/AccountingSystem/DB/AccountingDB.py index 01389df7171..7828cdb7996 100644 --- a/src/DIRAC/AccountingSystem/DB/AccountingDB.py +++ b/src/DIRAC/AccountingSystem/DB/AccountingDB.py @@ -5,6 +5,9 @@ import threading import time +from cachetools import cachedmethod, LRUCache, TTLCache, cached + + from DIRAC import S_ERROR, S_OK from DIRAC.Core.Base.DB import DB from DIRAC.Core.Utilities import DEncode, List, ThreadSafe, TimeUtilities @@ -15,11 +18,19 @@ class AccountingDB(DB): - def __init__(self, name="Accounting/AccountingDB", readOnly=False, parentLogger=None): + def __init__(self, name="Accounting/AccountingDB", readOnly=False, parentLogger=None, accounting_types=None): DB.__init__(self, "AccountingDB", name, parentLogger=parentLogger) + + # Cached method + self._addkeyvalue_cache = LRUCache(maxsize=1024) + self._addkeyvalue_lock = threading.Lock() + self._gettablename_cache = LRUCache(maxsize=128) + self._gettablename_lock = threading.Lock() + self.maxBucketTime = 604800 # 1 w self.autoCompact = False self.__readOnly = readOnly + self.__accounting_types = accounting_types if accounting_types else [] self.__doingCompaction = False self.__doingPendingLockTime = 0 self.__deadLockRetries = 2 @@ -128,6 +139,9 @@ def __loadCatalogFromDB(self): raise Exception(retVal["Message"]) for typesEntry in retVal["Value"]: typeName = typesEntry[0] + if self.__accounting_types and typeName not in self.__accounting_types: + self.log.info("Ignoring accounting type as not in the list", typeName) + continue keyFields = List.fromChar(typesEntry[1], ",") valueFields = List.fromChar(typesEntry[2], ",") bucketsLength = DEncode.decode(typesEntry[3].encode())[0] @@ -158,6 +172,8 @@ def loadPendingRecords(self): """ Load all records pending to insertion and generate threaded jobs """ + print(f"CHRIS addkeyvalue cache {self.__addKeyValue.cache.currsize}/{self.__addKeyValue.cache.maxsize}") + print(f"CHRIS gettablename{self._getTableName.cache.currsize}/{self._getTableName.cache.maxsize}") gSynchro.lock() try: now = time.time() @@ -185,7 +201,7 @@ def loadPendingRecords(self): % self.getWaitingRecordsLifeTime() ) req = "SELECT " - req += ",".join(sqlFields) + req += ", ".join([f"`{f}`" for f in sqlFields]) req += f" FROM {sqlTableName} " req += "WHERE taken = 0 or TIMESTAMPDIFF( SECOND, takenSince, UTC_TIMESTAMP() ) > %s " args = [self.getWaitingRecordsLifeTime()] @@ -197,6 +213,7 @@ def loadPendingRecords(self): "[PENDING] Error when trying to get pending records", f"for {typeName} : {result['Message']}", ) + self.__doingPendingLockTime = 0 return result self.log.info(f"[PENDING] Got {len(result['Value'])} pending records for type {typeName}") dbData = result["Value"] @@ -278,6 +295,9 @@ def registerType(self, name, definitionKeyFields, definitionAccountingFields, bu """ Register a new type """ + if self.__accounting_types and name not in self.__accounting_types: + self.log.info("Not registering accounting type as not in the list", name) + return S_OK(False) result = self.__loadTablesCreated() if not result["OK"]: @@ -439,6 +459,7 @@ def __getIdForKeyValue(self, typeName, keyName, keyValue, conn=False): return S_OK(retVal["Value"][0][0]) return S_ERROR(f"Key id {keyName} for value {keyValue} does not exist although it should") + @cachedmethod(lambda self: self._addkeyvalue_cache, lock=lambda self: self._addkeyvalue_lock) def __addKeyValue(self, typeName, keyName, keyValue): """ Adds a key value to a key table if not existant @@ -565,10 +586,12 @@ def __insertFromINTable(self, recordTuples): """ Do the real insert and delete from the in buffer table """ + if self.__readOnly: + return S_ERROR("ReadOnly mode enabled. No modification allowed") self.log.verbose("Received bundle to process", f"of {len(recordTuples)} elements") for record in recordTuples: iD, typeName, startTime, endTime, valuesList, insertionEpoch = record - result = self.insertRecordDirectly(typeName, startTime, endTime, valuesList) + result = self._insertRecordDirectly(typeName, startTime, endTime, valuesList) if not result["OK"]: req = "UPDATE " req += self._getTableName("in", typeName) @@ -583,12 +606,10 @@ def __insertFromINTable(self, recordTuples): if not result["OK"]: self.log.error("Can't delete row from the IN table", result["Message"]) - def insertRecordDirectly(self, typeName, startTime, endTime, valuesList): + def _insertRecordDirectly(self, typeName, startTime, endTime, valuesList): """ Add an entry to the type contents """ - if self.__readOnly: - return S_ERROR("ReadOnly mode enabled. No modification allowed") self.log.info( "Adding record", "for type %s\n [%s -> %s]" @@ -688,7 +709,7 @@ def __writeBuckets(self, typeName, buckets, keyValues, valuesList, connObj=False sqlFields.extend(self.dbCatalog[typeName]["keys"]) sqlFields.extend(self.dbCatalog[typeName]["values"]) sqlUpData = ["entriesInBucket=entriesInBucket+VALUES(entriesInBucket)"] - sqlUpData.extend([f"{x}={x}+VALUES({x})" for x in self.dbCatalog[typeName]["values"]]) + sqlUpData.extend([f"`{x}`=`{x}`+VALUES(`{x}`)" for x in self.dbCatalog[typeName]["values"]]) valueGroups = [] sqlValues = [] for bucketInfo in buckets: @@ -709,7 +730,7 @@ def __writeBuckets(self, typeName, buckets, keyValues, valuesList, connObj=False req = "INSERT INTO " req += self._getTableName("bucket", typeName) req += " (" - req += ",".join(sqlFields) + req += ", ".join([f"`{f}`" for f in sqlFields]) req += ") VALUES " req += ",".join(valueGroups) req += " ON DUPLICATE KEY UPDATE " @@ -1291,6 +1312,7 @@ def __commitTransaction(self, connObj): def __rollbackTransaction(self, connObj): return self._query("ROLLBACK", conn=connObj) + @cachedmethod(lambda self: self._gettablename_cache, lock=lambda self: self._gettablename_lock) def _getTableName(self, tableType, typeName, keyName=None): """ Generate table name diff --git a/src/DIRAC/AccountingSystem/DB/MultiAccountingDB.py b/src/DIRAC/AccountingSystem/DB/MultiAccountingDB.py index 64b1c85be6d..6e11d08a26d 100644 --- a/src/DIRAC/AccountingSystem/DB/MultiAccountingDB.py +++ b/src/DIRAC/AccountingSystem/DB/MultiAccountingDB.py @@ -6,10 +6,11 @@ class MultiAccountingDB: - def __init__(self, csPath, readOnly=False): + def __init__(self, csPath, readOnly=False, accounting_types=None): self.__csPath = csPath self.__readOnly = readOnly self.__dbByType = {} + self.__accountingTypes = accounting_types if accounting_types else [] self.__defaultDB = "AccountingDB/AccountingDB" self.__log = gLogger.getSubLogger(self.__class__.__name__) self.__generateDBs() @@ -17,7 +18,9 @@ def __init__(self, csPath, readOnly=False): def __generateDBs(self): self.__log.notice("Creating default AccountingDB...") - self.__allDBs = {self.__defaultDB: AccountingDB(readOnly=self.__readOnly)} + self.__allDBs = { + self.__defaultDB: AccountingDB(readOnly=self.__readOnly, accounting_types=self.__accountingTypes) + } result = gConfig.getOptionsDict(self.__csPath) if not result["OK"]: gLogger.verbose("No extra databases defined", f"in {self.__csPath}") diff --git a/src/DIRAC/AccountingSystem/Service/DataStoreHandler.py b/src/DIRAC/AccountingSystem/Service/DataStoreHandler.py index 69516b1b0e8..ed4f9ab36b4 100644 --- a/src/DIRAC/AccountingSystem/Service/DataStoreHandler.py +++ b/src/DIRAC/AccountingSystem/Service/DataStoreHandler.py @@ -28,7 +28,10 @@ class DataStoreHandler(RequestHandler): @classmethod def initializeHandler(cls, svcInfoDict): multiPath = PathFinder.getDatabaseSection("Accounting/MultiDB") - cls.__acDB = MultiAccountingDB(multiPath) + # we can focus on only some of the accoutning type + cls.accounting_types = getServiceOption(svcInfoDict, "AccountingTypes", []) + + cls.__acDB = MultiAccountingDB(multiPath, accounting_types=cls.accounting_types) # we can run multiple services in read only mode. In that case we do not bucket cls.runBucketing = getServiceOption(svcInfoDict, "RunBucketing", True) if cls.runBucketing: diff --git a/tests/Integration/AccountingSystem/Test_AccountingDB.py b/tests/Integration/AccountingSystem/Test_AccountingDB.py index 468eb50a008..4b85faa85cf 100644 --- a/tests/Integration/AccountingSystem/Test_AccountingDB.py +++ b/tests/Integration/AccountingSystem/Test_AccountingDB.py @@ -42,10 +42,10 @@ @pytest.fixture def inout(): - res = acDB.insertRecordDirectly("Pilot", startTime, middleTime, keyValues_1 + nonKeyValue_1) + res = acDB._insertRecordDirectly("Pilot", startTime, middleTime, keyValues_1 + nonKeyValue_1) assert res["OK"], res["Message"] - res = acDB.insertRecordDirectly("Pilot", middleTime, endTime, keyValues_2 + nonKeyValue_2) + res = acDB._insertRecordDirectly("Pilot", middleTime, endTime, keyValues_2 + nonKeyValue_2) assert res["OK"], res["Message"] yield