~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/btree_index.py

(jameinel) Allow 'bzr serve' to interpret SIGHUP as a graceful shutdown.
 (bug #795025) (John A Meinel)

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2008 Canonical Ltd
 
1
# Copyright (C) 2008-2011 Canonical Ltd
2
2
#
3
3
# This program is free software; you can redistribute it and/or modify
4
4
# it under the terms of the GNU General Public License as published by
12
12
#
13
13
# You should have received a copy of the GNU General Public License
14
14
# along with this program; if not, write to the Free Software
15
 
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
15
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
16
16
#
17
17
 
18
18
"""B+Tree indices"""
19
19
 
20
 
import array
 
20
import cStringIO
 
21
 
 
22
from bzrlib.lazy_import import lazy_import
 
23
lazy_import(globals(), """
21
24
import bisect
22
 
from bisect import bisect_right
23
 
from copy import deepcopy
24
25
import math
25
 
import struct
26
26
import tempfile
27
27
import zlib
 
28
""")
28
29
 
29
30
from bzrlib import (
30
31
    chunk_writer,
31
32
    debug,
32
33
    errors,
 
34
    fifo_cache,
33
35
    index,
34
36
    lru_cache,
35
37
    osutils,
 
38
    static_tuple,
36
39
    trace,
 
40
    transport,
37
41
    )
38
42
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
39
 
from bzrlib.transport import get_transport
40
43
 
41
44
 
42
45
_BTSIGNATURE = "B+Tree Graph Index 2\n"
63
66
    def __init__(self):
64
67
        """Create a _BuilderRow."""
65
68
        self.nodes = 0
66
 
        self.spool = tempfile.TemporaryFile()
 
69
        self.spool = None# tempfile.TemporaryFile(prefix='bzr-index-row-')
67
70
        self.writer = None
68
71
 
69
72
    def finish_node(self, pad=True):
70
73
        byte_lines, _, padding = self.writer.finish()
71
74
        if self.nodes == 0:
 
75
            self.spool = cStringIO.StringIO()
72
76
            # padded note:
73
77
            self.spool.write("\x00" * _RESERVED_HEADER_BYTES)
 
78
        elif self.nodes == 1:
 
79
            # We got bigger than 1 node, switch to a temp file
 
80
            spool = tempfile.TemporaryFile(prefix='bzr-index-row-')
 
81
            spool.write(self.spool.getvalue())
 
82
            self.spool = spool
74
83
        skipped_bytes = 0
75
84
        if not pad and padding:
76
85
            del byte_lines[-1]
153
162
        :param references: An iterable of iterables of keys. Each is a
154
163
            reference to another key.
155
164
        :param value: The value to associate with the key. It may be any
156
 
            bytes as long as it does not contain \0 or \n.
 
165
            bytes as long as it does not contain \\0 or \\n.
157
166
        """
 
167
        # Ensure that 'key' is a StaticTuple
 
168
        key = static_tuple.StaticTuple.from_sequence(key).intern()
158
169
        # we don't care about absent_references
159
170
        node_refs, _ = self._check_key_ref_value(key, references, value)
160
171
        if key in self._nodes:
161
172
            raise errors.BadIndexDuplicateKey(key, self)
162
 
        self._nodes[key] = (node_refs, value)
163
 
        self._keys.add(key)
 
173
        self._nodes[key] = static_tuple.StaticTuple(node_refs, value)
164
174
        if self._nodes_by_key is not None and self._key_length > 1:
165
175
            self._update_nodes_by_key(key, value, node_refs)
166
 
        if len(self._keys) < self._spill_at:
 
176
        if len(self._nodes) < self._spill_at:
167
177
            return
168
178
        self._spill_mem_keys_to_disk()
169
179
 
180
190
        combine mem with the first and second indexes, creating a new one of
181
191
        size 4x. On the fifth create a single new one, etc.
182
192
        """
 
193
        if self._combine_backing_indices:
 
194
            (new_backing_file, size,
 
195
             backing_pos) = self._spill_mem_keys_and_combine()
 
196
        else:
 
197
            new_backing_file, size = self._spill_mem_keys_without_combining()
 
198
        # Note: The transport here isn't strictly needed, because we will use
 
199
        #       direct access to the new_backing._file object
 
200
        new_backing = BTreeGraphIndex(transport.get_transport_from_path('.'),
 
201
                                      '<temp>', size)
 
202
        # GC will clean up the file
 
203
        new_backing._file = new_backing_file
 
204
        if self._combine_backing_indices:
 
205
            if len(self._backing_indices) == backing_pos:
 
206
                self._backing_indices.append(None)
 
207
            self._backing_indices[backing_pos] = new_backing
 
208
            for backing_pos in range(backing_pos):
 
209
                self._backing_indices[backing_pos] = None
 
210
        else:
 
211
            self._backing_indices.append(new_backing)
 
212
        self._nodes = {}
 
213
        self._nodes_by_key = None
 
214
 
 
215
    def _spill_mem_keys_without_combining(self):
 
216
        return self._write_nodes(self._iter_mem_nodes(), allow_optimize=False)
 
217
 
 
218
    def _spill_mem_keys_and_combine(self):
183
219
        iterators_to_combine = [self._iter_mem_nodes()]
184
220
        pos = -1
185
221
        for pos, backing in enumerate(self._backing_indices):
189
225
            iterators_to_combine.append(backing.iter_all_entries())
190
226
        backing_pos = pos + 1
191
227
        new_backing_file, size = \
192
 
            self._write_nodes(self._iter_smallest(iterators_to_combine))
193
 
        dir_path, base_name = osutils.split(new_backing_file.name)
194
 
        # Note: The transport here isn't strictly needed, because we will use
195
 
        #       direct access to the new_backing._file object
196
 
        new_backing = BTreeGraphIndex(get_transport(dir_path),
197
 
                                      base_name, size)
198
 
        # GC will clean up the file
199
 
        new_backing._file = new_backing_file
200
 
        if len(self._backing_indices) == backing_pos:
201
 
            self._backing_indices.append(None)
202
 
        self._backing_indices[backing_pos] = new_backing
203
 
        for pos in range(backing_pos):
204
 
            self._backing_indices[pos] = None
205
 
        self._keys = set()
206
 
        self._nodes = {}
207
 
        self._nodes_by_key = None
 
228
            self._write_nodes(self._iter_smallest(iterators_to_combine),
 
229
                              allow_optimize=False)
 
230
        return new_backing_file, size, backing_pos
208
231
 
209
232
    def add_nodes(self, nodes):
210
233
        """Add nodes to the index.
262
285
            except StopIteration:
263
286
                current_values[pos] = None
264
287
 
265
 
    def _add_key(self, string_key, line, rows):
 
288
    def _add_key(self, string_key, line, rows, allow_optimize=True):
266
289
        """Add a key to the current chunk.
267
290
 
268
291
        :param string_key: The key to add.
269
292
        :param line: The fully serialised key and value.
 
293
        :param allow_optimize: If set to False, prevent setting the optimize
 
294
            flag when writing out. This is used by the _spill_mem_keys_to_disk
 
295
            functionality.
270
296
        """
271
297
        if rows[-1].writer is None:
272
298
            # opening a new leaf chunk;
277
303
                    length = _PAGE_SIZE
278
304
                    if internal_row.nodes == 0:
279
305
                        length -= _RESERVED_HEADER_BYTES # padded
 
306
                    if allow_optimize:
 
307
                        optimize_for_size = self._optimize_for_size
 
308
                    else:
 
309
                        optimize_for_size = False
280
310
                    internal_row.writer = chunk_writer.ChunkWriter(length, 0,
281
 
                        optimize_for_size=self._optimize_for_size)
 
311
                        optimize_for_size=optimize_for_size)
282
312
                    internal_row.writer.write(_INTERNAL_FLAG)
283
313
                    internal_row.writer.write(_INTERNAL_OFFSET +
284
314
                        str(rows[pos + 1].nodes) + "\n")
296
326
            new_row = True
297
327
            for row in reversed(rows[:-1]):
298
328
                # Mark the start of the next node in the node above. If it
299
 
                # doesn't fit then propogate upwards until we find one that
 
329
                # doesn't fit then propagate upwards until we find one that
300
330
                # it does fit into.
301
331
                if row.writer.write(key_line):
302
332
                    row.finish_node()
322
352
                new_row.writer.write(_INTERNAL_OFFSET +
323
353
                    str(rows[1].nodes - 1) + "\n")
324
354
                new_row.writer.write(key_line)
325
 
            self._add_key(string_key, line, rows)
 
355
            self._add_key(string_key, line, rows, allow_optimize=allow_optimize)
326
356
 
327
 
    def _write_nodes(self, node_iterator):
 
357
    def _write_nodes(self, node_iterator, allow_optimize=True):
328
358
        """Write node_iterator out as a B+Tree.
329
359
 
330
360
        :param node_iterator: An iterator of sorted nodes. Each node should
331
361
            match the output given by iter_all_entries.
 
362
        :param allow_optimize: If set to False, prevent setting the optimize
 
363
            flag when writing out. This is used by the _spill_mem_keys_to_disk
 
364
            functionality.
332
365
        :return: A file handle for a temporary file containing a B+Tree for
333
366
            the nodes.
334
367
        """
344
377
        self.row_lengths = []
345
378
        # Loop over all nodes adding them to the bottom row
346
379
        # (rows[-1]). When we finish a chunk in a row,
347
 
        # propogate the key that didn't fit (comes after the chunk) to the
 
380
        # propagate the key that didn't fit (comes after the chunk) to the
348
381
        # row above, transitively.
349
382
        for node in node_iterator:
350
383
            if key_count == 0:
353
386
            key_count += 1
354
387
            string_key, line = _btree_serializer._flatten_node(node,
355
388
                                    self.reference_lists)
356
 
            self._add_key(string_key, line, rows)
 
389
            self._add_key(string_key, line, rows, allow_optimize=allow_optimize)
357
390
        for row in reversed(rows):
358
391
            pad = (type(row) != _LeafBuilderRow)
359
392
            row.finish_node(pad=pad)
360
 
        result = tempfile.NamedTemporaryFile()
361
393
        lines = [_BTSIGNATURE]
362
394
        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')
363
395
        lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n')
364
396
        lines.append(_OPTION_LEN + str(key_count) + '\n')
365
397
        row_lengths = [row.nodes for row in rows]
366
398
        lines.append(_OPTION_ROW_LENGTHS + ','.join(map(str, row_lengths)) + '\n')
 
399
        if row_lengths and row_lengths[-1] > 1:
 
400
            result = tempfile.NamedTemporaryFile(prefix='bzr-index-')
 
401
        else:
 
402
            result = cStringIO.StringIO()
367
403
        result.writelines(lines)
368
404
        position = sum(map(len, lines))
369
405
        root_row = True
380
416
            # Special case the first node as it may be prefixed
381
417
            node = row.spool.read(_PAGE_SIZE)
382
418
            result.write(node[reserved:])
383
 
            result.write("\x00" * (reserved - position))
 
419
            if len(node) == _PAGE_SIZE:
 
420
                result.write("\x00" * (reserved - position))
384
421
            position = 0 # Only the root row actually has an offset
385
422
            copied_len = osutils.pumpfile(row.spool, result)
386
423
            if copied_len != (row.nodes - 1) * _PAGE_SIZE:
405
442
    def iter_all_entries(self):
406
443
        """Iterate over all keys within the index
407
444
 
408
 
        :return: An iterable of (index, key, reference_lists, value). There is no
409
 
            defined order for the result iteration - it will be in the most
 
445
        :return: An iterable of (index, key, value, reference_lists). There is
 
446
            no defined order for the result iteration - it will be in the most
410
447
            efficient order for the index (in this case dictionary hash order).
411
448
        """
412
449
        if 'evil' in debug.debug_flags:
431
468
            efficient order for the index (keys iteration order in this case).
432
469
        """
433
470
        keys = set(keys)
434
 
        local_keys = keys.intersection(self._keys)
 
471
        # Note: We don't use keys.intersection() here. If you read the C api,
 
472
        #       set.intersection(other) special cases when other is a set and
 
473
        #       will iterate the smaller of the two and lookup in the other.
 
474
        #       It does *not* do this for any other type (even dict, unlike
 
475
        #       some other set functions.) Since we expect keys is generally <<
 
476
        #       self._nodes, it is faster to iterate over it in a list
 
477
        #       comprehension
 
478
        nodes = self._nodes
 
479
        local_keys = [key for key in keys if key in nodes]
435
480
        if self.reference_lists:
436
481
            for key in local_keys:
437
 
                node = self._nodes[key]
 
482
                node = nodes[key]
438
483
                yield self, key, node[1], node[0]
439
484
        else:
440
485
            for key in local_keys:
441
 
                node = self._nodes[key]
 
486
                node = nodes[key]
442
487
                yield self, key, node[1]
443
488
        # Find things that are in backing indices that have not been handled
444
489
        # yet.
527
572
                    else:
528
573
                        # yield keys
529
574
                        for value in key_dict.itervalues():
530
 
                            yield (self, ) + value
 
575
                            yield (self, ) + tuple(value)
531
576
            else:
532
577
                yield (self, ) + key_dict
533
578
 
554
599
 
555
600
        For InMemoryGraphIndex the estimate is exact.
556
601
        """
557
 
        return len(self._keys) + sum(backing.key_count() for backing in
 
602
        return len(self._nodes) + sum(backing.key_count() for backing in
558
603
            self._backing_indices if backing is not None)
559
604
 
560
605
    def validate(self):
561
606
        """In memory index's have no known corruption at the moment."""
562
607
 
563
608
 
564
 
class _LeafNode(object):
 
609
class _LeafNode(dict):
565
610
    """A leaf node for a serialised B+Tree index."""
566
611
 
 
612
    __slots__ = ('min_key', 'max_key', '_keys')
 
613
 
567
614
    def __init__(self, bytes, key_length, ref_list_length):
568
615
        """Parse bytes to create a leaf node object."""
569
616
        # splitlines mangles the \r delimiters.. don't use it.
570
 
        self.keys = dict(_btree_serializer._parse_leaf_lines(bytes,
571
 
            key_length, ref_list_length))
 
617
        key_list = _btree_serializer._parse_leaf_lines(bytes,
 
618
            key_length, ref_list_length)
 
619
        if key_list:
 
620
            self.min_key = key_list[0][0]
 
621
            self.max_key = key_list[-1][0]
 
622
        else:
 
623
            self.min_key = self.max_key = None
 
624
        super(_LeafNode, self).__init__(key_list)
 
625
        self._keys = dict(self)
 
626
 
 
627
    def all_items(self):
 
628
        """Return a sorted list of (key, (value, refs)) items"""
 
629
        items = self.items()
 
630
        items.sort()
 
631
        return items
 
632
 
 
633
    def all_keys(self):
 
634
        """Return a sorted list of all keys."""
 
635
        keys = self.keys()
 
636
        keys.sort()
 
637
        return keys
572
638
 
573
639
 
574
640
class _InternalNode(object):
575
641
    """An internal node for a serialised B+Tree index."""
576
642
 
 
643
    __slots__ = ('keys', 'offset')
 
644
 
577
645
    def __init__(self, bytes):
578
646
        """Parse bytes to create an internal node object."""
579
647
        # splitlines mangles the \r delimiters.. don't use it.
582
650
    def _parse_lines(self, lines):
583
651
        nodes = []
584
652
        self.offset = int(lines[1][7:])
 
653
        as_st = static_tuple.StaticTuple.from_sequence
585
654
        for line in lines[2:]:
586
655
            if line == '':
587
656
                break
588
 
            nodes.append(tuple(line.split('\0')))
 
657
            nodes.append(as_st(map(intern, line.split('\0'))).intern())
589
658
        return nodes
590
659
 
591
660
 
596
665
    memory except when very large walks are done.
597
666
    """
598
667
 
599
 
    def __init__(self, transport, name, size):
 
668
    def __init__(self, transport, name, size, unlimited_cache=False,
 
669
                 offset=0):
600
670
        """Create a B+Tree index object on the index name.
601
671
 
602
672
        :param transport: The transport to read data for the index from.
606
676
            the initial read (to read the root node header) can be done
607
677
            without over-reading even on empty indices, and on small indices
608
678
            allows single-IO to read the entire index.
 
679
        :param unlimited_cache: If set to True, then instead of using an
 
680
            LRUCache with size _NODE_CACHE_SIZE, we will use a dict and always
 
681
            cache all leaf nodes.
 
682
        :param offset: The start of the btree index data isn't byte 0 of the
 
683
            file. Instead it starts at some point later.
609
684
        """
610
685
        self._transport = transport
611
686
        self._name = name
613
688
        self._file = None
614
689
        self._recommended_pages = self._compute_recommended_pages()
615
690
        self._root_node = None
 
691
        self._base_offset = offset
 
692
        self._leaf_factory = _LeafNode
616
693
        # Default max size is 100,000 leave values
617
694
        self._leaf_value_cache = None # lru_cache.LRUCache(100*1000)
618
 
        self._leaf_node_cache = lru_cache.LRUCache(_NODE_CACHE_SIZE)
619
 
        self._internal_node_cache = lru_cache.LRUCache()
 
695
        if unlimited_cache:
 
696
            self._leaf_node_cache = {}
 
697
            self._internal_node_cache = {}
 
698
        else:
 
699
            self._leaf_node_cache = lru_cache.LRUCache(_NODE_CACHE_SIZE)
 
700
            # We use a FIFO here just to prevent possible blowout. However, a
 
701
            # 300k record btree has only 3k leaf nodes, and only 20 internal
 
702
            # nodes. A value of 100 scales to ~100*100*100 = 1M records.
 
703
            self._internal_node_cache = fifo_cache.FIFOCache(100)
620
704
        self._key_count = None
621
705
        self._row_lengths = None
622
706
        self._row_offsets = None # Start of each row, [-1] is the end
654
738
                if start_of_leaves is None:
655
739
                    start_of_leaves = self._row_offsets[-2]
656
740
                if node_pos < start_of_leaves:
657
 
                    self._internal_node_cache.add(node_pos, node)
 
741
                    self._internal_node_cache[node_pos] = node
658
742
                else:
659
 
                    self._leaf_node_cache.add(node_pos, node)
 
743
                    self._leaf_node_cache[node_pos] = node
660
744
            found[node_pos] = node
661
745
        return found
662
746
 
801
885
            new_tips = next_tips
802
886
        return final_offsets
803
887
 
 
888
    def clear_cache(self):
 
889
        """Clear out any cached/memoized values.
 
890
 
 
891
        This can be called at any time, but generally it is used when we have
 
892
        extracted some information, but don't expect to be requesting any more
 
893
        from this index.
 
894
        """
 
895
        # Note that we don't touch self._root_node or self._internal_node_cache
 
896
        # We don't expect either of those to be big, and it can save
 
897
        # round-trips in the future. We may re-evaluate this if InternalNode
 
898
        # memory starts to be an issue.
 
899
        self._leaf_node_cache.clear()
 
900
 
 
901
    def external_references(self, ref_list_num):
 
902
        if self._root_node is None:
 
903
            self._get_root_node()
 
904
        if ref_list_num + 1 > self.node_ref_lists:
 
905
            raise ValueError('No ref list %d, index has %d ref lists'
 
906
                % (ref_list_num, self.node_ref_lists))
 
907
        keys = set()
 
908
        refs = set()
 
909
        for node in self.iter_all_entries():
 
910
            keys.add(node[1])
 
911
            refs.update(node[3][ref_list_num])
 
912
        return refs - keys
 
913
 
804
914
    def _find_layer_first_and_end(self, offset):
805
915
        """Find the start/stop nodes for the layer corresponding to offset.
806
916
 
858
968
        """Cache directly from key => value, skipping the btree."""
859
969
        if self._leaf_value_cache is not None:
860
970
            for node in nodes.itervalues():
861
 
                for key, value in node.keys.iteritems():
 
971
                for key, value in node.all_items():
862
972
                    if key in self._leaf_value_cache:
863
973
                        # Don't add the rest of the keys, we've seen this node
864
974
                        # before.
888
998
        if self._row_offsets[-1] == 1:
889
999
            # There is only the root node, and we read that via key_count()
890
1000
            if self.node_ref_lists:
891
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
1001
                for key, (value, refs) in self._root_node.all_items():
892
1002
                    yield (self, key, value, refs)
893
1003
            else:
894
 
                for key, (value, refs) in sorted(self._root_node.keys.items()):
 
1004
                for key, (value, refs) in self._root_node.all_items():
895
1005
                    yield (self, key, value)
896
1006
            return
897
1007
        start_of_leaves = self._row_offsets[-2]
907
1017
        # for spilling index builds to disk.
908
1018
        if self.node_ref_lists:
909
1019
            for _, node in nodes:
910
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1020
                for key, (value, refs) in node.all_items():
911
1021
                    yield (self, key, value, refs)
912
1022
        else:
913
1023
            for _, node in nodes:
914
 
                for key, (value, refs) in sorted(node.keys.items()):
 
1024
                for key, (value, refs) in node.all_items():
915
1025
                    yield (self, key, value)
916
1026
 
917
1027
    @staticmethod
941
1051
        # iter_steps = len(in_keys) + len(fixed_keys)
942
1052
        # bisect_steps = len(in_keys) * math.log(len(fixed_keys), 2)
943
1053
        if len(in_keys) == 1: # Bisect will always be faster for M = 1
944
 
            return [(bisect_right(fixed_keys, in_keys[0]), in_keys)]
 
1054
            return [(bisect.bisect_right(fixed_keys, in_keys[0]), in_keys)]
945
1055
        # elif bisect_steps < iter_steps:
946
1056
        #     offsets = {}
947
1057
        #     for key in in_keys:
996
1106
            output.append(cur_out)
997
1107
        return output
998
1108
 
 
1109
    def _walk_through_internal_nodes(self, keys):
 
1110
        """Take the given set of keys, and find the corresponding LeafNodes.
 
1111
 
 
1112
        :param keys: An unsorted iterable of keys to search for
 
1113
        :return: (nodes, index_and_keys)
 
1114
            nodes is a dict mapping {index: LeafNode}
 
1115
            keys_at_index is a list of tuples of [(index, [keys for Leaf])]
 
1116
        """
 
1117
        # 6 seconds spent in miss_torture using the sorted() line.
 
1118
        # Even with out of order disk IO it seems faster not to sort it when
 
1119
        # large queries are being made.
 
1120
        keys_at_index = [(0, sorted(keys))]
 
1121
 
 
1122
        for row_pos, next_row_start in enumerate(self._row_offsets[1:-1]):
 
1123
            node_indexes = [idx for idx, s_keys in keys_at_index]
 
1124
            nodes = self._get_internal_nodes(node_indexes)
 
1125
 
 
1126
            next_nodes_and_keys = []
 
1127
            for node_index, sub_keys in keys_at_index:
 
1128
                node = nodes[node_index]
 
1129
                positions = self._multi_bisect_right(sub_keys, node.keys)
 
1130
                node_offset = next_row_start + node.offset
 
1131
                next_nodes_and_keys.extend([(node_offset + pos, s_keys)
 
1132
                                           for pos, s_keys in positions])
 
1133
            keys_at_index = next_nodes_and_keys
 
1134
        # We should now be at the _LeafNodes
 
1135
        node_indexes = [idx for idx, s_keys in keys_at_index]
 
1136
 
 
1137
        # TODO: We may *not* want to always read all the nodes in one
 
1138
        #       big go. Consider setting a max size on this.
 
1139
        nodes = self._get_leaf_nodes(node_indexes)
 
1140
        return nodes, keys_at_index
 
1141
 
999
1142
    def iter_entries(self, keys):
1000
1143
        """Iterate over keys within the index.
1001
1144
 
1039
1182
        needed_keys = keys
1040
1183
        if not needed_keys:
1041
1184
            return
1042
 
        # 6 seconds spent in miss_torture using the sorted() line.
1043
 
        # Even with out of order disk IO it seems faster not to sort it when
1044
 
        # large queries are being made.
1045
 
        needed_keys = sorted(needed_keys)
1046
 
 
1047
 
        nodes_and_keys = [(0, needed_keys)]
1048
 
 
1049
 
        for row_pos, next_row_start in enumerate(self._row_offsets[1:-1]):
1050
 
            node_indexes = [idx for idx, s_keys in nodes_and_keys]
1051
 
            nodes = self._get_internal_nodes(node_indexes)
1052
 
 
1053
 
            next_nodes_and_keys = []
1054
 
            for node_index, sub_keys in nodes_and_keys:
1055
 
                node = nodes[node_index]
1056
 
                positions = self._multi_bisect_right(sub_keys, node.keys)
1057
 
                node_offset = next_row_start + node.offset
1058
 
                next_nodes_and_keys.extend([(node_offset + pos, s_keys)
1059
 
                                           for pos, s_keys in positions])
1060
 
            nodes_and_keys = next_nodes_and_keys
1061
 
        # We should now be at the _LeafNodes
1062
 
        node_indexes = [idx for idx, s_keys in nodes_and_keys]
1063
 
 
1064
 
        # TODO: We may *not* want to always read all the nodes in one
1065
 
        #       big go. Consider setting a max size on this.
1066
 
 
1067
 
        nodes = self._get_leaf_nodes(node_indexes)
 
1185
        nodes, nodes_and_keys = self._walk_through_internal_nodes(needed_keys)
1068
1186
        for node_index, sub_keys in nodes_and_keys:
1069
1187
            if not sub_keys:
1070
1188
                continue
1071
1189
            node = nodes[node_index]
1072
1190
            for next_sub_key in sub_keys:
1073
 
                if next_sub_key in node.keys:
1074
 
                    value, refs = node.keys[next_sub_key]
 
1191
                if next_sub_key in node:
 
1192
                    value, refs = node[next_sub_key]
1075
1193
                    if self.node_ref_lists:
1076
1194
                        yield (self, next_sub_key, value, refs)
1077
1195
                    else:
1078
1196
                        yield (self, next_sub_key, value)
1079
1197
 
 
1198
    def _find_ancestors(self, keys, ref_list_num, parent_map, missing_keys):
 
1199
        """Find the parent_map information for the set of keys.
 
1200
 
 
1201
        This populates the parent_map dict and missing_keys set based on the
 
1202
        queried keys. It also can fill out an arbitrary number of parents that
 
1203
        it finds while searching for the supplied keys.
 
1204
 
 
1205
        It is unlikely that you want to call this directly. See
 
1206
        "CombinedGraphIndex.find_ancestry()" for a more appropriate API.
 
1207
 
 
1208
        :param keys: A keys whose ancestry we want to return
 
1209
            Every key will either end up in 'parent_map' or 'missing_keys'.
 
1210
        :param ref_list_num: This index in the ref_lists is the parents we
 
1211
            care about.
 
1212
        :param parent_map: {key: parent_keys} for keys that are present in this
 
1213
            index. This may contain more entries than were in 'keys', that are
 
1214
            reachable ancestors of the keys requested.
 
1215
        :param missing_keys: keys which are known to be missing in this index.
 
1216
            This may include parents that were not directly requested, but we
 
1217
            were able to determine that they are not present in this index.
 
1218
        :return: search_keys    parents that were found but not queried to know
 
1219
            if they are missing or present. Callers can re-query this index for
 
1220
            those keys, and they will be placed into parent_map or missing_keys
 
1221
        """
 
1222
        if not self.key_count():
 
1223
            # We use key_count() to trigger reading the root node and
 
1224
            # determining info about this BTreeGraphIndex
 
1225
            # If we don't have any keys, then everything is missing
 
1226
            missing_keys.update(keys)
 
1227
            return set()
 
1228
        if ref_list_num >= self.node_ref_lists:
 
1229
            raise ValueError('No ref list %d, index has %d ref lists'
 
1230
                % (ref_list_num, self.node_ref_lists))
 
1231
 
 
1232
        # The main trick we are trying to accomplish is that when we find a
 
1233
        # key listing its parents, we expect that the parent key is also likely
 
1234
        # to sit on the same page. Allowing us to expand parents quickly
 
1235
        # without suffering the full stack of bisecting, etc.
 
1236
        nodes, nodes_and_keys = self._walk_through_internal_nodes(keys)
 
1237
 
 
1238
        # These are parent keys which could not be immediately resolved on the
 
1239
        # page where the child was present. Note that we may already be
 
1240
        # searching for that key, and it may actually be present [or known
 
1241
        # missing] on one of the other pages we are reading.
 
1242
        # TODO:
 
1243
        #   We could try searching for them in the immediate previous or next
 
1244
        #   page. If they occur "later" we could put them in a pending lookup
 
1245
        #   set, and then for each node we read thereafter we could check to
 
1246
        #   see if they are present.
 
1247
        #   However, we don't know the impact of keeping this list of things
 
1248
        #   that I'm going to search for every node I come across from here on
 
1249
        #   out.
 
1250
        #   It doesn't handle the case when the parent key is missing on a
 
1251
        #   page that we *don't* read. So we already have to handle being
 
1252
        #   re-entrant for that.
 
1253
        #   Since most keys contain a date string, they are more likely to be
 
1254
        #   found earlier in the file than later, but we would know that right
 
1255
        #   away (key < min_key), and wouldn't keep searching it on every other
 
1256
        #   page that we read.
 
1257
        #   Mostly, it is an idea, one which should be benchmarked.
 
1258
        parents_not_on_page = set()
 
1259
 
 
1260
        for node_index, sub_keys in nodes_and_keys:
 
1261
            if not sub_keys:
 
1262
                continue
 
1263
            # sub_keys is all of the keys we are looking for that should exist
 
1264
            # on this page, if they aren't here, then they won't be found
 
1265
            node = nodes[node_index]
 
1266
            parents_to_check = set()
 
1267
            for next_sub_key in sub_keys:
 
1268
                if next_sub_key not in node:
 
1269
                    # This one is just not present in the index at all
 
1270
                    missing_keys.add(next_sub_key)
 
1271
                else:
 
1272
                    value, refs = node[next_sub_key]
 
1273
                    parent_keys = refs[ref_list_num]
 
1274
                    parent_map[next_sub_key] = parent_keys
 
1275
                    parents_to_check.update(parent_keys)
 
1276
            # Don't look for things we've already found
 
1277
            parents_to_check = parents_to_check.difference(parent_map)
 
1278
            # this can be used to test the benefit of having the check loop
 
1279
            # inlined.
 
1280
            # parents_not_on_page.update(parents_to_check)
 
1281
            # continue
 
1282
            while parents_to_check:
 
1283
                next_parents_to_check = set()
 
1284
                for key in parents_to_check:
 
1285
                    if key in node:
 
1286
                        value, refs = node[key]
 
1287
                        parent_keys = refs[ref_list_num]
 
1288
                        parent_map[key] = parent_keys
 
1289
                        next_parents_to_check.update(parent_keys)
 
1290
                    else:
 
1291
                        # This parent either is genuinely missing, or should be
 
1292
                        # found on another page. Perf test whether it is better
 
1293
                        # to check if this node should fit on this page or not.
 
1294
                        # in the 'everything-in-one-pack' scenario, this *not*
 
1295
                        # doing the check is 237ms vs 243ms.
 
1296
                        # So slightly better, but I assume the standard 'lots
 
1297
                        # of packs' is going to show a reasonable improvement
 
1298
                        # from the check, because it avoids 'going around
 
1299
                        # again' for everything that is in another index
 
1300
                        # parents_not_on_page.add(key)
 
1301
                        # Missing for some reason
 
1302
                        if key < node.min_key:
 
1303
                            # in the case of bzr.dev, 3.4k/5.3k misses are
 
1304
                            # 'earlier' misses (65%)
 
1305
                            parents_not_on_page.add(key)
 
1306
                        elif key > node.max_key:
 
1307
                            # This parent key would be present on a different
 
1308
                            # LeafNode
 
1309
                            parents_not_on_page.add(key)
 
1310
                        else:
 
1311
                            # assert key != node.min_key and key != node.max_key
 
1312
                            # If it was going to be present, it would be on
 
1313
                            # *this* page, so mark it missing.
 
1314
                            missing_keys.add(key)
 
1315
                parents_to_check = next_parents_to_check.difference(parent_map)
 
1316
                # Might want to do another .difference() from missing_keys
 
1317
        # parents_not_on_page could have been found on a different page, or be
 
1318
        # known to be missing. So cull out everything that has already been
 
1319
        # found.
 
1320
        search_keys = parents_not_on_page.difference(
 
1321
            parent_map).difference(missing_keys)
 
1322
        return search_keys
 
1323
 
1080
1324
    def iter_entries_prefix(self, keys):
1081
1325
        """Iterate over keys within the index using prefix matching.
1082
1326
 
1272
1516
        # list of (offset, length) regions of the file that should, evenually
1273
1517
        # be read in to data_ranges, either from 'bytes' or from the transport
1274
1518
        ranges = []
 
1519
        base_offset = self._base_offset
1275
1520
        for index in nodes:
1276
 
            offset = index * _PAGE_SIZE
 
1521
            offset = (index * _PAGE_SIZE)
1277
1522
            size = _PAGE_SIZE
1278
1523
            if index == 0:
1279
1524
                # Root node - special case
1283
1528
                    # The only case where we don't know the size, is for very
1284
1529
                    # small indexes. So we read the whole thing
1285
1530
                    bytes = self._transport.get_bytes(self._name)
1286
 
                    self._size = len(bytes)
 
1531
                    num_bytes = len(bytes)
 
1532
                    self._size = num_bytes - base_offset
1287
1533
                    # the whole thing should be parsed out of 'bytes'
1288
 
                    ranges.append((0, len(bytes)))
 
1534
                    ranges = [(start, min(_PAGE_SIZE, num_bytes - start))
 
1535
                        for start in xrange(base_offset, num_bytes, _PAGE_SIZE)]
1289
1536
                    break
1290
1537
            else:
1291
1538
                if offset > self._size:
1293
1540
                                         ' of the file %s > %s'
1294
1541
                                         % (offset, self._size))
