memory management

Author: max-ostapenko

infra/bigquery-export/firestore.js

@@ -20,22 +20,64 @@ export class FirestoreBatch {
     // Configuration constants
     this.config = {
       timeout: 10 * 60 * 1000, // 10 minutes
-      progressReportInterval: 200000, // Report progress every N operations
-      flushThreshold: 200000 // Flush BulkWriter every N operations
+      progressReportInterval: 100000, // Report progress every N operations
+      flushThreshold: 50000, // Flush BulkWriter every N operations
+      gcInterval: 50000 // Force garbage collection interval
     }
 
     this.reset()
   }
 
+  // Memory monitoring utility
+  logMemoryUsage (operation = '') {
+    const used = process.memoryUsage()
+    const memoryInfo = {
+      rss: Math.round(used.rss / 1024 / 1024 * 100) / 100,
+      heapTotal: Math.round(used.heapTotal / 1024 / 1024 * 100) / 100,
+      heapUsed: Math.round(used.heapUsed / 1024 / 1024 * 100) / 100,
+      external: Math.round(used.external / 1024 / 1024 * 100) / 100
+    }
+
+    console.log(`Memory usage ${operation}: RSS ${memoryInfo.rss}MB, Heap Used ${memoryInfo.heapUsed}MB, Heap Total ${memoryInfo.heapTotal}MB, External ${memoryInfo.external}MB`)
+
+    // Configurable memory warning threshold from environment
+    const warningThreshold = parseInt(process.env.MEMORY_WARNING_THRESHOLD_MB || '1500')
+    if (memoryInfo.heapUsed > warningThreshold) {
+      console.warn(`⚠️ High memory usage detected: ${memoryInfo.heapUsed}MB heap used (threshold: ${warningThreshold}MB)`)
+    }
+
+    return memoryInfo
+  }
+
+  // Enhanced reset with memory cleanup
   reset () {
     this.processedDocs = 0
     this.totalDocs = 0
+
+    // Clean up existing BulkWriter if it exists
+    if (this.bulkWriter) {
+      try {
+        this.bulkWriter.close()
+      } catch (error) {
+        console.warn('Error closing existing BulkWriter:', error.message)
+      }
+    }
     this.bulkWriter = null
+
+    // Force garbage collection if available
+    if (global.gc) {
+      global.gc()
+    }
+
+    // Log memory usage after reset
+    this.logMemoryUsage('after reset')
   }
 
   createBulkWriter (operation) {
     const bulkWriter = this.firestore.bulkWriter()
 
+    bulkWriter.maxBatchSize = 500 // Reduce batch size for memory efficiency
+
     // Configure error handling with progress info
     bulkWriter.onWriteError((error) => {
       const progressInfo = this.totalDocs > 0 ? ` (${this.processedDocs}/${this.totalDocs})` : ''
@@ -54,6 +96,11 @@ export class FirestoreBatch {
       if (this.processedDocs % this.config.progressReportInterval === 0) {
         const progressInfo = this.totalDocs > 0 ? ` (${this.processedDocs}/${this.totalDocs})` : ` (${this.processedDocs} processed)`
         console.log(`Progress${progressInfo} - ${operation}ing documents in ${this.collectionName}`)
+
+        // Force garbage collection periodically
+        if (this.processedDocs % this.config.gcInterval === 0 && global.gc) {
+          global.gc()
+        }
       }
     })
 
@@ -115,20 +162,26 @@ export class FirestoreBatch {
       if (snapshot.empty) break
 
       // Add all delete operations to BulkWriter
-      snapshot.docs.forEach(doc => {
+      for (const doc of snapshot.docs) {
         this.bulkWriter.delete(doc.ref)
         deletedCount++
-      })
 
-      // Periodically flush to manage memory
-      // if (deletedCount % this.config.flushThreshold === 0) {
-      console.log(`Flushing BulkWriter at ${deletedCount} operations...`)
-      await this.bulkWriter.flush()
-      // }
+        // Frequent flushing to prevent memory buildup
+        if (deletedCount % this.config.flushThreshold === 0) {
+          console.log(`Flushing BulkWriter at ${deletedCount} operations...`)
+          await this.bulkWriter.flush()
+
+          // Force garbage collection after flush
+          if (global.gc) {
+            global.gc()
+          }
+        }
+      }
     }
 
     // Final flush and close
     console.log('Finalizing deletion operations...')
+    await this.bulkWriter.flush()
     await this.bulkWriter.close()
 
     const duration = (Date.now() - startTime) / 1000
@@ -144,32 +197,52 @@ export class FirestoreBatch {
     this.bulkWriter = this.createBulkWriter('writ')
 
     let rowCount = 0
+    let batchCount = 0
     const collectionRef = this.firestore.collection(this.collectionName)
 
-    for await (const row of rowStream) {
-      // Add document to BulkWriter
-      const docRef = collectionRef.doc()
-      this.bulkWriter.set(docRef, row)
-
-      rowCount++
-      this.totalDocs = rowCount // Update total as we go since we can't predict BigQuery result size
-
-      // Periodically flush to manage memory
-      if (rowCount % this.config.flushThreshold === 0) {
-        console.log(`Flushing BulkWriter at ${rowCount} operations...`)
-        await this.bulkWriter.flush()
+    try {
+      for await (const row of rowStream) {
+        // Add document to BulkWriter
+        const docRef = collectionRef.doc()
+        this.bulkWriter.set(docRef, row)
+
+        rowCount++
+        this.totalDocs = rowCount // Update totalDocs for progress tracking
+
+        if (rowCount % this.config.flushThreshold === 0) {
+          console.log(`Flushing BulkWriter at ${rowCount} operations...`)
+          await this.bulkWriter.flush()
+          batchCount++
+
+          if (batchCount % 5 === 0 && global.gc) {
+            console.log(`Forcing garbage collection at batch ${batchCount}...`)
+            global.gc()
+          }
+        }
       }
+    } catch (error) {
+      console.error('Error during BigQuery streaming:', error)
+      throw error
     }
 
     // Final flush and close
     console.log('Finalizing write operations...')
+    await this.bulkWriter.flush()
     await this.bulkWriter.close()
 
+    // Final garbage collection
+    if (global.gc) {
+      global.gc()
+    }
+
     const duration = (Date.now() - startTime) / 1000
     console.info(`Transfer to ${this.collectionName} complete. Total rows processed: ${this.processedDocs}. Time: ${duration} seconds`)
   }
 
   async export (query, exportConfig) {
+    console.log(`Starting export to ${exportConfig.collection}...`)
+    this.logMemoryUsage('at start')
+
     // Configure Firestore settings
     this.firestore.settings({
       databaseId: exportConfig.database,
@@ -183,9 +256,19 @@ export class FirestoreBatch {
       date: exportConfig.date
     })
 
-    await this.batchDelete()
+    try {
+      await this.batchDelete()
+      this.logMemoryUsage('after deletion')
 
-    const rowStream = await this.bigquery.queryResultsStream(query)
-    await this.streamFromBigQuery(rowStream)
+      const rowStream = await this.bigquery.queryResultsStream(query)
+      await this.streamFromBigQuery(rowStream)
+
+      this.logMemoryUsage('at completion')
+      console.log(`✅ Export to ${exportConfig.collection} completed successfully`)
+    } catch (error) {
+      this.logMemoryUsage('on error')
+      console.error(`❌ Export to ${exportConfig.collection} failed:`, error)
+      throw error
+    }
   }
 }

infra/tf/bigquery_export/main.tf

@@ -20,8 +20,14 @@ resource "google_cloud_run_v2_job" "bigquery_export" {
   deletion_protection = false
 
   template {
-    parallelism = 0
+    parallelism = 1
+    task_count  = 1  # Ensure single task execution
+
     template {
+      timeout         = "10800s"  # 3 hours
+      service_account = var.function_identity
+      max_retries     = 0   # No retries
+
       containers {
         image = "${var.location}.gcr.io/${var.project}/cloud-run/${var.function_name}:latest"
         resources {
@@ -34,10 +40,21 @@ resource "google_cloud_run_v2_job" "bigquery_export" {
           name  = "EXPORT_CONFIG"
           value = ""
         }
+        env {
+          name  = "NODE_OPTIONS"
+          value = "--expose-gc --max-old-space-size=6144"  # 6GB heap limit with GC
+        }
+
+        env {
+          name  = "MEMORY_WARNING_THRESHOLD_MB"
+          value = "4915"  # 80% of max heap size = 6144MB
+        }
+
+        env {
+          name  = "LOG_LEVEL"
+          value = "info"
+        }
       }
-      timeout         = "7200s"
-      service_account = var.function_identity
-      max_retries     = 1
     }
   }
 }