Support cas to return nodes

README.md

@@ -116,14 +116,16 @@ If you're inserting a series of data atomically or want more performance then ch
 `options` includes:
 ```js
 {
-  cas (prev, next) { return true }
+  cas (prev, next) { return next }
 }
 ```
 
 ##### Compare And Swap (cas)
 `cas` option is a function comparator to control whether the `put` succeeds.
 
-By returning `true` it will insert the value, otherwise it won't.
+By returning `next` it will insert the value, otherwise it won't.
+
+It's allowed to change `next.value` only.
 
 It receives two args: `prev` is the current node entry, and `next` is the potential new node.
 
@@ -140,7 +142,7 @@ console.log(await db.get('number')) // => { seq: 2, key: 'number', value: '456'
 
 function cas (prev, next) {
   // You can use same-data or same-object lib, depending on the value complexity
-  return prev.value !== next.value
+  return prev.value !== next.value ? next : prev
 }
 ```
 
@@ -157,7 +159,7 @@ Delete a key.
 `options` include:
 ```js
 {
-  cas (prev, next) { return true }
+  cas (prev, next) { return next }
 }
 ```
 
@@ -180,7 +182,7 @@ await db.del('number', { cas })
 console.log(await db.get('number')) // => null
 
 function cas (prev) {
-  return prev.value === 'can-be-deleted'
+  return prev.value === 'can-be-deleted' ? next : prev
 }
 ```
 

index.js