1295
1542
                size = min(size, self._size - offset)
1296
 
            ranges.append((offset, size))
 
1543
            ranges.append((base_offset + offset, size))
1297
1544
        if not ranges:
1298
1545
            return
1299
1546
        elif bytes is not None:
1300
1547
            # already have the whole file
1301
 
            data_ranges = [(start, bytes[start:start+_PAGE_SIZE])
1302
 
                           for start in xrange(0, len(bytes), _PAGE_SIZE)]
 
1548
            data_ranges = [(start, bytes[start:start+size])
 
1549
                           for start, size in ranges]
1303
1550
        elif self._file is None:
1304
1551
            data_ranges = self._transport.readv(self._name, ranges)
1305
1552
        else:
1308
1555
                self._file.seek(offset)
1309
1556
                data_ranges.append((offset, self._file.read(size)))
1310
1557
        for offset, data in data_ranges:
 
1558
            offset -= base_offset
1311
1559
            if offset == 0:
1312
1560
                # extract the header
1313
1561
                offset, data = self._parse_header_from_bytes(data)
1315
1563
                    continue
1316
1564
            bytes = zlib.decompress(data)
1317
1565
            if bytes.startswith(_LEAF_FLAG):
1318
 
                node = _LeafNode(bytes, self._key_length, self.node_ref_lists)
 
1566
                node = self._leaf_factory(bytes, self._key_length,
 
1567
                                          self.node_ref_lists)
1319
1568
            elif bytes.startswith(_INTERNAL_FLAG):
1320
1569
                node = _InternalNode(bytes)
1321
1570
            else:
1340
1589
            pass
1341
1590
 
1342
1591
 
 
1592
_gcchk_factory = _LeafNode
 
1593
 
1343
1594
try:
1344
 
    from bzrlib import _btree_serializer_c as _btree_serializer
1345
 
except ImportError:
 
1595
    from bzrlib import _btree_serializer_pyx as _btree_serializer
 
1596
    _gcchk_factory = _btree_serializer._parse_into_chk
 
1597
except ImportError, e:
 
1598
    osutils.failed_to_load_extension(e)
1346
1599
    from bzrlib import _btree_serializer_py as _btree_serializer