bulltrackers-module 1.0.105 → 1.0.107

package/functions/task-engine/utils/firestore_batch_manager.js

@@ -1,259 +1,259 @@
-/**
- * @fileoverview Utility class to manage stateful Firestore write batches.
- * REFACTORED: Renamed 'firestore' to 'db' for consistency.
- * OPTIMIZED: Added logic to handle speculator timestamp fixes within the batch.
- */
-
-const { FieldValue } = require('@google-cloud/firestore');
-// No longer requires logger from sharedsetup, it's passed in.
-
-class FirestoreBatchManager {
-    /**
-     * @param {Firestore} db - A Firestore instance.
-     * @param {IntelligentHeaderManager} headerManager - An IntelligentHeaderManager instance.
-     * @param {object} logger - A logger instance.
-     * @param {object} config - Configuration object.
-     */
-    constructor(db, headerManager, logger, config) {
-        this.db = db; // Renamed from firestore
-        this.headerManager = headerManager;
-        this.logger = logger; // Added logger
-        this.config = config;
-
-        this.portfolioBatch = {};
-        this.timestampBatch = {};
-        this.processedSpeculatorCids = new Set();
-        this.speculatorTimestampFixBatch = {}; // <-- OPTIMIZATION: ADD THIS
-        this.batchTimeout = null;
-
-        this.logger.log('INFO', 'FirestoreBatchManager initialized.');
-    }
-
-    /**
-     * Schedules a flush if one isn't already scheduled.
-     */
-    _scheduleFlush() {
-        if (!this.batchTimeout) {
-            this.batchTimeout = setTimeout(
-                () => this.flushBatches(),
-                this.config.TASK_ENGINE_FLUSH_INTERVAL_MS
-            );
-        }
-    }
-
-    /**
-     * Adds a portfolio to the batch.
-     * @param {string} userId 
-     * @param {string} blockId 
-     * @param {string} date
-     * @param {object} portfolioData 
-     * @param {string} userType 
-     * @param {string|null} instrumentId 
-     */
-    async addToPortfolioBatch(userId, blockId, date, portfolioData, userType, instrumentId = null) {
-        const collection = userType === 'speculator'
-            ? this.config.FIRESTORE_COLLECTION_SPECULATOR_PORTFOLIOS
-            : this.config.FIRESTORE_COLLECTION_NORMAL_PORTFOLIOS;
-        const basePath = `${collection}/${blockId}/snapshots/${date}`;
-
-        if (!this.portfolioBatch[basePath]) {
-            this.portfolioBatch[basePath] = {};
-        }
-
-        this.portfolioBatch[basePath][userId] = portfolioData;
-
-        const totalUsersInBatch = Object.values(this.portfolioBatch).reduce((sum, users) => sum + Object.keys(users).length, 0);
-
-        if (totalUsersInBatch >= this.config.TASK_ENGINE_MAX_BATCH_SIZE) {
-            await this.flushBatches();
-        } else {
-            this._scheduleFlush();
-        }
-    }
-
-    /**
-     * Adds a user timestamp update to the batch.
-     * @param {string} userId 
-     * @param {string} userType 
-     * @param {string|null} instrumentId 
-     */
-    async updateUserTimestamp(userId, userType, instrumentId = null) {
-        const collection = userType === 'speculator'
-            ? this.config.FIRESTORE_COLLECTION_SPECULATOR_PORTFOLIOS
-            : this.config.FIRESTORE_COLLECTION_NORMAL_PORTFOLIOS;
-        const docId = userType === 'speculator' ? 'speculators' : 'normal';
-        const docPath = `${collection}/${docId}`;
-
-        if (!this.timestampBatch[docPath]) {
-            this.timestampBatch[docPath] = {};
-        }
-
-        const timestampKey = userType === 'speculator' ? `${userId}_${instrumentId}` : userId;
-        this.timestampBatch[docPath][timestampKey] = new Date();
-
-        if (Object.keys(this.timestampBatch[docPath]).length >= this.config.TASK_ENGINE_MAX_BATCH_SIZE) {
-            await this.flushBatches();
-        } else {
-            this._scheduleFlush();
-        }
-    }
-
-    /**
-     * Removes a user timestamp from the batch (e.g., if user is private).
-     * @param {string} userId 
-     * @param {string} userType 
-     * @param {string|null} instrumentId 
-     */
-    deleteFromTimestampBatch(userId, userType, instrumentId) {
-        const collection = userType === 'speculator'
-            ? this.config.FIRESTORE_COLLECTION_SPECULATOR_PORTFOLIOS
-            : this.config.FIRESTORE_COLLECTION_NORMAL_PORTFOLIOS;
-        const docId = userType === 'speculator' ? 'speculators' : 'normal';
-        const docPath = `${collection}/${docId}`;
-
-        if (this.timestampBatch[docPath]) {
-            const timestampKey = userType === 'speculator' ? `${userId}_${instrumentId}` : userId;
-            delete this.timestampBatch[docPath][timestampKey];
-        }
-    }
-
-    /**
-     * Adds discovered speculator CIDs to the in-memory set for later deletion.
-     * @param {Array<string>} cids 
-     */
-    addProcessedSpeculatorCids(cids) {
-        cids.forEach(cid => this.processedSpeculatorCids.add(cid));
-    }
-
-    /**
-     * --- OPTIMIZATION: ADD THIS NEW METHOD ---
-     * Adds a speculator timestamp fix to the batch.
-     * @param {string} userId
-     * @param {string} orchestratorBlockId (e.g., "1000000")
-     */
-    async addSpeculatorTimestampFix(userId, orchestratorBlockId) {
-        const docPath = `${this.config.FIRESTORE_COLLECTION_SPECULATOR_BLOCKS}/${orchestratorBlockId}`;
-
-        if (!this.speculatorTimestampFixBatch[docPath]) {
-            this.speculatorTimestampFixBatch[docPath] = {};
-        }
-
-        // These are the two fields the orchestrator reads
-        this.speculatorTimestampFixBatch[docPath][`users.${userId}.lastVerified`] = new Date();
-        this.speculatorTimestampFixBatch[docPath][`users.${userId}.lastHeldSpeculatorAsset`] = new Date();
-        
-        this.logger.log('TRACE', `[BATCH] Queued speculator timestamp fix for user ${userId} in block ${orchestratorBlockId}`);
-        this._scheduleFlush(); // Schedule a flush
-    }
-    // --- END NEW METHOD ---
-
-    /**
-     * Flushes all pending writes to Firestore and updates header performance.
-     */
-    async flushBatches() {
-        if (this.batchTimeout) {
-            clearTimeout(this.batchTimeout);
-            this.batchTimeout = null;
-        }
-
-        const promises = [];
-        const firestoreBatch = this.db.batch(); // Use this.db
-        let batchOperationCount = 0;
-
-        for (const basePath in this.portfolioBatch) {
-            const userPortfolios = this.portfolioBatch[basePath];
-            const userIds = Object.keys(userPortfolios);
-
-            if (userIds.length > 0) {
-                for (let i = 0; i < userIds.length; i += this.config.TASK_ENGINE_MAX_USERS_PER_SHARD) {
-                    const chunkUserIds = userIds.slice(i, i + this.config.TASK_ENGINE_MAX_USERS_PER_SHARD);
-                    const shardIndex = Math.floor(i / this.config.TASK_ENGINE_MAX_USERS_PER_SHARD);
-
-                    const chunkData = {};
-                    chunkUserIds.forEach(userId => {
-                        chunkData[userId] = userPortfolios[userId];
-                    });
-
-                    const docRef = this.db.collection(`${basePath}/parts`).doc(); // Now uses a reandomised ID from firestore to solve an overwrite issue with contention
-                    firestoreBatch.set(docRef, chunkData); // No merge: true needed, it's a new doc
-                    batchOperationCount++;
-                }
-                    // (Optional) Update the log message to reflect the path used for the data write
-                this.logger.log('INFO', `[BATCH] Staged ${userIds.length} users into ${Math.ceil(userIds.length / this.config.TASK_ENGINE_MAX_USERS_PER_SHARD)} unique shard documents for ${basePath}.`);
-            }
-            delete this.portfolioBatch[basePath];
-        }
-
-        for (const docPath in this.timestampBatch) {
-            if (Object.keys(this.timestampBatch[docPath]).length > 0) {
-                const [collection] = docPath.split('/');
-                const docRef = this.db.collection(collection).doc('timestamps').collection('users').doc('normal'); // Use this.db
-
-                if (batchOperationCount < 450) { 
-                    firestoreBatch.set(docRef, { users: this.timestampBatch[docPath] }, { merge: true });
-                    batchOperationCount++;
-                } else {
-                    promises.push(firestoreBatch.commit());
-                    const newBatch = this.db.batch(); // Use this.db
-                    newBatch.set(docRef, { users: this.timestampBatch[docPath] }, { merge: true });
-                    promises.push(newBatch.commit());
-                    batchOperationCount = 1;
-                }
-                delete this.timestampBatch[docPath];
-            }
-        }
-
-        if (this.processedSpeculatorCids.size > 0) {
-            this.logger.log('INFO', `[BATCH] Flushing ${this.processedSpeculatorCids.size} processed speculator CIDs from pending list documents.`);
-            const cidsToDelete = Array.from(this.processedSpeculatorCids);
-            this.processedSpeculatorCids.clear(); 
-
-            const pendingDocsSnapshot = await this.db.collection(this.config.PENDING_SPECULATORS_COLLECTION).get(); // Use this.db
-
-            if (!pendingDocsSnapshot.empty) {
-                const deletePromises = [];
-                pendingDocsSnapshot.forEach(doc => {
-                    const docData = doc.data().users || {};
-                    const cidsInThisDoc = cidsToDelete.filter(cid => docData.hasOwnProperty(cid));
-
-                    if (cidsInThisDoc.length > 0) {
-                        const deleteBatch = this.db.batch(); // Use this.db
-                        const updates = {};
-                        cidsInThisDoc.forEach(cid => {
-                            updates[`users.${cid}`] = FieldValue.delete();
-                        });
-                        deleteBatch.update(doc.ref, updates);
-                        deletePromises.push(deleteBatch.commit());
-                        this.logger.log('INFO', `[BATCH] Staged deletion of ${cidsInThisDoc.length} CIDs from document: ${doc.id}`);
-                    }
-                });
-                promises.push(...deletePromises);
-            }
-        }
-
-        // --- OPTIMIZATION: ADD THIS LOOP ---
-        for (const docPath in this.speculatorTimestampFixBatch) {
-            const updates = this.speculatorTimestampFixBatch[docPath];
-            if (Object.keys(updates).length > 0) {
-                const docRef = this.db.doc(docPath);
-                
-                // Add to the existing batch
-                firestoreBatch.set(docRef, updates, { merge: true });
-                batchOperationCount++;
-            }
-            delete this.speculatorTimestampFixBatch[docPath];
-        }
-        // --- END OPTIMIZATION ---
-
-        if (batchOperationCount > 0) {
-            promises.push(firestoreBatch.commit());
-        }
-
-        promises.push(this.headerManager.flushPerformanceUpdates());
-
-        await Promise.all(promises);
-        this.logger.log('INFO', '[BATCH] All batches flushed successfully.');
-    }
-}
-
+/**
+ * @fileoverview Utility class to manage stateful Firestore write batches.
+ * REFACTORED: Renamed 'firestore' to 'db' for consistency.
+ * OPTIMIZED: Added logic to handle speculator timestamp fixes within the batch.
+ */
+
+const { FieldValue } = require('@google-cloud/firestore');
+// No longer requires logger from sharedsetup, it's passed in.
+
+class FirestoreBatchManager {
+    /**
+     * @param {Firestore} db - A Firestore instance.
+     * @param {IntelligentHeaderManager} headerManager - An IntelligentHeaderManager instance.
+     * @param {object} logger - A logger instance.
+     * @param {object} config - Configuration object.
+     */
+    constructor(db, headerManager, logger, config) {
+        this.db = db; // Renamed from firestore
+        this.headerManager = headerManager;
+        this.logger = logger; // Added logger
+        this.config = config;
+
+        this.portfolioBatch = {};
+        this.timestampBatch = {};
+        this.processedSpeculatorCids = new Set();
+        this.speculatorTimestampFixBatch = {}; // <-- OPTIMIZATION: ADD THIS
+        this.batchTimeout = null;
+
+        this.logger.log('INFO', 'FirestoreBatchManager initialized.');
+    }
+
+    /**
+     * Schedules a flush if one isn't already scheduled.
+     */
+    _scheduleFlush() {
+        if (!this.batchTimeout) {
+            this.batchTimeout = setTimeout(
+                () => this.flushBatches(),
+                this.config.TASK_ENGINE_FLUSH_INTERVAL_MS
+            );
+        }
+    }
+
+    /**
+     * Adds a portfolio to the batch.
+     * @param {string} userId 
+     * @param {string} blockId 
+     * @param {string} date
+     * @param {object} portfolioData 
+     * @param {string} userType 
+     * @param {string|null} instrumentId 
+     */
+    async addToPortfolioBatch(userId, blockId, date, portfolioData, userType, instrumentId = null) {
+        const collection = userType === 'speculator'
+            ? this.config.FIRESTORE_COLLECTION_SPECULATOR_PORTFOLIOS
+            : this.config.FIRESTORE_COLLECTION_NORMAL_PORTFOLIOS;
+        const basePath = `${collection}/${blockId}/snapshots/${date}`;
+
+        if (!this.portfolioBatch[basePath]) {
+            this.portfolioBatch[basePath] = {};
+        }
+
+        this.portfolioBatch[basePath][userId] = portfolioData;
+
+        const totalUsersInBatch = Object.values(this.portfolioBatch).reduce((sum, users) => sum + Object.keys(users).length, 0);
+
+        if (totalUsersInBatch >= this.config.TASK_ENGINE_MAX_BATCH_SIZE) {
+            await this.flushBatches();
+        } else {
+            this._scheduleFlush();
+        }
+    }
+
+    /**
+     * Adds a user timestamp update to the batch.
+     * @param {string} userId 
+     * @param {string} userType 
+     * @param {string|null} instrumentId 
+     */
+    async updateUserTimestamp(userId, userType, instrumentId = null) {
+        const collection = userType === 'speculator'
+            ? this.config.FIRESTORE_COLLECTION_SPECULATOR_PORTFOLIOS
+            : this.config.FIRESTORE_COLLECTION_NORMAL_PORTFOLIOS;
+        const docId = userType === 'speculator' ? 'speculators' : 'normal';
+        const docPath = `${collection}/${docId}`;
+
+        if (!this.timestampBatch[docPath]) {
+            this.timestampBatch[docPath] = {};
+        }
+
+        const timestampKey = userType === 'speculator' ? `${userId}_${instrumentId}` : userId;
+        this.timestampBatch[docPath][timestampKey] = new Date();
+
+        if (Object.keys(this.timestampBatch[docPath]).length >= this.config.TASK_ENGINE_MAX_BATCH_SIZE) {
+            await this.flushBatches();
+        } else {
+            this._scheduleFlush();
+        }
+    }
+
+    /**
+     * Removes a user timestamp from the batch (e.g., if user is private).
+     * @param {string} userId 
+     * @param {string} userType 
+     * @param {string|null} instrumentId 
+     */
+    deleteFromTimestampBatch(userId, userType, instrumentId) {
+        const collection = userType === 'speculator'
+            ? this.config.FIRESTORE_COLLECTION_SPECULATOR_PORTFOLIOS
+            : this.config.FIRESTORE_COLLECTION_NORMAL_PORTFOLIOS;
+        const docId = userType === 'speculator' ? 'speculators' : 'normal';
+        const docPath = `${collection}/${docId}`;
+
+        if (this.timestampBatch[docPath]) {
+            const timestampKey = userType === 'speculator' ? `${userId}_${instrumentId}` : userId;
+            delete this.timestampBatch[docPath][timestampKey];
+        }
+    }
+
+    /**
+     * Adds discovered speculator CIDs to the in-memory set for later deletion.
+     * @param {Array<string>} cids 
+     */
+    addProcessedSpeculatorCids(cids) {
+        cids.forEach(cid => this.processedSpeculatorCids.add(cid));
+    }
+
+    /**
+     * --- OPTIMIZATION: ADD THIS NEW METHOD ---
+     * Adds a speculator timestamp fix to the batch.
+     * @param {string} userId
+     * @param {string} orchestratorBlockId (e.g., "1000000")
+     */
+    async addSpeculatorTimestampFix(userId, orchestratorBlockId) {
+        const docPath = `${this.config.FIRESTORE_COLLECTION_SPECULATOR_BLOCKS}/${orchestratorBlockId}`;
+
+        if (!this.speculatorTimestampFixBatch[docPath]) {
+            this.speculatorTimestampFixBatch[docPath] = {};
+        }
+
+        // These are the two fields the orchestrator reads
+        this.speculatorTimestampFixBatch[docPath][`users.${userId}.lastVerified`] = new Date();
+        this.speculatorTimestampFixBatch[docPath][`users.${userId}.lastHeldSpeculatorAsset`] = new Date();
+        
+        this.logger.log('TRACE', `[BATCH] Queued speculator timestamp fix for user ${userId} in block ${orchestratorBlockId}`);
+        this._scheduleFlush(); // Schedule a flush
+    }
+    // --- END NEW METHOD ---
+
+    /**
+     * Flushes all pending writes to Firestore and updates header performance.
+     */
+    async flushBatches() {
+        if (this.batchTimeout) {
+            clearTimeout(this.batchTimeout);
+            this.batchTimeout = null;
+        }
+
+        const promises = [];
+        const firestoreBatch = this.db.batch(); // Use this.db
+        let batchOperationCount = 0;
+
+        for (const basePath in this.portfolioBatch) {
+            const userPortfolios = this.portfolioBatch[basePath];
+            const userIds = Object.keys(userPortfolios);
+
+            if (userIds.length > 0) {
+                for (let i = 0; i < userIds.length; i += this.config.TASK_ENGINE_MAX_USERS_PER_SHARD) {
+                    const chunkUserIds = userIds.slice(i, i + this.config.TASK_ENGINE_MAX_USERS_PER_SHARD);
+                    const shardIndex = Math.floor(i / this.config.TASK_ENGINE_MAX_USERS_PER_SHARD);
+
+                    const chunkData = {};
+                    chunkUserIds.forEach(userId => {
+                        chunkData[userId] = userPortfolios[userId];
+                    });
+
+                    const docRef = this.db.collection(`${basePath}/parts`).doc(); // Now uses a reandomised ID from firestore to solve an overwrite issue with contention
+                    firestoreBatch.set(docRef, chunkData); // No merge: true needed, it's a new doc
+                    batchOperationCount++;
+                }
+                    // (Optional) Update the log message to reflect the path used for the data write
+                this.logger.log('INFO', `[BATCH] Staged ${userIds.length} users into ${Math.ceil(userIds.length / this.config.TASK_ENGINE_MAX_USERS_PER_SHARD)} unique shard documents for ${basePath}.`);
+            }
+            delete this.portfolioBatch[basePath];
+        }
+
+        for (const docPath in this.timestampBatch) {
+            if (Object.keys(this.timestampBatch[docPath]).length > 0) {
+                const [collection] = docPath.split('/');
+                const docRef = this.db.collection(collection).doc('timestamps').collection('users').doc('normal'); // Use this.db
+
+                if (batchOperationCount < 450) { 
+                    firestoreBatch.set(docRef, { users: this.timestampBatch[docPath] }, { merge: true });
+                    batchOperationCount++;
+                } else {
+                    promises.push(firestoreBatch.commit());
+                    const newBatch = this.db.batch(); // Use this.db
+                    newBatch.set(docRef, { users: this.timestampBatch[docPath] }, { merge: true });
+                    promises.push(newBatch.commit());
+                    batchOperationCount = 1;
+                }
+                delete this.timestampBatch[docPath];
+            }
+        }
+
+        if (this.processedSpeculatorCids.size > 0) {
+            this.logger.log('INFO', `[BATCH] Flushing ${this.processedSpeculatorCids.size} processed speculator CIDs from pending list documents.`);
+            const cidsToDelete = Array.from(this.processedSpeculatorCids);
+            this.processedSpeculatorCids.clear(); 
+
+            const pendingDocsSnapshot = await this.db.collection(this.config.PENDING_SPECULATORS_COLLECTION).get(); // Use this.db
+
+            if (!pendingDocsSnapshot.empty) {
+                const deletePromises = [];
+                pendingDocsSnapshot.forEach(doc => {
+                    const docData = doc.data().users || {};
+                    const cidsInThisDoc = cidsToDelete.filter(cid => docData.hasOwnProperty(cid));
+
+                    if (cidsInThisDoc.length > 0) {
+                        const deleteBatch = this.db.batch(); // Use this.db
+                        const updates = {};
+                        cidsInThisDoc.forEach(cid => {
+                            updates[`users.${cid}`] = FieldValue.delete();
+                        });
+                        deleteBatch.update(doc.ref, updates);
+                        deletePromises.push(deleteBatch.commit());
+                        this.logger.log('INFO', `[BATCH] Staged deletion of ${cidsInThisDoc.length} CIDs from document: ${doc.id}`);
+                    }
+                });
+                promises.push(...deletePromises);
+            }
+        }
+
+        // --- OPTIMIZATION: ADD THIS LOOP ---
+        for (const docPath in this.speculatorTimestampFixBatch) {
+            const updates = this.speculatorTimestampFixBatch[docPath];
+            if (Object.keys(updates).length > 0) {
+                const docRef = this.db.doc(docPath);
+                
+                // Add to the existing batch
+                firestoreBatch.set(docRef, updates, { merge: true });
+                batchOperationCount++;
+            }
+            delete this.speculatorTimestampFixBatch[docPath];
+        }
+        // --- END OPTIMIZATION ---
+
+        if (batchOperationCount > 0) {
+            promises.push(firestoreBatch.commit());
+        }
+
+        promises.push(this.headerManager.flushPerformanceUpdates());
+
+        await Promise.all(promises);
+        this.logger.log('INFO', '[BATCH] All batches flushed successfully.');
+    }
+}
+
 module.exports = { FirestoreBatchManager };