Handle recovery when many windows save the same file simultaneously

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", () => {

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)
   }
 }