Add preupgrade hook

index.js

@@ -12,6 +12,7 @@ const codecs = require('codecs')
 const Replicator = require('./lib/replicator')
 const Core = require('./lib/core')
 const BlockEncryption = require('./lib/block-encryption')
+const Upgrader = require('./lib/upgrader')
 const { ReadStream, WriteStream } = require('./lib/streams')
 const { BAD_ARGUMENT, SESSION_CLOSED, SESSION_NOT_WRITABLE, SNAPSHOT_NOT_AVAILABLE } = require('./lib/errors')
 
@@ -78,6 +79,8 @@ module.exports = class Hypercore extends EventEmitter {
     this._preappend = preappend.bind(this)
     this._snapshot = null
     this._findingPeers = 0
+
+    this._upgrader = new Upgrader(this, opts)
   }
 
   [inspect] (depth, opts) {
@@ -519,6 +522,8 @@ module.exports = class Hypercore extends EventEmitter {
 
           // If snapshotted, make sure to update our compat so we can fail gets
           if (s._snapshot && bitfield.start < s._snapshot.compatLength) s._snapshot.compatLength = bitfield.start
+
+          s._upgrader.ontruncate()
         }
 
         if (s.sparse ? truncated : truncatedNonSparse) {
@@ -613,22 +618,7 @@ module.exports = class Hypercore extends EventEmitter {
     const activeRequests = (opts && opts.activeRequests) || this.activeRequests
     const req = this.replicator.addUpgrade(activeRequests)
 
-    let upgraded = await req.promise
-
-    if (!this.sparse) {
-      // Download all available blocks in non-sparse mode
-      const start = this.length
-      const end = this.core.tree.length
-      const contig = this.contiguousLength
-
-      await this.download({ start, end, ifAvailable: true }).downloaded()
-
-      if (!upgraded) upgraded = this.contiguousLength !== contig
-    }
-
-    if (!upgraded) return false
-    if (this.snapshotted) return this._updateSnapshot()
-    return true
+    return this._upgrader.onupgrade(req.promise)
   }
 
   async seek (bytes, opts) {

lib/upgrader.js

@@ -0,0 +1,97 @@
+module.exports = class Upgrader {
+  constructor (session, hooks = {}) {
+    this.session = session
+    this.active = null
+    this.preupgrade = hooks.preupgrade || null
+  }
+
+  cancel () {
+    if (this.active) {
+      this.active.cancel()
+      this.active = null
+    }
+  }
+
+  onupgrade (request) {
+    if (this.active) this.active.cancel()
+
+    const upgrade = this.active = new Upgrade(this)
+    upgrade.run(request)
+
+    return this.active.upgraded
+  }
+
+  ontruncate () {
+    // TODO: handle truncation
+  }
+}
+
+class Upgrade {
+  constructor (upgrader) {
+    this.upgrader = upgrader
+    this.active = true
+    this.upgraded = new Promise((resolve) => {
+      this.done = resolve
+    })
+  }
+
+  cancel () {
+    this.active = false
+  }
+
+  async run (request) {
+    if (!this.active) return this.done(false)
+
+    const session = this.upgrader.session
+    const preupgradeLength = session.length
+
+    if (this.upgrader.preupgrade) session._updateSnapshot()
+
+    let upgraded = await request
+
+    if (!this.active) return this.done(false)
+
+    if (this.upgrader.preupgrade) {
+      const latest = session.session()
+
+      let length
+      try {
+        length = await this.upgrader.preupgrade(latest)
+      } catch {
+        length = session.length
+      } finally {
+        await latest.close()
+      }
+
+      if (!this.active) return this.done(false)
+
+      if (typeof length === 'number' && length >= preupgradeLength && length < session.core.tree.length) {
+        session._snapshot = {
+          length,
+          byteLength: 0,
+          fork: session.core.tree.fork,
+          compatLength: length
+        }
+
+        return this.done(length !== preupgradeLength)
+      }
+    }
+
+    if (!session.sparse) {
+      // Download all available blocks in non-sparse mode
+      const start = session.length
+      const end = session.core.tree.length
+      const contig = session.contiguousLength
+
+      await session.download({ start, end, ifAvailable: true }).downloaded()
+
+      if (!upgraded) upgraded = session.contiguousLength !== contig
+    }
+
+    if (!upgraded) return this.done(false)
+
+    if (this.upgrader.preupgrade || session.snapshotted) return this.done(session._updateSnapshot())
+
+    return this.done(true)
+  }
+}

test/preupgrade.js

@@ -0,0 +1,80 @@
+const test = require('brittle')
+const { create, replicate } = require('./helpers')
+
+test('upgrade to latest length', async function (t) {
+  const a = await create()
+  await a.append(['a', 'b', 'c'])
+
+  const b = await create(a.key, {
+    preupgrade (latest) {
+      t.is(b.length, 0)
+      t.is(latest.length, 3)
+
+      return latest.length
+    }
+  })
+
+  replicate(a, b, t)
+  t.is(b.length, 0)
+
+  t.ok(await b.update())
+  t.is(b.length, 3)
+})
+
+test('stay on previous length', async function (t) {
+  const a = await create()
+  await a.append(['a', 'b', 'c'])
+
+  const b = await create(a.key, {
+    preupgrade (latest) {
+      t.is(b.length, 0)
+      t.is(latest.length, 3)
+
+      return 0
+    }
+  })
+
+  replicate(a, b, t)
+  t.is(b.length, 0)
+
+  t.absent(await b.update())
+  t.is(b.length, 0)
+})
+
+test('return no length', async function (t) {
+  const a = await create()
+  await a.append(['a', 'b', 'c'])
+
+  const b = await create(a.key, {
+    preupgrade (latest) {
+      t.is(b.length, 0)
+      t.is(latest.length, 3)
+    }
+  })
+
+  replicate(a, b, t)
+  t.is(b.length, 0)
+
+  t.ok(await b.update())
+  t.is(b.length, 3)
+})
+
+test('return invalid length', async function (t) {
+  const a = await create()
+  await a.append(['a', 'b', 'c'])
+
+  const b = await create(a.key, {
+    preupgrade (latest) {
+      t.is(b.length, 0)
+      t.is(latest.length, 3)
+
+      return -1
+    }
+  })
+
+  replicate(a, b, t)
+  t.is(b.length, 0)
+
+  t.ok(await b.update())
+  t.is(b.length, 3)
+})