fix: upgrade idb

Author: sheepzh

package.json

@@ -39,7 +39,7 @@
         "@types/chrome": "0.1.37",
         "@types/decompress": "^4.2.7",
         "@types/jest": "^30.0.0",
-        "@types/node": "^25.3.0",
+        "@types/node": "^25.3.1",
         "@types/punycode": "^2.1.4",
         "@vue/babel-plugin-jsx": "^2.0.1",
         "babel-loader": "^10.0.0",

src/background/index.ts

@@ -15,7 +15,7 @@ import badgeTextManager from "./badge-manager"
 import initBrowserAction from "./browser-action-manager"
 import initCsHandler from "./content-script-handler"
 import initDataCleaner from "./data-cleaner"
-import handleInstall from "./install-handler"
+import handleInstall from './install-handler'
 import initLimitProcessor from "./limit-processor"
 import MessageDispatcher from "./message-dispatcher"
 import VersionMigrator from "./migrator"

src/background/migrator/index.ts

@@ -42,7 +42,7 @@ class VersionManager {
         }
     }
 
-    init() {
+    async init() {
         onInstalled(reason => this.onChromeInstalled(reason))
     }
 }

src/background/migrator/indexed-migrator.ts

@@ -1,10 +1,28 @@
+import { BaseIDBStorage } from '@db/common/indexed-storage'
+import { IDBStatDatabase } from '@db/stat-database/idb'
 import timelineDatabase from '@db/timeline-database'
+import IDBTimelineDatabase from '@db/timeline-database/idb'
 import type { Migrator } from './types'
 
