Rework BTreeData class.

Add inline test cases.
Fix several bugs brought up by test cases.
Alter PersistentString instances when they already exist rather than
creating new ones.

product/ERP5Type/BTreeData.py

@@ -10,94 +10,191 @@ class PersistentString(Persistent):
         return self.value
 
 class BTreeData(Persistent):
+    """
+    In-ZODB (non-BLOB) storage of arbitrary binary data.
+
+    File is managed as chunks, each with a starting offset. Chunks are
+    individually persistent (so they are loaded individually when accessed),
+    and organised in a BTree (so access to any part of the file is in
+    O=log(N)).
+
+    Each call to write() creates a new chunk, so the number and size of chunks
+    is (and must be) controled outside this class.
+
+    It supports sparse files, ie writing one byte at 10M offset will not use
+    10MB on disk. Sparse bytes read as 0x00 (NULL-bytes).
+    """
     def __init__(self):
         self._tree = LOBTree()
 
+    def __len__(self):
+        """
+        Return the position of last data chunk in file.
+        Does not tell how many bytes are actually used.
+        """
+        tree = self._tree
+        try:
+            result = tree.maxKey()
+        except ValueError:
+            return 0
+        return result + len(tree[result])
+
     def write(self, buf, offset):
+        """
+        Create a new chunk at given offset, with given data.
+
+        buf (string)
+          Data to write
+        offset (int)
+          Offset of first data byte.
+        """
         # TODO: auto-aggregation of continuous keys when overwriting
         tree = self._tree
-        buf_len = len(buf)
+        key = offset
         try:
             lower_key = tree.maxKey(offset)
         except ValueError:
-            # No early-enough entry
-            key = offset
+            pass
         else:
+            # Reuse data from an existing and overlapping entry, if any.
+            # Avoids fragmenting the file when overwriting with unaligned
+            # writes.
             if lower_key < offset:
-                value = str(tree[lower_key])
-                value_len = len(value)
+                chunk = tree[lower_key]
+                value_len = len(chunk)
                 if lower_key + value_len > offset:
                     key = lower_key
-                    buf = value[:offset - key] + buf
-                else:
-                    key = offset
+                    buf = chunk.value[:offset - key] + buf
+        while buf:
+            try:
+                next_key = tree.minKey(key + 1)
+            except ValueError:
+                to_write_len = len(buf)
+                next_key = None
             else:
-                assert lower_key == offset, (lower_key, offset)
-                key = lower_key
-        buf_offset = 0
-        to_apply = {}
-        for next_key in tree.iterkeys(key):
-            if buf_offset >= buf_len:
-                break
-            next_buf_offset = buf_offset + (next_key - key)
-            to_apply[key] = PersistentString(buf[buf_offset:next_buf_offset])
-            buf_offset = next_buf_offset
+                to_write_len = next_key - key
+            try:
+                chunk = tree[key]
+            except KeyError:
+                tree[key] = chunk = PersistentString('')
+            entry_size = len(chunk)
+            to_write = buf[:to_write_len]
+            buf = buf[to_write_len:]
+            if to_write_len < entry_size:
+                assert not buf, (key, to_write_len, entry_size)
+                to_write += chunk.value[to_write_len:]
+            chunk.value = to_write
             key = next_key
-        else:
-            to_add = buf[buf_offset:]
-            if to_add:
-                tree[offset + buf_offset] = PersistentString(to_add)
-        for key, value in to_apply.iteritems():
-            tree[key] = value
-        return buf_len
 
     def read(self, offset, size):
-        start_offset = offset
+        """
+        Read data back from object.
+
+        offset (int)
+          Offset of first byte to read.
+        size (int)
+          Number of bytes to read.
+
+        Returns string of read data.
+        """
         tree = self._tree
         try:
             key = tree.maxKey(offset)
         except ValueError:
-            return ''
-        else:
-            offset -= key
+            key = offset
+        # (supposedly) marginal optimisations possible:
+        # - use key found by maxKey
+        # - avoid loading last key if its past the end of read request
+        # Would simlify the loop, but might duplicate code... And CPU is not
+        # expected to be the bottleneck here.
         result = StringIO()
         write = result.write
         written = 0
         for key in tree.iterkeys(key):
-            padding = min(size, key - start_offset - written)
-            if padding:
+            next_byte = offset + written
+            padding = min(size, key - next_byte)
+            if padding > 0:
                 write('\x00' * padding)
                 written += padding
                 size -= padding
+                chunk_offset = 0
+            else:
+                chunk_offset = next_byte - key
             if size == 0:
                 break
             chunk = tree[key]
-            to_write = str(chunk)[offset:offset+size]
+            to_write = chunk.value[chunk_offset:chunk_offset + size]
+            # Free memory used by chunk. Helps avoiding thrashing connection
+            # cache by discarding chunks earlier.
             chunk._p_deactivate()
             to_write_len = len(to_write)
             write(to_write)
             size -= to_write_len
             written += to_write_len
-            offset = 0
         return result.getvalue()
 
     def truncate(self, offset):
+        """
+        Truncate data at given offset.
+
+        offset (int)
+          Offet of the first byte to discard.
+        """
         tree = self._tree
         try:
             key = tree.maxKey(offset)
         except ValueError:
-            # No key below offset, flush everything.
             tree.clear()
         else:
-            value = str(tree[key])
+            chunk = tree[key]
+            chunk_len = len(chunk)
             value_len = offset - key
-            if len(value) > value_len:
-                tree[key] = PersistentString(value[:value_len])
+            if chunk_len > value_len:
+                chunk.value = chunk.value[:value_len]
             minKey = tree.minKey
+            # It is not possible to drop keys as we iterate when using
+            # iterkeys, so call minKey repeatedly.
             while True:
                 next_key = minKey(offset)
                 if next_key is None:
                     break
                 del tree[key]
 
+if __name__ == '__main__':
+
+    def check(tree, length, read_offset, read_length, data, keys=None):
+        tree_length = len(tree)
+        tree_data = tree.read(read_offset, read_length)
+        assert tree_length == length, tree_length
+        assert tree_data == data, repr(tree_data)
+        if keys is not None:
+            tree_keys = list(tree._tree.keys())
+            assert tree_keys == keys, tree_keys
+
+    data = BTreeData()
+
+    data.write('', 10)
+    check(data, 0, 0, 20, '', [])
+
+    data.write('a', 5)
+    check(data, 6, 4, 3, '\x00a', [5])
+
+    data.write('b', 5)
+    check(data, 6, 4, 3, '\x00b', [5])
+
+    data.write('0123456', 0)
+    check(data, 7, 0, 10, '0123456', [0, 5])
+    check(data, 7, 0, 1, '0')
+    check(data, 7, 1, 1, '1')
+    check(data, 7, 2, 1, '2')
+    check(data, 7, 5, 1, '5')
+    check(data, 7, 6, 1, '6')
+
+    data.write('XY', 4)
+    check(data, 7, 0, 10, '0123XY6', [0, 5])
+
+    data.write('a', 10)
+    data.write('8', 8)
+    check(data, 11, 0, 10, '0123XY6\x008\x00', [0, 5, 8, 10])
+    check(data, 11, 7, 10, '\x008\x00a')