From 48ce7df43a559e8331ed4616b437a4c3afb7dbfc Mon Sep 17 00:00:00 2001
From: Julien Fontanet <julien.fontanet@isonoe.net>
Date: Thu, 7 May 2020 17:36:37 +0200
Subject: [PATCH] feat(xo-server/file restore): dedupe mount/unmount (#4961)

---
 packages/xo-server/src/_dedupeUnmount.js      | 45 +++++++++
 .../src/xo-mixins/file-restore-ng.js          | 94 +++++++++++--------
 2 files changed, 99 insertions(+), 40 deletions(-)
 create mode 100644 packages/xo-server/src/_dedupeUnmount.js

diff --git a/packages/xo-server/src/_dedupeUnmount.js b/packages/xo-server/src/_dedupeUnmount.js
new file mode 100644
index 000000000..b83cee4e0
--- /dev/null
+++ b/packages/xo-server/src/_dedupeUnmount.js
@@ -0,0 +1,45 @@
+import assert from 'assert'
+
+import ensureArray from './_ensureArray'
+import MultiKeyMap from './_MultiKeyMap'
+
+function State() {
+  this.i = 0
+  this.value = undefined
+}
+
+export const dedupeUnmount = (fn, keyFn) => {
+  const states = new MultiKeyMap()
+
+  return function() {
+    const keys = ensureArray(keyFn.apply(this, arguments))
+    let state = states.get(keys)
+    if (state === undefined) {
+      state = new State()
+      states.set(keys, state)
+
+      const mount = async () => {
+        try {
+          const value = await fn.apply(this, arguments)
+          return {
+            __proto__: value,
+            async unmount() {
+              assert(state.i > 0)
+              if (--state.i === 0) {
+                states.delete(keys)
+                await value.unmount()
+              }
+            },
+          }
+        } catch (error) {
+          states.delete(keys)
+          throw error
+        }
+      }
+
+      state.value = mount()
+    }
+    ++state.i
+    return state.value
+  }
+}
diff --git a/packages/xo-server/src/xo-mixins/file-restore-ng.js b/packages/xo-server/src/xo-mixins/file-restore-ng.js
index 26f45bbf3..dab5517c4 100644
--- a/packages/xo-server/src/xo-mixins/file-restore-ng.js
+++ b/packages/xo-server/src/xo-mixins/file-restore-ng.js
@@ -6,9 +6,15 @@ import { normalize } from 'path'
 import { readdir, rmdir, stat } from 'fs-extra'
 import { ZipFile } from 'yazl'
 
+import { decorateWith } from '../_decorateWith'
+import { dedupeUnmount } from '../_dedupeUnmount'
 import { lvs, pvs } from '../lvm'
 import { resolveSubpath, tmpDir } from '../utils'
 
+const compose = (...fns) => value => fns.reduce((value, fn) => fn(value), value)
+
+const dedupeUnmountWithArgs = fn => dedupeUnmount(fn, (...args) => args)
+
 const IGNORED_PARTITION_TYPES = {
   // https://github.com/jhermsmeier/node-mbr/blob/master/lib/partition.js#L38
   0x05: true,
@@ -60,50 +66,56 @@ const parsePartxLine = createPairsParser({
       : value,
 })
 
-const listLvmLogicalVolumes = defer(
-  async ($defer, devicePath, partition, results = []) => {
-    const pv = await mountLvmPhysicalVolume(devicePath, partition)
-    $defer(pv.unmount)
+const listLvmLogicalVolumes = compose(
+  defer,
+  dedupeUnmountWithArgs
+)(async ($defer, devicePath, partition, results = []) => {
+  const pv = await mountLvmPhysicalVolume(devicePath, partition)
+  $defer(pv.unmount)
 
-    const lvs = await pvs(['lv_name', 'lv_path', 'lv_size', 'vg_name'], pv.path)
-    const partitionId = partition !== undefined ? partition.id : ''
-    lvs.forEach((lv, i) => {
-      const name = lv.lv_name
-      if (name !== '') {
-        results.push({
-          id: `${partitionId}/${lv.vg_name}/${name}`,
-          name,
-          size: lv.lv_size,
-        })
-      }
-    })
-    return results
+  const lvs = await pvs(['lv_name', 'lv_path', 'lv_size', 'vg_name'], pv.path)
+  const partitionId = partition !== undefined ? partition.id : ''
+  lvs.forEach((lv, i) => {
+    const name = lv.lv_name
+    if (name !== '') {
+      results.push({
+        id: `${partitionId}/${lv.vg_name}/${name}`,
+        name,
+        size: lv.lv_size,
+      })
+    }
+  })
+  return results
+})
+
+const mountLvmPhysicalVolume = dedupeUnmountWithArgs(
+  async (devicePath, partition) => {
+    const args = []
+    if (partition !== undefined) {
+      args.push('-o', partition.start * 512)
+    }
+    args.push('--show', '-f', devicePath)
+    const path = (await execa('losetup', args)).stdout.trim()
+    await execa('pvscan', ['--cache', path])
+
+    return {
+      path,
+      unmount: async () => {
+        try {
+          const vgNames = await pvs('vg_name', path)
+          await execa('vgchange', ['-an', ...vgNames])
+        } finally {
+          await execa('losetup', ['-d', path])
+        }
+      },
+    }
   }
 )
 
-async function mountLvmPhysicalVolume(devicePath, partition) {
-  const args = []
-  if (partition !== undefined) {
-    args.push('-o', partition.start * 512)
-  }
-  args.push('--show', '-f', devicePath)
-  const path = (await execa('losetup', args)).stdout.trim()
-  await execa('pvscan', ['--cache', path])
-
-  return {
-    path,
-    unmount: async () => {
-      try {
-        const vgNames = await pvs('vg_name', path)
-        await execa('vgchange', ['-an', ...vgNames])
-      } finally {
-        await execa('losetup', ['-d', path])
-      }
-    },
-  }
-}
-
-const mountPartition = defer(async ($defer, devicePath, partition) => {
+const mountPartition = compose(
+  defer,
+  dedupeUnmountWithArgs
+)(async ($defer, devicePath, partition) => {
   const options = ['loop', 'ro']
 
   if (partition !== undefined) {
@@ -280,6 +292,7 @@ export default class BackupNgFileRestore {
     return partitions
   }
 
+  @decorateWith(dedupeUnmountWithArgs)
   @defer
   async _mountDisk($defer, remoteId, diskId) {
     const handler = await this._app.getRemoteHandler(remoteId)
@@ -321,6 +334,7 @@ export default class BackupNgFileRestore {
     }
   }
 
+  @decorateWith(dedupeUnmountWithArgs)
   @defer
   async _mountPartition($defer, devicePath, partitionId) {
     if (partitionId === undefined) {