brainy/tests/unit/storage/hnswConcurrency.test.ts

/**
 * Unit Tests for HNSW Concurrency Bug Fix (v4.10.1)
 *
 * Tests atomic write strategies across storage adapters to prevent race conditions
 * during concurrent HNSW neighbor updates.
 *
 * Test Coverage:
 * 1. MemoryStorage - Mutex locking
 * 2. FileSystemStorage - Atomic rename
 * 3. Concurrent saveHNSWData() calls on same entity
 * 4. Data integrity verification (no lost connections)
 *
 * NO MOCKS - All tests use real storage operations and verify actual behavior
 */

import { describe, it, expect, beforeEach, afterEach } from 'vitest'
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'
import { FileSystemStorage } from '../../../src/storage/adapters/fileSystemStorage.js'
import * as fs from 'fs/promises'
import * as path from 'path'
import * as os from 'os'

const TEST_ROOT = path.join(os.tmpdir(), 'brainy-hnsw-concurrency-tests')

describe('HNSW Concurrency Bug Fix (v4.10.1)', () => {
  describe('MemoryStorage - Mutex Locking', () => {
    let storage: MemoryStorage

    beforeEach(async () => {
      storage = new MemoryStorage()
      await storage.init()
    })

    it('should serialize concurrent saveHNSWData() calls on same entity', async () => {
      console.log('🧪 Testing MemoryStorage mutex locking...')

      const nounId = '00000000-0000-0000-0000-000000000001'

      // Simulate 20 concurrent neighbor connections (like bulk import)
      const concurrentUpdates = []
      for (let i = 0; i < 20; i++) {
        const connections: Record<string, string[]> = {
          '0': [`neighbor-${i}`]
        }

        concurrentUpdates.push(
          storage.saveHNSWData(nounId, {
            level: 0,
            connections
          })
        )
      }

      // Execute all updates concurrently
      await Promise.all(concurrentUpdates)

      // Verify final state - should have the last update's data
      const finalData = await storage.getHNSWData(nounId)
      expect(finalData).toBeDefined()
      expect(finalData!.level).toBe(0)
      expect(finalData!.connections['0']).toBeDefined()

      // Due to mutex serialization, last writer wins
      // Important: verify NO crash and data is consistent
      expect(finalData!.connections['0'].length).toBeGreaterThan(0)

      console.log('✅ MemoryStorage mutex prevented race condition')
    })

    it('should preserve existing data when updating HNSW connections', async () => {
      console.log('🧪 Testing MemoryStorage data preservation...')

      const nounId = '00000000-0000-0000-0000-000000000002'

      // Initial state: 5 connections at level 0
      await storage.saveHNSWData(nounId, {
        level: 0,
        connections: {
          '0': ['conn-1', 'conn-2', 'conn-3', 'conn-4', 'conn-5']
        }
      })

      // Update: Add connection at level 1 (should preserve level 0)
      await storage.saveHNSWData(nounId, {
        level: 1,
        connections: {
          '0': ['conn-1', 'conn-2', 'conn-3', 'conn-4', 'conn-5'],
          '1': ['conn-6']
        }
      })

      const finalData = await storage.getHNSWData(nounId)
      expect(finalData!.level).toBe(1)
      expect(finalData!.connections['0']).toHaveLength(5)
      expect(finalData!.connections['1']).toHaveLength(1)

      console.log('✅ MemoryStorage preserved existing connections')
    })

    it('should handle saveHNSWSystem() concurrent calls', async () => {
      console.log('🧪 Testing MemoryStorage saveHNSWSystem() mutex...')

      // Simulate concurrent system updates (entry point changes)
      const concurrentUpdates = []
      for (let i = 0; i < 10; i++) {
        concurrentUpdates.push(
          storage.saveHNSWSystem({
            entryPointId: `entry-${i}`,
            maxLevel: i
          })
        )
      }

      await Promise.all(concurrentUpdates)

      const systemData = await storage.getHNSWSystem()
      expect(systemData).toBeDefined()
      expect(systemData!.entryPointId).toBeDefined()
      expect(systemData!.maxLevel).toBeGreaterThanOrEqual(0)

      console.log('✅ MemoryStorage saveHNSWSystem() mutex working')
    })
  })

  describe('FileSystemStorage - Atomic Rename', () => {
    let storage: FileSystemStorage
    let testDir: string

    beforeEach(async () => {
      testDir = path.join(TEST_ROOT, `test-${Date.now()}-${Math.random().toString(36).substring(2)}`)
      await fs.mkdir(testDir, { recursive: true })

      storage = new FileSystemStorage(testDir)
      await storage.init()
    })

    afterEach(async () => {
      try {
        await fs.rm(testDir, { recursive: true, force: true })
      } catch (error) {
        // Ignore cleanup errors
      }
    })

    it('should use atomic rename for concurrent saveHNSWData() calls', async () => {
      console.log('🧪 Testing FileSystemStorage atomic rename...')

      const nounId = 'ab000000-0000-0000-0000-000000000001'

      // Create initial noun data (so file exists)
      await storage.saveHNSWData(nounId, {
        level: 0,
        connections: { '0': ['initial'] }
      })

      // Simulate 20 concurrent updates
      const concurrentUpdates = []
      for (let i = 0; i < 20; i++) {
        const connections: Record<string, string[]> = {
          '0': [`neighbor-${i}`]
        }

        concurrentUpdates.push(
          storage.saveHNSWData(nounId, {
            level: 0,
            connections
          })
        )
      }

      // Execute all updates concurrently
      await Promise.all(concurrentUpdates)

      // Verify final state
      const finalData = await storage.getHNSWData(nounId)
      expect(finalData).toBeDefined()
      expect(finalData!.level).toBe(0)
      expect(finalData!.connections['0']).toBeDefined()
      expect(finalData!.connections['0'].length).toBeGreaterThan(0)

      // Verify no temp files left behind
      const nounsDir = path.join(testDir, 'entities', 'nouns', 'hnsw', 'ab')
      try {
        const files = await fs.readdir(nounsDir)
        const tempFiles = files.filter(f => f.includes('.tmp.'))
        expect(tempFiles).toHaveLength(0)
      } catch (error: any) {
        // Directory doesn't exist = no temp files leaked (good!)
        if (error.code !== 'ENOENT') throw error
      }

      console.log('✅ FileSystemStorage atomic rename working, no temp files leaked')
    })

    it('should preserve existing node data during HNSW updates', async () => {
      console.log('🧪 Testing FileSystemStorage data preservation...')

      const nounId = 'cd000000-0000-0000-0000-000000000002'

      // Manually create a noun file with id and vector (simulating real entity)
      const shardDir = path.join(testDir, 'entities', 'nouns', 'hnsw', 'cd')
      await fs.mkdir(shardDir, { recursive: true })

      const initialNode = {
        id: nounId,
        vector: [0.1, 0.2, 0.3, 0.4],
        someOtherField: 'should-be-preserved'
      }

      await fs.writeFile(
        path.join(shardDir, `${nounId}.json`),
        JSON.stringify(initialNode, null, 2)
      )

      // Update HNSW data
      await storage.saveHNSWData(nounId, {
        level: 2,
        connections: {
          '0': ['conn-1', 'conn-2'],
          '1': ['conn-3'],
          '2': ['conn-4']
        }
      })

      // Read file and verify id and vector are preserved
      const updatedContent = await fs.readFile(
        path.join(shardDir, `${nounId}.json`),
        'utf-8'
      )
      const updatedNode = JSON.parse(updatedContent)

      expect(updatedNode.id).toBe(nounId)
      expect(updatedNode.vector).toEqual([0.1, 0.2, 0.3, 0.4])
      expect(updatedNode.someOtherField).toBe('should-be-preserved')
      expect(updatedNode.level).toBe(2)
      expect(updatedNode.connections['0']).toHaveLength(2)
      expect(updatedNode.connections['1']).toHaveLength(1)
      expect(updatedNode.connections['2']).toHaveLength(1)

      console.log('✅ FileSystemStorage preserved existing node data')
    })

    it('should handle saveHNSWSystem() with atomic rename', async () => {
      console.log('🧪 Testing FileSystemStorage saveHNSWSystem() atomic rename...')

      // Concurrent system updates
      const concurrentUpdates = []
      for (let i = 0; i < 10; i++) {
        concurrentUpdates.push(
          storage.saveHNSWSystem({
            entryPointId: `entry-${i}`,
            maxLevel: i
          })
        )
      }

      await Promise.all(concurrentUpdates)

      const systemData = await storage.getHNSWSystem()
      expect(systemData).toBeDefined()
      expect(systemData!.entryPointId).toBeDefined()

      // Verify no temp files left
      const systemDir = path.join(testDir, 'system')
      try {
        const files = await fs.readdir(systemDir)
        const tempFiles = files.filter(f => f.includes('.tmp.'))
        expect(tempFiles).toHaveLength(0)
      } catch (error: any) {
        // Directory doesn't exist = no temp files leaked (good!)
        if (error.code !== 'ENOENT') throw error
      }

      console.log('✅ FileSystemStorage saveHNSWSystem() atomic rename working')
    })

    it('should clean up temp files on error', async () => {
      console.log('🧪 Testing FileSystemStorage temp file cleanup on error...')

      const nounId = 'ef000000-0000-0000-0000-000000000003'

      // This should succeed normally
      await storage.saveHNSWData(nounId, {
        level: 0,
        connections: { '0': ['test'] }
      })

      // Verify temp files are cleaned up
      const shardDir = path.join(testDir, 'entities', 'nouns', 'hnsw', 'ef')
      try {
        const files = await fs.readdir(shardDir)
        const tempFiles = files.filter(f => f.includes('.tmp.'))
        expect(tempFiles).toHaveLength(0)
      } catch (error: any) {
        // Directory doesn't exist = no temp files leaked (good!)
        if (error.code !== 'ENOENT') throw error
      }

      console.log('✅ FileSystemStorage cleans up temp files')
    })
  })

  describe('Cross-Adapter Consistency', () => {
    it('should produce same results across MemoryStorage and FileSystemStorage', async () => {
      console.log('🧪 Testing cross-adapter consistency...')

      const nounId = '12000000-0000-0000-0000-000000000001'
      const hnswData = {
        level: 3,
        connections: {
          '0': ['n1', 'n2', 'n3'],
          '1': ['n4', 'n5'],
          '2': ['n6'],
          '3': ['n7']
        }
      }

      // Test MemoryStorage
      const memStorage = new MemoryStorage()
      await memStorage.init()
      await memStorage.saveHNSWData(nounId, hnswData)
      const memResult = await memStorage.getHNSWData(nounId)

      // Test FileSystemStorage
      const testDir = path.join(TEST_ROOT, `cross-test-${Date.now()}`)
      await fs.mkdir(testDir, { recursive: true })

      try {
        const fsStorage = new FileSystemStorage(testDir)
        await fsStorage.init()
        await fsStorage.saveHNSWData(nounId, hnswData)
        const fsResult = await fsStorage.getHNSWData(nounId)

        // Verify both produce same results
        expect(memResult).toEqual(fsResult)
        expect(memResult!.level).toBe(3)
        expect(Object.keys(memResult!.connections)).toHaveLength(4)

        console.log('✅ Both storage adapters produce consistent results')
      } finally {
        await fs.rm(testDir, { recursive: true, force: true })
      }
    })
  })

  // Cleanup test root after all tests
  afterEach(async () => {
    try {
      const exists = await fs.access(TEST_ROOT).then(() => true).catch(() => false)
      if (exists) {
        await fs.rm(TEST_ROOT, { recursive: true, force: true })
      }
    } catch (error) {
      // Ignore cleanup errors
    }
  })
})