+async function upgradeIndexedDB() {
+    try {
+        const storages: BaseIDBStorage<any>[] = [new IDBStatDatabase(), new IDBTimelineDatabase()]
+        for (const storage of storages) {
+            await storage.upgrade()
+        }
+        console.log('IndexedDB upgraded successfully')
+    } catch (error) {
+        console.error('Failed to upgrade IndexedDB', error)
+    }
+}
+
 class IndexedMigrator implements Migrator {
     onInstall(): void {
     }
-    onUpdate(_version: string): void {
+
+    async onUpdate(_version: string): Promise<void> {
+        await upgradeIndexedDB()
+
         timelineDatabase.migrateFromClassic()
             .then(() => console.log('Timeline data migrated to IndexedDB'))
             .catch(e => console.error('Failed to migrate timeline data to IndexedDB', e))

src/database/common/indexed-storage.ts

@@ -87,53 +87,131 @@ export abstract class BaseIDBStorage<T = Record<string, unknown>> {
     private DB_NAME = `tt4b_${chrome.runtime.id}` as const
 
     private db: IDBDatabase | undefined
+    private static initPromises = new Map<string, Promise<IDBDatabase>>()
+
     abstract indexes: Index<T>[]
     abstract key: Key<T> | Key<T>[]
     abstract table: Table
 
     protected async initDb(): Promise<IDBDatabase> {
         if (this.db) return this.db
 
+        let initPromise = BaseIDBStorage.initPromises.get(this.table)
+        if (!initPromise) {
+            initPromise = this.doInitDb()
+            BaseIDBStorage.initPromises.set(this.table, initPromise)
+        }
+
+        try {
+            this.db = await initPromise
+            this.setupDbCloseHandler(this.db)
+            return this.db
+        } catch (error) {
+            BaseIDBStorage.initPromises.delete(this.table)
+            throw error
+        }
+    }
+
+    private setupDbCloseHandler(db: IDBDatabase): void {
+        db.onclose = () => {
+            if (this.db !== db) return
+
+            this.db = undefined
+            BaseIDBStorage.initPromises.delete(this.table)
+        }
+    }
+
+    private async doInitDb(): Promise<IDBDatabase> {
         const factory = typeof window !== 'undefined' ? window.indexedDB : globalThis.indexedDB
-        const checkRequest = factory.open(this.DB_NAME)
 
-        return new Promise((resolve, reject) => {
-            checkRequest.onsuccess = () => {
-                const db = checkRequest.result
-                const storeExisted = db.objectStoreNames.contains(this.table)
-                const needUpgrade = !storeExisted || this.needUpgradeIndexes(db)
+        const checkDb = await new Promise<IDBDatabase>((resolve, reject) => {
+            const checkRequest = factory.open(this.DB_NAME)
+            checkRequest.onsuccess = () => resolve(checkRequest.result)
+            checkRequest.onerror = () => reject(checkRequest.error || new Error("Failed to open database"))
+        })
 
-                if (!needUpgrade) {
-                    this.db = db
-                    return resolve(db)
-                }
+        return checkDb
+    }
 
-                const currentVersion = db.version
-                db.close()
+    // Only used for testing, be careful when using in production
+    public async clear(): Promise<void> {
+        await this.withStore(store => store.clear(), 'readwrite')
+    }
 
-                const upgradeRequest = factory.open(this.DB_NAME, currentVersion + 1)
+    async upgrade(): Promise<void> {
+        const factory = typeof window !== 'undefined' ? window.indexedDB : globalThis.indexedDB
+
+        const checkDb = await new Promise<IDBDatabase>((resolve, reject) => {
+            const checkRequest = factory.open(this.DB_NAME)
+            checkRequest.onsuccess = () => resolve(checkRequest.result)
+            checkRequest.onerror = () => reject(checkRequest.error || new Error("Failed to open database"))
+            checkRequest.onblocked = () => {
+                console.warn(`Database check blocked for "${this.table}" (DB: ${this.DB_NAME}), waiting for other connections to close`)
+            }
+        })
 
-                upgradeRequest.onupgradeneeded = () => {
+        const storeExisted = checkDb.objectStoreNames.contains(this.table)
+        const needUpgrade = !storeExisted || this.needUpgradeIndexes(checkDb)
+
+        if (!needUpgrade) {
+            checkDb.close()
+            return
+        }
+
+        const currentVersion = checkDb.version
+        checkDb.close()
+
+        return new Promise<void>((resolve, reject) => {
+            const upgradeRequest = factory.open(this.DB_NAME, currentVersion + 1)
+
+            upgradeRequest.onupgradeneeded = () => {
+                try {
                     const upgradeDb = upgradeRequest.result
                     const transaction = upgradeRequest.transaction
-                    if (!transaction) return reject("Failed to get transaction of upgrading request")
+                    if (!transaction) {
+                        reject(new Error("Failed to get transaction of upgrading request"))
+                        return
+                    }
+
+                    transaction.onerror = () => {
+                        reject(transaction.error || new Error("Transaction failed"))
+                    }
+
+                    transaction.onabort = () => {
+                        reject(new Error("Upgrade transaction was aborted"))
+                    }
 
                     let store = upgradeDb.objectStoreNames.contains(this.table)
                         ? transaction.objectStore(this.table)
-                        : upgradeDb.createObjectStore(this.table, { keyPath: this.key })
+                        : upgradeDb.createObjectStore(this.table, { keyPath: this.key as string | string[] })
                     this.createIndexes(store)
+                } catch (error) {
+                    console.error("Failed to upgrade database in onupgradeneeded", error)
+                    upgradeRequest.transaction?.abort()
+                    reject(error instanceof Error ? error : new Error(String(error)))
                 }
+            }
 
-                upgradeRequest.onsuccess = () => {
-                    console.log("IndexedDB upgraded")
-                    this.db = upgradeRequest.result
-                    resolve(upgradeRequest.result)
-                }
+            upgradeRequest.onsuccess = () => {
+                console.log(`IndexedDB upgraded for table "${this.table}"`)
+                upgradeRequest.result.close()
+                resolve()
+            }
 
-                upgradeRequest.onerror = () => reject(upgradeRequest.error)
+            upgradeRequest.onerror = (event) => {
+                console.error("Failed to upgrade database", event, upgradeRequest.error)
+                reject(upgradeRequest.error || new Error("Failed to upgrade database"))
             }
 
-            checkRequest.onerror = () => reject(checkRequest.error)
+            upgradeRequest.onblocked = () => {
+                const blockingTables = Array.from(BaseIDBStorage.initPromises.keys())
+                    .filter(table => table !== this.table)
+                console.warn(
+                    `Database upgrade blocked for table "${this.table}" (DB: ${this.DB_NAME}), ` +
+                    `waiting for other connections to close. ` +
+                    `Other tables with active connections: ${blockingTables.length > 0 ? blockingTables.join(', ') : 'none'}`
+                )
+            }
         })
     }
 
@@ -170,27 +248,41 @@ export abstract class BaseIDBStorage<T = Record<string, unknown>> {
     }
 
     protected async withStore<T = unknown>(operation: (store: IDBObjectStore) => T | Promise<T>, mode?: IDBTransactionMode): Promise<T> {
-        const db = await this.initDb()
-        const trans = db.transaction(this.table, mode ?? 'readwrite')
-        try {
-            const store = trans.objectStore(this.table)
-            const result = await operation(store)
-            // Waiting for transaction completed
-            await new Promise<void>((resolve, reject) => {
-                trans.oncomplete = () => resolve()
-                trans.onerror = () => reject(trans.error)
-                trans.onabort = () => reject(new Error('Transaction aborted'))
-            })
-            return result
-        } catch (e) {
-            console.error("Failed to process with transaction", e)
-            if (!trans.error && trans.mode !== 'readonly') {
-                try {
-                    trans.abort()
-                } catch (ignored) { }
+        let db = await this.initDb()
+
+        for (let retryCount = 0; retryCount < 2; retryCount++) {
+            let trans: IDBTransaction | undefined
+            try {
+                trans = db.transaction(this.table, mode ?? 'readwrite')
+                const store = trans.objectStore(this.table)
+                const result = await operation(store)
+                const transaction = trans
+                await new Promise<void>((resolve, reject) => {
+                    transaction.oncomplete = () => resolve()
+                    transaction.onerror = () => reject(transaction.error)
+                    transaction.onabort = () => reject(new Error('Transaction aborted'))
+                })
+                return result
+            } catch (e) {
+                const isConnectionError = e instanceof DOMException &&
+                    (e.name === 'InvalidStateError' || e.name === 'AbortError')
+
+                if (isConnectionError) {
+                    this.db = undefined
+                    BaseIDBStorage.initPromises.delete(this.table)
+                    db = await this.initDb()
+                } else {
+                    console.error("Failed to process with transaction", e)
+                    if (trans && !trans.error && trans.mode !== 'readonly') {
+                        try {
+                            trans.abort()
+                        } catch (ignored) { }
+                    }
+                    throw e
+                }
             }
-            throw e
         }
+        throw new Error("Max retries exceeded")
     }
 
     protected assertIndex(store: IDBObjectStore, key: Key<T> | Key<T>[]): IDBIndex {

test/database/idb.ts

@@ -2,23 +2,25 @@ import { zeroResult, zeroRow } from '@db/stat-database/common'
 import { IDBStatDatabase } from '@db/stat-database/idb'
 import 'fake-indexeddb/auto'
 import { mockRuntime } from '../../__mock__/runtime'
-import { cleanupIDB } from '../idb'
 
-let db: IDBStatDatabase
 const GOOGLE = 'www.google.com'
 const GITHUB = 'www.github.com'
 const GITHUB_VIRTUAL = 'www.github.com/sheepzh/**'
 const GROUP_1 = 1
 const GROUP_2 = 2
 const MAYBE_GROUP_1 = '1'
 
+let db: IDBStatDatabase
+
 describe('stat-database/idb', () => {
-    beforeAll(() => mockRuntime())
-    beforeEach(async () => {
-        await cleanupIDB()
+    beforeAll(async () => {
+        mockRuntime()
         db = new IDBStatDatabase()
+        await db.upgrade()
     })
 
+    beforeEach(() => db.clear())
+
     test('accumulate', async () => {
         await db.accumulate(GITHUB, '20240601', { focus: 10, time: 20 })
         await db.accumulate(GOOGLE, new Date(2025, 10, 1), { focus: 1, time: 0 })