diff --git a/spec/main-process/file-recovery-service.spec.js b/spec/main-process/file-recovery-service.spec.js index 7b60a9d7d..1cc0062f1 100644 --- a/spec/main-process/file-recovery-service.spec.js +++ b/spec/main-process/file-recovery-service.spec.js @@ -42,7 +42,7 @@ describe("FileRecoveryService", () => { assert.equal(fs.readFileSync(filePath, 'utf8'), "changed") }) - it("creates many recovery files and deletes them when many windows attempt to save the same file", () => { + it("creates only one recovery file when many windows attempt to save the same file, deleting it when the last one finishes saving it", () => { const mockWindow = createWindow() const anotherMockWindow = createWindow() const filePath = temp.path() @@ -50,7 +50,7 @@ describe("FileRecoveryService", () => { fs.writeFileSync(filePath, "some content") recoveryService.willSavePath({sender: mockWindow.webContents}, filePath) recoveryService.willSavePath({sender: anotherMockWindow.webContents}, filePath) - assert.equal(fs.listTreeSync(recoveryDirectory).length, 2) + assert.equal(fs.listTreeSync(recoveryDirectory).length, 1) fs.writeFileSync(filePath, "changed") recoveryService.didSavePath({sender: mockWindow.webContents}, filePath) @@ -78,26 +78,42 @@ describe("FileRecoveryService", () => { assert.equal(fs.readFileSync(filePath, 'utf8'), "some content") }) - it("restores the created recovery files and deletes them in the order in which windows crash when they attempt to save the same file", () => { - const mockWindow = createWindow() - const anotherMockWindow = createWindow() - const filePath = temp.path() + describe("when many windows attempt to save the same file", () => { + it("recovers the file when the window that initiated the save crashes", () => { + const mockWindow = createWindow() + const anotherMockWindow = createWindow() + const filePath = temp.path() - fs.writeFileSync(filePath, "window 1") - recoveryService.willSavePath({sender: mockWindow.webContents}, filePath) - fs.writeFileSync(filePath, "window 2") - recoveryService.willSavePath({sender: anotherMockWindow.webContents}, filePath) - assert.equal(fs.listTreeSync(recoveryDirectory).length, 2) + fs.writeFileSync(filePath, "window 1") + recoveryService.willSavePath({sender: mockWindow.webContents}, filePath) + fs.writeFileSync(filePath, "window 2") + recoveryService.willSavePath({sender: anotherMockWindow.webContents}, filePath) + assert.equal(fs.listTreeSync(recoveryDirectory).length, 1) - fs.writeFileSync(filePath, "changed") + fs.writeFileSync(filePath, "changed") - mockWindow.webContents.emit("crashed") - assert.equal(fs.readFileSync(filePath, 'utf8'), "window 1") - assert.equal(fs.listTreeSync(recoveryDirectory).length, 1) + mockWindow.webContents.emit("crashed") + assert.equal(fs.readFileSync(filePath, 'utf8'), "window 1") + assert.equal(fs.listTreeSync(recoveryDirectory).length, 0) + }) - anotherMockWindow.webContents.emit("crashed") - assert.equal(fs.readFileSync(filePath, 'utf8'), "window 2") - assert.equal(fs.listTreeSync(recoveryDirectory).length, 0) + it("recovers the file when a window that did not initiate the save crashes", () => { + const mockWindow = createWindow() + const anotherMockWindow = createWindow() + const filePath = temp.path() + + fs.writeFileSync(filePath, "window 1") + recoveryService.willSavePath({sender: mockWindow.webContents}, filePath) + fs.writeFileSync(filePath, "window 2") + recoveryService.willSavePath({sender: anotherMockWindow.webContents}, filePath) + assert.equal(fs.listTreeSync(recoveryDirectory).length, 1) + + fs.writeFileSync(filePath, "changed") + + anotherMockWindow.webContents.emit("crashed") + assert.equal(fs.readFileSync(filePath, 'utf8'), "window 1") + assert.equal(fs.listTreeSync(recoveryDirectory).length, 0) + }) }) it("emits a warning when a file can't be recovered", () => { diff --git a/src/main-process/file-recovery-service.js b/src/main-process/file-recovery-service.js index 56cc5e2fe..71ae6200d 100644 --- a/src/main-process/file-recovery-service.js +++ b/src/main-process/file-recovery-service.js @@ -5,10 +5,47 @@ import crypto from 'crypto' import Path from 'path' import fs from 'fs-plus' +class RecoveryFile { + constructor (originalPath, recoveryPath) { + this.originalPath = originalPath + this.recoveryPath = recoveryPath + this.refCount = 0 + } + + storeSync () { + fs.writeFileSync(this.recoveryPath, fs.readFileSync(this.originalPath)) + } + + recoverSync () { + fs.writeFileSync(this.originalPath, fs.readFileSync(this.recoveryPath)) + this.removeSync() + this.refCount = 0 + } + + removeSync () { + fs.unlinkSync(this.recoveryPath) + } + + retain () { + if (this.refCount === 0) this.storeSync() + this.refCount++ + } + + release () { + this.refCount-- + if (this.refCount === 0) this.removeSync() + } + + isReleased () { + return this.refCount === 0 + } +} + export default class FileRecoveryService { constructor (recoveryDirectory) { this.recoveryDirectory = recoveryDirectory - this.recoveryPathsByWindowAndFilePath = new WeakMap() + this.recoveryFilesByFilePath = new Map() + this.recoveryFilesByWindow = new WeakMap() this.observedWindows = new WeakSet() } @@ -25,22 +62,24 @@ export default class FileRecoveryService { } const window = BrowserWindow.fromWebContents(event.sender) - const recoveryFileName = crypto.randomBytes(5).toString('hex') - const recoveryPath = Path.join(this.recoveryDirectory, recoveryFileName) - fs.writeFileSync(recoveryPath, fs.readFileSync(path)) - - if (!this.recoveryPathsByWindowAndFilePath.has(window)) { - this.recoveryPathsByWindowAndFilePath.set(window, new Map()) + let recoveryFile = this.recoveryFilesByFilePath.get(path) + if (recoveryFile == null) { + const recoveryPath = Path.join(this.recoveryDirectory, crypto.randomBytes(5).toString('hex')) + recoveryFile = new RecoveryFile(path, recoveryPath) + this.recoveryFilesByFilePath.set(path, recoveryFile) } - this.recoveryPathsByWindowAndFilePath.get(window).set(path, recoveryPath) + recoveryFile.retain() + + if (!this.recoveryFilesByWindow.has(window)) this.recoveryFilesByWindow.set(window, new Set()) + this.recoveryFilesByWindow.get(window).add(recoveryFile) if (!this.observedWindows.has(window)) { + this.observedWindows.add(window) window.webContents.on("crashed", () => this.recoverFilesForWindow(window)) window.on("closed", () => { this.observedWindows.delete(window) - this.recoveryPathsByWindowAndFilePath.delete(window) + this.recoveryFilesByWindow.delete(window) }) - this.observedWindows.add(window) } event.returnValue = true @@ -48,29 +87,29 @@ export default class FileRecoveryService { didSavePath (event, path) { const window = BrowserWindow.fromWebContents(event.sender) - const recoveryPathsByFilePath = this.recoveryPathsByWindowAndFilePath.get(window) - if (recoveryPathsByFilePath == null || !recoveryPathsByFilePath.has(path)) { - event.returnValue = false - return + const recoveryFile = this.recoveryFilesByFilePath.get(path) + if (recoveryFile != null) { + recoveryFile.release() + if (recoveryFile.isReleased()) this.recoveryFilesByFilePath.delete(path) + this.recoveryFilesByWindow.get(window).delete(recoveryFile) } - const recoveryPath = recoveryPathsByFilePath.get(path) - fs.unlinkSync(recoveryPath) - recoveryPathsByFilePath.delete(path) event.returnValue = true } recoverFilesForWindow (window) { - const recoveryPathsByFilePath = this.recoveryPathsByWindowAndFilePath.get(window) - for (let [filePath, recoveryPath] of recoveryPathsByFilePath) { + if (!this.recoveryFilesByWindow.has(window)) return + + for (const recoveryFile of this.recoveryFilesByWindow.get(window)) { try { - fs.writeFileSync(filePath, fs.readFileSync(recoveryPath)) - fs.unlinkSync(recoveryPath) + recoveryFile.recoverSync() } catch (error) { - console.log(`Cannot recover ${filePath}. A recovery file has been saved here: ${recoveryPath}.`) + console.log(`Cannot recover ${recoveryFile.originalPath}. A recovery file has been saved here: ${recoveryFile.recoveryPath}.`) + } finally { + this.recoveryFilesByFilePath.delete(recoveryFile.originalPath) } } - recoveryPathsByFilePath.clear() + this.recoveryFilesByWindow.delete(window) } }