perf settings refactored

Author: max-ostapenko

infra/bigquery-export/firestore.js

@@ -4,11 +4,10 @@ import { BigQueryExport } from './bigquery.js'
 export class FirestoreBatch {
   constructor () {
     this.firestore = new Firestore({
-      // Increase timeout to 10 minutes for large batch operations
       gaxOptions: {
         grpc: {
-          max_receive_message_length: 100 * 1024 * 1024, // 100MB
-          max_send_message_length: 100 * 1024 * 1024, // 100MB
+          max_receive_message_length: 500 * 1024 * 1024, // 500MB
+          max_send_message_length: 500 * 1024 * 1024, // 500MB
           'grpc.max_connection_idle_ms': 5 * 60 * 1000, // 5 minutes
           'grpc.keepalive_time_ms': 30 * 1000, // 30 seconds
           'grpc.keepalive_timeout_ms': 60 * 1000, // 1 minute
@@ -17,12 +16,54 @@ export class FirestoreBatch {
       }
     })
     this.bigquery = new BigQueryExport()
-    this.batchSizeDelete = 500
-    this.batchSizeWrite = 400 // Reduced batch size for better performance
-    this.maxConcurrentBatches = 100 // Reduced concurrent batches to avoid overwhelming
+
+    // Configuration constants
+    this.config = {
+      batchSize: {
+        delete: 500,
+        write: 400
+      },
+      maxConcurrentBatches: 200,
+      retryCount: 5,
+      timeout: 10 * 60 * 1000 // 10 minutes
+    }
+
+    this.reset()
   }
 
-  queueBatch (operation) {
+  reset () {
+    this.currentBatch = []
+    this.batchPromises = []
+  }
+
+  getCurrentBatchSize (operation) {
+    return this.config.batchSize[operation === 'delete' ? 'delete' : 'write']
+  }
+
+  async commitWithRetry (batch, index) {
+    let lastError
+
+    for (let attempt = 1; attempt <= this.config.retryCount; attempt++) {
+      try {
+        await batch.commit()
+        return
+      } catch (error) {
+        lastError = error
+        console.warn(`Batch ${index} attempt ${attempt} failed:`, error.message)
+
+        if (attempt < this.config.retryCount) {
+          const delayMs = Math.pow(2, attempt) * 500
+          console.log(`Retrying batch ${index} in ${delayMs}ms...`)
+          await new Promise(resolve => setTimeout(resolve, delayMs))
+        }
+      }
+    }
+
+    console.error(`Batch ${index} failed after ${this.config.retryCount} attempts:`, lastError)
+    throw lastError
+  }
+
+  createBatch (operation) {
     const batch = this.firestore.batch()
 
     this.currentBatch.forEach((doc) => {
@@ -32,140 +73,126 @@ export class FirestoreBatch {
         const docRef = this.firestore.collection(this.collectionName).doc()
         batch.set(docRef, doc)
       } else {
-        throw new Error('Invalid operation')
+        throw new Error(`Invalid operation: ${operation}`)
       }
     })
+
+    return batch
+  }
+
+  queueBatch (operation) {
+    const batch = this.createBatch(operation)
     this.batchPromises.push(batch)
     this.currentBatch = []
   }
 
   async commitBatches () {
+    if (this.batchPromises.length === 0) return
+
     console.log(`Committing ${this.batchPromises.length} batches to ${this.collectionName}`)
 
     await Promise.all(
-      this.batchPromises.map(async (batchPromise, index) => {
-        const retryCount = 3
-        let lastError
-
-        for (let attempt = 1; attempt <= retryCount; attempt++) {
-          try {
-            await batchPromise.commit()
-            return
-          } catch (error) {
-            lastError = error
-            console.warn(`Batch ${index} attempt ${attempt} failed:`, error.message)
-
-            if (attempt < retryCount) {
-              // Exponential backoff: 2^attempt seconds
-              const delayMs = Math.pow(2, attempt) * 1000
-              console.log(`Retrying batch ${index} in ${delayMs}ms...`)
-              await new Promise(resolve => setTimeout(resolve, delayMs))
-            }
-          }
-        }
-
-        console.error(`Batch ${index} failed after ${retryCount} attempts:`, lastError)
-        throw lastError
-      })
+      this.batchPromises.map((batch, index) =>
+        this.commitWithRetry(batch, index)
+      )
     )
 
     this.batchPromises = []
   }
 
-  async finalFlush (operation) {
-    if (this.currentBatch.length > 0) {
+  async processInBatches (operation, shouldFlush = false) {
+    const batchSize = this.getCurrentBatchSize(operation)
+
+    if (this.currentBatch.length >= batchSize || shouldFlush) {
       this.queueBatch(operation)
     }
 
-    if (this.batchPromises.length > 0) {
+    if (this.batchPromises.length >= this.config.maxConcurrentBatches || shouldFlush) {
       await this.commitBatches()
     }
   }
 
+  buildQuery (collectionRef) {
+    const queryMap = {
+      report: () => {
+        console.info(`Deleting documents from ${this.collectionName} for date ${this.date}`)
+        return collectionRef.where('date', '==', this.date)
+      },
+      dict: () => {
+        console.info(`Deleting documents from ${this.collectionName}`)
+        return collectionRef
+      }
+    }
+
+    const queryBuilder = queryMap[this.collectionType]
+    if (!queryBuilder) {
+      throw new Error(`Invalid collection type: ${this.collectionType}`)
+    }
+
+    return queryBuilder()
+  }
+
   async batchDelete () {
     console.info('Starting batch deletion...')
     const startTime = Date.now()
-    this.currentBatch = []
-    this.batchPromises = []
+    this.reset()
 
     let totalDocsDeleted = 0
     const collectionRef = this.firestore.collection(this.collectionName)
-
-    let collectionQuery
-    if (this.collectionType === 'report') {
-      console.info('Deleting documents from ' + this.collectionName + ' for date ' + this.date)
-      // Query to fetch monthly documents
-      collectionQuery = collectionRef.where('date', '==', this.date)
-    } else if (this.collectionType === 'dict') {
-      console.info('Deleting documents from ' + this.collectionName)
-      collectionQuery = collectionRef
-    } else {
-      throw new Error('Invalid collection type')
-    }
+    const collectionQuery = this.buildQuery(collectionRef)
+    const batchSize = this.getCurrentBatchSize('delete')
 
     while (true) {
-      const snapshot = await collectionQuery.limit(this.batchSizeDelete * this.maxConcurrentBatches).get()
-      if (snapshot.empty) {
-        break
-      }
+      const snapshot = await collectionQuery.limit(batchSize * this.config.maxConcurrentBatches).get()
+      if (snapshot.empty) break
 
-      for await (const doc of snapshot.docs) {
+      for (const doc of snapshot.docs) {
         this.currentBatch.push(doc)
-
-        if (this.currentBatch.length >= this.batchSize) {
-          this.queueBatch('delete')
-        }
-        if (this.batchPromises.length >= this.maxConcurrentBatches) {
-          await this.commitBatches()
-        }
+        await this.processInBatches('delete')
         totalDocsDeleted++
       }
     }
-    await this.finalFlush('delete')
+
+    // Final flush
+    await this.processInBatches('delete', true)
 
     const duration = (Date.now() - startTime) / 1000
     console.info(`Deletion complete. Total docs deleted: ${totalDocsDeleted}. Time: ${duration} seconds`)
   }
 
-  /**
-   * Streams BigQuery query results into a Firestore collection using batch commits.
-   * @param {string} query - The BigQuery SQL query.
-   */
   async streamFromBigQuery (rowStream) {
     console.info('Starting BigQuery to Firestore transfer...')
     const startTime = Date.now()
     let totalRowsProcessed = 0
 
-    this.currentBatch = []
-    this.batchPromises = []
+    this.reset()
 
     for await (const row of rowStream) {
       this.currentBatch.push(row)
-
-      // Write batch when it reaches specified size
-      if (this.currentBatch.length >= this.batchSize) {
-        this.queueBatch('set')
-      }
-
-      if (this.batchPromises.length >= this.maxConcurrentBatches) {
-        await this.commitBatches()
-      }
+      await this.processInBatches('set')
       totalRowsProcessed++
     }
-    await this.finalFlush('set')
+
+    // Final flush
+    await this.processInBatches('set', true)
 
     const duration = (Date.now() - startTime) / 1000
     console.info(`Transfer to ${this.collectionName} complete. Total rows processed: ${totalRowsProcessed}. Time: ${duration} seconds`)
   }
 
   async export (query, exportConfig) {
+    // Configure Firestore settings
     this.firestore.settings({
       databaseId: exportConfig.database,
-      timeout: 10 * 60 * 1000 // 10 minutes timeout
+      timeout: this.config.timeout
+    })
+
+    // Set instance properties
+    Object.assign(this, {
+      collectionName: exportConfig.collection,
+      collectionType: exportConfig.type,
+      date: exportConfig.date
     })
-    this.collectionName = exportConfig.collection
-    this.collectionType = exportConfig.type
-    this.date = exportConfig.date
 
     await this.batchDelete()
 

infra/tf/.terraform.lock.hcl

@@ -35,7 +35,7 @@ resource "google_cloud_run_v2_job" "bigquery_export" {
           value = ""
         }
       }
-      timeout         = "3600s"
+      timeout         = "7200s"
       service_account = var.function_identity
       max_retries     = 1
     }