@@ -109,7 +109,12 @@ class TreeNode {
       c = b4a.compare(key.value, await this.getKey(mid))
 
       if (c === 0) {
-        if (cas && !(await cas((await this.getKeyNode(mid)).final(encoding), node))) return true
+        if (cas) {
+          const older = (await this.getKeyNode(mid)).final(encoding)
+          const swap = await cas(older, node)
+          if (!swap || swap === older) return true
+        }
+
         this.changed = true
         this.keys[mid] = key
         return true
@@ -119,6 +124,11 @@ class TreeNode {
       else s = mid + 1
     }
 
+    if (cas && !(await cas(null, node))) {
+      this.changed = false
+      return true
+    }
+
     const i = c < 0 ? e : s
     this.keys.splice(i, 0, key)
     if (child) this.children.splice(i + 1, 0, new Child(0, 0, child))
@@ -736,7 +746,6 @@ class Batch {
       value
     }
     key = enc(encoding.key, key)
-    value = enc(encoding.value, value)
 
     const stack = []
 
@@ -760,10 +769,16 @@ class Batch {
         c = b4a.compare(target.value, await node.getKey(mid))
 
         if (c === 0) {
-          if (cas && !(await cas((await node.getKeyNode(mid)).final(encoding), newNode))) return this._unlockMaybe()
+          if (cas) {
+            const older = (await node.getKeyNode(mid)).final(encoding)
+            const swap = await cas(older, newNode)
+            if (!swap || swap === older) return this._unlockMaybe()
+          }
 
           node.setKey(mid, target)
-          return this._append(root, seq, key, value)
+
+          const valueEncoded = enc(encoding.value, newNode.value)
+          return this._append(root, seq, key, valueEncoded)
         }
 
         if (c < 0) e = mid
@@ -793,7 +808,8 @@ class Batch {
       }
     }
 
-    return this._append(root, seq, key, value)
+    const valueEncoded = enc(encoding.value, newNode.value)
+    return this._append(root, seq, key, valueEncoded)
   }
 
   async del (key, opts) {
@@ -839,7 +855,12 @@ class Batch {
         c = b4a.compare(key, await node.getKey(mid))
 
         if (c === 0) {
-          if (cas && !(await cas((await node.getKeyNode(mid)).final(encoding), delNode))) return this._unlockMaybe()
+          if (cas) {
+            const older = (await node.getKeyNode(mid)).final(encoding)
+            const swap = await cas(older, delNode)
+            if (!swap || swap === older) return this._unlockMaybe()
+          }
+
           if (node.children.length) await setKeyToNearestLeaf(node, mid, stack)
           else node.removeKey(mid)
           // we mark these as changed late, so we don't rewrite them if it is a 404

test/cas.js

@@ -2,6 +2,168 @@ const test = require('brittle')
 const b4a = require('b4a')
 const { create } = require('./helpers')
 
+test('cas - swap with a new value', async function (t) {
+  const db = create({ valueEncoding: 'json' })
+
+  await db.put('/a', 0, {
+    cas: function (prev, next) {
+      t.is(prev, null)
+      t.alike(next, { seq: 1, key: '/a', value: 0 })
+      return next
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 1, key: '/a', value: 0 })
+
+  await db.put('/a', 99, {
+    cas: function (prev, next) {
+      t.alike(prev, { seq: 1, key: '/a', value: 0 })
+      t.alike(next, { seq: 2, key: '/a', value: 99 })
+      next.value = prev.value + 1 // Overwrites so it's not 99 anymore
+      return next
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 2, key: '/a', value: 1 })
+
+  await db.put('/a', 99, {
+    cas: function (prev, next) {
+      t.alike(prev, { seq: 2, key: '/a', value: 1 })
+      t.alike(next, { seq: 3, key: '/a', value: 99 })
+      next.value = prev.value + 1
+      return next
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 3, key: '/a', value: 2 })
+})
+
+test('cas - should not swap', async function (t) {
+  t.plan(4)
+
+  const db = create({ valueEncoding: 'json' })
+
+  await db.put('/a', 1)
+
+  t.alike(await db.get('/a'), { seq: 1, key: '/a', value: 1 })
+
+  await db.put('/a', 2, {
+    cas: function (prev, next) {
+      t.alike(prev, { seq: 1, key: '/a', value: 1 })
+      t.alike(next, { seq: 2, key: '/a', value: 2 })
+      return null
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 1, key: '/a', value: 1 })
+})
+
+test('cas - swap but keep older one', async function (t) {
+  const db = create({ valueEncoding: 'json' })
+
+  await db.put('/a', 0)
+
+  t.alike(await db.get('/a'), { seq: 1, key: '/a', value: 0 })
+
+  await db.put('/a', 99, {
+    cas: function (prev, next) {
+      t.alike(prev, { seq: 1, key: '/a', value: 0 })
+      t.alike(next, { seq: 2, key: '/a', value: 99 })
+      return prev
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 1, key: '/a', value: 0 })
+})
+
+test('cas - swap deletion', async function (t) {
+  const db = create({ valueEncoding: 'json' })
+
+  await db.put('/a', 0)
+
+  await db.del('/a', {
+    cas: function (prev, next) {
+      t.alike(prev, { seq: 1, key: '/a', value: 0 })
+      t.alike(next, { seq: 2, key: '/a', value: null })
+      return prev
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 1, key: '/a', value: 0 })
+
+  await db.del('/a', {
+    cas: function (prev, next) {
+      t.alike(prev, { seq: 1, key: '/a', value: 0 })
+      t.alike(next, { seq: 2, key: '/a', value: null })
+      return next
+    }
+  })
+
+  t.alike(await db.get('/a'), null)
+})
+
+test('cas is called when prev does not exists', async function (t) {
+  t.plan(6)
+
+  const db = create()
+
+  t.comment('first put')
+
+  await db.put('/a', '1', {
+    cas: function (prev, next) {
+      t.comment('first cb')
+
+      t.is(prev, null)
+      t.alike(next, { seq: 1, key: '/a', value: '1' })
+
+      return true
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 1, key: '/a', value: '1' })
+
+  t.comment('second put')
+
+  await db.put('/a', '2', {
+    cas: function (prev, next) {
+      t.comment('second cb')
+
+      t.alike(prev, { seq: 1, key: '/a', value: '1' })
+      t.alike(next, { seq: 2, key: '/a', value: '2' })
+
+      return true
+    }
+  })
+
+  t.alike(await db.get('/a'), { seq: 2, key: '/a', value: '2' })
+})
+
+test('cas is respected when prev does not exists', async function (t) {
+  t.plan(6)
+
+  const db = create()
+
+  await db.put('/a', '1', {
+    cas: function (prev, next) {
+      t.is(prev, null)
+      t.alike(next, { seq: 1, key: '/a', value: '1' })
+      return false
+    }
+  })
+
+  t.is(await db.get('/a'), null)
+
+  await db.put('/a', '2', {
+    cas: function (prev, next) {
+      t.is(prev, null)
+      t.alike(next, { seq: 1, key: '/a', value: '2' })
+      return false
+    }
+  })
+
+  t.is(await db.get('/a'), null)
+})
+
 test('bee.put({ cas }) succeds if cas(last, next) returns truthy', async function (t) {
   const key = 'key'
   const value = 'value'