~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/chk_map.py

  • Committer: John Arbash Meinel
  • Date: 2010-08-02 17:16:12 UTC
  • mto: This revision was merged to the branch mainline in revision 5369.
  • Revision ID: john@arbash-meinel.com-20100802171612-rdh5ods70w2bl3j7
We also have to re-implement it for _simple_set_pyx.pyx

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2008 Canonical Ltd
 
1
# Copyright (C) 2008, 2009, 2010 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
26
26
 
27
27
Updates to a CHKMap are done preferentially via the apply_delta method, to
28
28
allow optimisation of the update operation; but individual map/unmap calls are
29
 
possible and supported. All changes via map/unmap are buffered in memory until
30
 
the _save method is called to force serialisation of the tree. apply_delta
31
 
performs a _save implicitly.
 
29
possible and supported. Individual changes via map/unmap are buffered in memory
 
30
until the _save method is called to force serialisation of the tree.
 
31
apply_delta records its changes immediately by performing an implicit _save.
32
32
 
33
33
TODO:
34
34
-----
38
38
"""
39
39
 
40
40
import heapq
41
 
import time
 
41
import threading
42
42
 
43
43
from bzrlib import lazy_import
44
44
lazy_import.lazy_import(globals(), """
45
 
from bzrlib import versionedfile
 
45
from bzrlib import (
 
46
    errors,
 
47
    versionedfile,
 
48
    )
46
49
""")
47
50
from bzrlib import (
48
 
    errors,
49
51
    lru_cache,
50
52
    osutils,
51
53
    registry,
 
54
    static_tuple,
52
55
    trace,
53
56
    )
 
57
from bzrlib.static_tuple import StaticTuple
54
58
 
55
59
# approx 4MB
56
60
# If each line is 50 bytes, and you have 255 internal pages, with 255-way fan
57
61
# out, it takes 3.1MB to cache the layer.
58
62
_PAGE_CACHE_SIZE = 4*1024*1024
59
 
# We are caching bytes so len(value) is perfectly accurate
60
 
_page_cache = lru_cache.LRUSizeCache(_PAGE_CACHE_SIZE)
 
63
# Per thread caches for 2 reasons:
 
64
# - in the server we may be serving very different content, so we get less
 
65
#   cache thrashing.
 
66
# - we avoid locking on every cache lookup.
 
67
_thread_caches = threading.local()
 
68
# The page cache.
 
69
_thread_caches.page_cache = None
 
70
 
 
71
def _get_cache():
 
72
    """Get the per-thread page cache.
 
73
 
 
74
    We need a function to do this because in a new thread the _thread_caches
 
75
    threading.local object does not have the cache initialized yet.
 
76
    """
 
77
    page_cache = getattr(_thread_caches, 'page_cache', None)
 
78
    if page_cache is None:
 
79
        # We are caching bytes so len(value) is perfectly accurate
 
80
        page_cache = lru_cache.LRUSizeCache(_PAGE_CACHE_SIZE)
 
81
        _thread_caches.page_cache = page_cache
 
82
    return page_cache
 
83
 
 
84
 
 
85
def clear_cache():
 
86
    _get_cache().clear()
 
87
 
61
88
 
62
89
# If a ChildNode falls below this many bytes, we check for a remap
63
90
_INTERESTING_NEW_SIZE = 50
79
106
class CHKMap(object):
80
107
    """A persistent map from string to string backed by a CHK store."""
81
108
 
 
109
    __slots__ = ('_store', '_root_node', '_search_key_func')
 
110
 
82
111
    def __init__(self, store, root_key, search_key_func=None):
83
112
        """Create a CHKMap object.
84
113
 
107
136
            of old_key is removed.
108
137
        """
109
138
        delete_count = 0
 
139
        # Check preconditions first.
 
140
        as_st = StaticTuple.from_sequence
 
141
        new_items = set([as_st(key) for (old, key, value) in delta
 
142
                         if key is not None and old is None])
 
143
        existing_new = list(self.iteritems(key_filter=new_items))
 
144
        if existing_new:
 
145
            raise errors.InconsistentDeltaDelta(delta,
 
146
                "New items are already in the map %r." % existing_new)
 
147
        # Now apply changes.
110
148
        for old, new, value in delta:
111
149
            if old is not None and old != new:
112
150
                self.unmap(old, check_remap=False)
121
159
 
122
160
    def _ensure_root(self):
123
161
        """Ensure that the root node is an object not a key."""
124
 
        if type(self._root_node) is tuple:
 
162
        if type(self._root_node) is StaticTuple:
125
163
            # Demand-load the root
126
164
            self._root_node = self._get_node(self._root_node)
127
165
 
135
173
        :param node: A tuple key or node object.
136
174
        :return: A node object.
137
175
        """
138
 
        if type(node) is tuple:
 
176
        if type(node) is StaticTuple:
139
177
            bytes = self._read_bytes(node)
140
178
            return _deserialise(bytes, node,
141
179
                search_key_func=self._search_key_func)
144
182
 
145
183
    def _read_bytes(self, key):
146
184
        try:
147
 
            return _page_cache[key]
 
185
            return _get_cache()[key]
148
186
        except KeyError:
149
187
            stream = self._store.get_record_stream([key], 'unordered', True)
150
188
            bytes = stream.next().get_bytes_as('fulltext')
151
 
            _page_cache[key] = bytes
 
189
            _get_cache()[key] = bytes
152
190
            return bytes
153
191
 
154
192
    def _dump_tree(self, include_keys=False):
182
220
            for key, value in sorted(node._items.iteritems()):
183
221
                # Don't use prefix nor indent here to line up when used in
184
222
                # tests in conjunction with assertEqualDiff
185
 
                result.append('      %r %r' % (key, value))
 
223
                result.append('      %r %r' % (tuple(key), value))
186
224
        return result
187
225
 
188
226
    @classmethod
203
241
            multiple pages.
204
242
        :return: The root chk of the resulting CHKMap.
205
243
        """
206
 
        result = CHKMap(store, None, search_key_func=search_key_func)
 
244
        root_key = klass._create_directly(store, initial_value,
 
245
            maximum_size=maximum_size, key_width=key_width,
 
246
            search_key_func=search_key_func)
 
247
        if type(root_key) is not StaticTuple:
 
248
            raise AssertionError('we got a %s instead of a StaticTuple'
 
249
                                 % (type(root_key),))
 
250
        return root_key
 
251
 
 
252
    @classmethod
 
253
    def _create_via_map(klass, store, initial_value, maximum_size=0,
 
254
                        key_width=1, search_key_func=None):
 
255
        result = klass(store, None, search_key_func=search_key_func)
207
256
        result._root_node.set_maximum_size(maximum_size)
208
257
        result._root_node._key_width = key_width
209
258
        delta = []
210
259
        for key, value in initial_value.items():
211
260
            delta.append((None, key, value))
212
 
        return result.apply_delta(delta)
 
261
        root_key = result.apply_delta(delta)
 
262
        return root_key
 
263
 
 
264
    @classmethod
 
265
    def _create_directly(klass, store, initial_value, maximum_size=0,
 
266
                         key_width=1, search_key_func=None):
 
267
        node = LeafNode(search_key_func=search_key_func)
 
268
        node.set_maximum_size(maximum_size)
 
269
        node._key_width = key_width
 
270
        as_st = StaticTuple.from_sequence
 
271
        node._items = dict([(as_st(key), val) for key, val
 
272
                                               in initial_value.iteritems()])
 
273
        node._raw_size = sum([node._key_value_len(key, value)
 
274
                              for key,value in node._items.iteritems()])
 
275
        node._len = len(node._items)
 
276
        node._compute_search_prefix()
 
277
        node._compute_serialised_prefix()
 
278
        if (node._len > 1
 
279
            and maximum_size
 
280
            and node._current_size() > maximum_size):
 
281
            prefix, node_details = node._split(store)
 
282
            if len(node_details) == 1:
 
283
                raise AssertionError('Failed to split using node._split')
 
284
            node = InternalNode(prefix, search_key_func=search_key_func)
 
285
            node.set_maximum_size(maximum_size)
 
286
            node._key_width = key_width
 
287
            for split, subnode in node_details:
 
288
                node.add_node(split, subnode)
 
289
        keys = list(node.serialise(store))
 
290
        return keys[-1]
213
291
 
214
292
    def iter_changes(self, basis):
215
293
        """Iterate over the changes between basis and self.
435
513
    def iteritems(self, key_filter=None):
436
514
        """Iterate over the entire CHKMap's contents."""
437
515
        self._ensure_root()
 
516
        if key_filter is not None:
 
517
            as_st = StaticTuple.from_sequence
 
518
            key_filter = [as_st(key) for key in key_filter]
438
519
        return self._root_node.iteritems(self._store, key_filter=key_filter)
439
520
 
440
521
    def key(self):
441
522
        """Return the key for this map."""
442
 
        if type(self._root_node) is tuple:
 
523
        if type(self._root_node) is StaticTuple:
443
524
            return self._root_node
444
525
        else:
445
526
            return self._root_node._key
449
530
        return len(self._root_node)
450
531
 
451
532
    def map(self, key, value):
452
 
        """Map a key tuple to value."""
 
533
        """Map a key tuple to value.
 
534
        
 
535
        :param key: A key to map.
 
536
        :param value: The value to assign to key.
 
537
        """
 
538
        key = StaticTuple.from_sequence(key)
453
539
        # Need a root object.
454
540
        self._ensure_root()
455
541
        prefix, node_details = self._root_node.map(self._store, key, value)
466
552
    def _node_key(self, node):
467
553
        """Get the key for a node whether it's a tuple or node."""
468
554
        if type(node) is tuple:
 
555
            node = StaticTuple.from_sequence(node)
 
556
        if type(node) is StaticTuple:
469
557
            return node
470
558
        else:
471
559
            return node._key
472
560
 
473
561
    def unmap(self, key, check_remap=True):
474
562
        """remove key from the map."""
 
563
        key = StaticTuple.from_sequence(key)
475
564
        self._ensure_root()
476
565
        if type(self._root_node) is InternalNode:
477
566
            unmapped = self._root_node.unmap(self._store, key,
491
580
 
492
581
        :return: The key of the root node.
493
582
        """
494
 
        if type(self._root_node) is tuple:
 
583
        if type(self._root_node) is StaticTuple:
495
584
            # Already saved.
496
585
            return self._root_node
497
586
        keys = list(self._root_node.serialise(self._store))
505
594
        adding the header bytes, and without prefix compression.
506
595
    """
507
596
 
 
597
    __slots__ = ('_key', '_len', '_maximum_size', '_key_width',
 
598
                 '_raw_size', '_items', '_search_prefix', '_search_key_func'
 
599
                )
 
600
 
508
601
    def __init__(self, key_width=1):
509
602
        """Create a node.
510
603
 
599
692
        the key/value pairs.
600
693
    """
601
694
 
 
695
    __slots__ = ('_common_serialised_prefix',)
 
696
 
602
697
    def __init__(self, search_key_func=None):
603
698
        Node.__init__(self)
604
699
        # All of the keys in this leaf node share this common prefix
605
700
        self._common_serialised_prefix = None
606
 
        self._serialise_key = '\x00'.join
607
701
        if search_key_func is None:
608
702
            self._search_key_func = _search_key_plain
609
703
        else:
647
741
        :param bytes: The bytes of the node.
648
742
        :param key: The key that the serialised node has.
649
743
        """
 
744
        key = static_tuple.expect_static_tuple(key)
650
745
        return _deserialise_leaf_node(bytes, key,
651
746
                                      search_key_func=search_key_func)
652
747
 
764
859
                result[prefix] = node
765
860
            else:
766
861
                node = result[prefix]
767
 
            node.map(store, key, value)
 
862
            sub_prefix, node_details = node.map(store, key, value)
 
863
            if len(node_details) > 1:
 
864
                if prefix != sub_prefix:
 
865
                    # This node has been split and is now found via a different
 
866
                    # path
 
867
                    result.pop(prefix)
 
868
                new_node = InternalNode(sub_prefix,
 
869
                    search_key_func=self._search_key_func)
 
870
                new_node.set_maximum_size(self._maximum_size)
 
871
                new_node._key_width = self._key_width
 
872
                for split, node in node_details:
 
873
                    new_node.add_node(split, node)
 
874
                result[prefix] = new_node
768
875
        return common_prefix, result.items()
769
876
 
770
877
    def map(self, store, key, value):
780
887
                raise AssertionError('%r must be known' % self._search_prefix)
781
888
            return self._search_prefix, [("", self)]
782
889
 
 
890
    _serialise_key = '\x00'.join
 
891
 
783
892
    def serialise(self, store):
784
893
        """Serialise the LeafNode to store.
785
894
 
810
919
            lines.append(serialized[prefix_len:])
811
920
            lines.extend(value_lines)
812
921
        sha1, _, _ = store.add_lines((None,), (), lines)
813
 
        self._key = ("sha1:" + sha1,)
 
922
        self._key = StaticTuple("sha1:" + sha1,).intern()
814
923
        bytes = ''.join(lines)
815
924
        if len(bytes) != self._current_size():
816
925
            raise AssertionError('Invalid _current_size')
817
 
        _page_cache.add(self._key, bytes)
 
926
        _get_cache().add(self._key, bytes)
818
927
        return [self._key]
819
928
 
820
929
    def refs(self):
884
993
        LeafNode or InternalNode.
885
994
    """
886
995
 
 
996
    __slots__ = ('_node_width',)
 
997
 
887
998
    def __init__(self, prefix='', search_key_func=None):
888
999
        Node.__init__(self)
889
1000
        # The size of an internalnode with default values and no children.
931
1042
        :param key: The key that the serialised node has.
932
1043
        :return: An InternalNode instance.
933
1044
        """
 
1045
        key = static_tuple.expect_static_tuple(key)
934
1046
        return _deserialise_internal_node(bytes, key,
935
1047
                                          search_key_func=search_key_func)
936
1048
 
961
1073
            # for whatever we are missing
962
1074
            shortcut = True
963
1075
            for prefix, node in self._items.iteritems():
964
 
                if node.__class__ is tuple:
 
1076
                if node.__class__ is StaticTuple:
965
1077
                    keys[node] = (prefix, None)
966
1078
                else:
967
1079
                    yield node, None
996
1108
                    # A given key can only match 1 child node, if it isn't
997
1109
                    # there, then we can just return nothing
998
1110
                    return
999
 
                if node.__class__ is tuple:
 
1111
                if node.__class__ is StaticTuple:
1000
1112
                    keys[node] = (search_prefix, [key])
1001
1113
                else:
1002
1114
                    # This is loaded, and the only thing that can match,
1029
1141
                        # We can ignore this one
1030
1142
                        continue
1031
1143
                    node_key_filter = prefix_to_keys[search_prefix]
1032
 
                    if node.__class__ is tuple:
 
1144
                    if node.__class__ is StaticTuple:
1033
1145
                        keys[node] = (search_prefix, node_key_filter)
1034
1146
                    else:
1035
1147
                        yield node, node_key_filter
1044
1156
                        if sub_prefix in length_filter:
1045
1157
                            node_key_filter.extend(prefix_to_keys[sub_prefix])
1046
1158
                    if node_key_filter: # this key matched something, yield it
1047
 
                        if node.__class__ is tuple:
 
1159
                        if node.__class__ is StaticTuple:
1048
1160
                            keys[node] = (prefix, node_key_filter)
1049
1161
                        else:
1050
1162
                            yield node, node_key_filter
1053
1165
            found_keys = set()
1054
1166
            for key in keys:
1055
1167
                try:
1056
 
                    bytes = _page_cache[key]
 
1168
                    bytes = _get_cache()[key]
1057
1169
                except KeyError:
1058
1170
                    continue
1059
1171
                else:
1084
1196
                    prefix, node_key_filter = keys[record.key]
1085
1197
                    node_and_filters.append((node, node_key_filter))
1086
1198
                    self._items[prefix] = node
1087
 
                    _page_cache.add(record.key, bytes)
 
1199
                    _get_cache().add(record.key, bytes)
1088
1200
                for info in node_and_filters:
1089
1201
                    yield info
1090
1202
 
1182
1294
        :return: An iterable of the keys inserted by this operation.
1183
1295
        """
1184
1296
        for node in self._items.itervalues():
1185
 
            if type(node) is tuple:
 
1297
            if type(node) is StaticTuple:
1186
1298
                # Never deserialised.
1187
1299
                continue
1188
1300
            if node._key is not None:
1199
1311
        lines.append('%s\n' % (self._search_prefix,))
1200
1312
        prefix_len = len(self._search_prefix)
1201
1313
        for prefix, node in sorted(self._items.items()):
1202
 
            if type(node) is tuple:
 
1314
            if type(node) is StaticTuple:
1203
1315
                key = node[0]
1204
1316
            else:
1205
1317
                key = node._key[0]
1209
1321
                    % (serialised, self._search_prefix))
1210
1322
            lines.append(serialised[prefix_len:])
1211
1323
        sha1, _, _ = store.add_lines((None,), (), lines)
1212
 
        self._key = ("sha1:" + sha1,)
1213
 
        _page_cache.add(self._key, ''.join(lines))
 
1324
        self._key = StaticTuple("sha1:" + sha1,).intern()
 
1325
        _get_cache().add(self._key, ''.join(lines))
1214
1326
        yield self._key
1215
1327
 
1216
1328
    def _search_key(self, key):
1244
1356
            raise AssertionError("unserialised nodes have no refs.")
1245
1357
        refs = []
1246
1358
        for value in self._items.itervalues():
1247
 
            if type(value) is tuple:
 
1359
            if type(value) is StaticTuple:
1248
1360
                refs.append(value)
1249
1361
            else:
1250
1362
                refs.append(value.key())
1351
1463
    return node
1352
1464
 
1353
1465
 
1354
 
def _find_children_info(store, interesting_keys, uninteresting_keys, pb):
1355
 
    """Read the associated records, and determine what is interesting."""
1356
 
    uninteresting_keys = set(uninteresting_keys)
1357
 
    chks_to_read = uninteresting_keys.union(interesting_keys)
1358
 
    next_uninteresting = set()
1359
 
    next_interesting = set()
1360
 
    uninteresting_items = set()
1361
 
    interesting_items = set()
1362
 
    interesting_to_yield = []
1363
 
    for record in store.get_record_stream(chks_to_read, 'unordered', True):
1364
 
        # records_read.add(record.key())
1365
 
        if pb is not None:
1366
 
            pb.tick()
1367
 
        bytes = record.get_bytes_as('fulltext')
1368
 
        # We don't care about search_key_func for this code, because we only
1369
 
        # care about external references.
1370
 
        node = _deserialise(bytes, record.key, search_key_func=None)
1371
 
        if record.key in uninteresting_keys:
1372
 
            if type(node) is InternalNode:
1373
 
                next_uninteresting.update(node.refs())
1374
 
            else:
1375
 
                # We know we are at a LeafNode, so we can pass None for the
1376
 
                # store
1377
 
                uninteresting_items.update(node.iteritems(None))
1378
 
        else:
1379
 
            interesting_to_yield.append(record.key)
1380
 
            if type(node) is InternalNode:
1381
 
                next_interesting.update(node.refs())
1382
 
            else:
1383
 
                interesting_items.update(node.iteritems(None))
1384
 
    return (next_uninteresting, uninteresting_items,
1385
 
            next_interesting, interesting_to_yield, interesting_items)
1386
 
 
1387
 
 
1388
 
def _find_all_uninteresting(store, interesting_root_keys,
1389
 
                            uninteresting_root_keys, pb):
1390
 
    """Determine the full set of uninteresting keys."""
1391
 
    # What about duplicates between interesting_root_keys and
1392
 
    # uninteresting_root_keys?
1393
 
    if not uninteresting_root_keys:
1394
 
        # Shortcut case. We know there is nothing uninteresting to filter out
1395
 
        # So we just let the rest of the algorithm do the work
1396
 
        # We know there is nothing uninteresting, and we didn't have to read
1397
 
        # any interesting records yet.
1398
 
        return (set(), set(), set(interesting_root_keys), [], set())
1399
 
    all_uninteresting_chks = set(uninteresting_root_keys)
1400
 
    all_uninteresting_items = set()
1401
 
 
1402
 
    # First step, find the direct children of both the interesting and
1403
 
    # uninteresting set
1404
 
    (uninteresting_keys, uninteresting_items,
1405
 
     interesting_keys, interesting_to_yield,
1406
 
     interesting_items) = _find_children_info(store, interesting_root_keys,
1407
 
                                              uninteresting_root_keys,
1408
 
                                              pb=pb)
1409
 
    all_uninteresting_chks.update(uninteresting_keys)
1410
 
    all_uninteresting_items.update(uninteresting_items)
1411
 
    del uninteresting_items
1412
 
    # Note: Exact matches between interesting and uninteresting do not need
1413
 
    #       to be search further. Non-exact matches need to be searched in case
1414
 
    #       there is a future exact-match
1415
 
    uninteresting_keys.difference_update(interesting_keys)
1416
 
 
1417
 
    # Second, find the full set of uninteresting bits reachable by the
1418
 
    # uninteresting roots
1419
 
    chks_to_read = uninteresting_keys
1420
 
    while chks_to_read:
1421
 
        next_chks = set()
1422
 
        for record in store.get_record_stream(chks_to_read, 'unordered', False):
1423
 
            # TODO: Handle 'absent'
1424
 
            if pb is not None:
1425
 
                pb.tick()
 
1466
class CHKMapDifference(object):
 
1467
    """Iterate the stored pages and key,value pairs for (new - old).
 
1468
 
 
1469
    This class provides a generator over the stored CHK pages and the
 
1470
    (key, value) pairs that are in any of the new maps and not in any of the
 
1471
    old maps.
 
1472
 
 
1473
    Note that it may yield chk pages that are common (especially root nodes),
 
1474
    but it won't yield (key,value) pairs that are common.
 
1475
    """
 
1476
 
 
1477
    def __init__(self, store, new_root_keys, old_root_keys,
 
1478
                 search_key_func, pb=None):
 
1479
        # TODO: Should we add a StaticTuple barrier here? It would be nice to
 
1480
        #       force callers to use StaticTuple, because there will often be
 
1481
        #       lots of keys passed in here. And even if we cast it locally,
 
1482
        #       that just meanst that we will have *both* a StaticTuple and a
 
1483
        #       tuple() in memory, referring to the same object. (so a net
 
1484
        #       increase in memory, not a decrease.)
 
1485
        self._store = store
 
1486
        self._new_root_keys = new_root_keys
 
1487
        self._old_root_keys = old_root_keys
 
1488
        self._pb = pb
 
1489
        # All uninteresting chks that we have seen. By the time they are added
 
1490
        # here, they should be either fully ignored, or queued up for
 
1491
        # processing
 
1492
        # TODO: This might grow to a large size if there are lots of merge
 
1493
        #       parents, etc. However, it probably doesn't scale to O(history)
 
1494
        #       like _processed_new_refs does.
 
1495
        self._all_old_chks = set(self._old_root_keys)
 
1496
        # All items that we have seen from the old_root_keys
 
1497
        self._all_old_items = set()
 
1498
        # These are interesting items which were either read, or already in the
 
1499
        # interesting queue (so we don't need to walk them again)
 
1500
        # TODO: processed_new_refs becomes O(all_chks), consider switching to
 
1501
        #       SimpleSet here.
 
1502
        self._processed_new_refs = set()
 
1503
        self._search_key_func = search_key_func
 
1504
 
 
1505
        # The uninteresting and interesting nodes to be searched
 
1506
        self._old_queue = []
 
1507
        self._new_queue = []
 
1508
        # Holds the (key, value) items found when processing the root nodes,
 
1509
        # waiting for the uninteresting nodes to be walked
 
1510
        self._new_item_queue = []
 
1511
        self._state = None
 
1512
 
 
1513
    def _read_nodes_from_store(self, keys):
 
1514
        # We chose not to use _get_cache(), because we think in
 
1515
        # terms of records to be yielded. Also, we expect to touch each page
 
1516
        # only 1 time during this code. (We may want to evaluate saving the
 
1517
        # raw bytes into the page cache, which would allow a working tree
 
1518
        # update after the fetch to not have to read the bytes again.)
 
1519
        as_st = StaticTuple.from_sequence
 
1520
        stream = self._store.get_record_stream(keys, 'unordered', True)
 
1521
        for record in stream:
 
1522
            if self._pb is not None:
 
1523
                self._pb.tick()
 
1524
            if record.storage_kind == 'absent':
 
1525
                raise errors.NoSuchRevision(self._store, record.key)
1426
1526
            bytes = record.get_bytes_as('fulltext')
1427
 
            # We don't care about search_key_func for this code, because we
1428
 
            # only care about external references.
1429
 
            node = _deserialise(bytes, record.key, search_key_func=None)
 
1527
            node = _deserialise(bytes, record.key,
 
1528
                                search_key_func=self._search_key_func)
1430
1529
            if type(node) is InternalNode:
1431
 
                # uninteresting_prefix_chks.update(node._items.iteritems())
1432
 
                chks = node._items.values()
1433
 
                # TODO: We remove the entries that are already in
1434
 
                #       uninteresting_chks ?
1435
 
                next_chks.update(chks)
1436
 
                all_uninteresting_chks.update(chks)
 
1530
                # Note we don't have to do node.refs() because we know that
 
1531
                # there are no children that have been pushed into this node
 
1532
                # Note: Using as_st() here seemed to save 1.2MB, which would
 
1533
                #       indicate that we keep 100k prefix_refs around while
 
1534
                #       processing. They *should* be shorter lived than that...
 
1535
                #       It does cost us ~10s of processing time
 
1536
                #prefix_refs = [as_st(item) for item in node._items.iteritems()]
 
1537
                prefix_refs = node._items.items()
 
1538
                items = []
1437
1539
            else:
1438
 
                all_uninteresting_items.update(node._items.iteritems())
1439
 
        chks_to_read = next_chks
1440
 
    return (all_uninteresting_chks, all_uninteresting_items,
1441
 
            interesting_keys, interesting_to_yield, interesting_items)
 
1540
                prefix_refs = []
 
1541
                # Note: We don't use a StaticTuple here. Profiling showed a
 
1542
                #       minor memory improvement (0.8MB out of 335MB peak 0.2%)
 
1543
                #       But a significant slowdown (15s / 145s, or 10%)
 
1544
                items = node._items.items()
 
1545
            yield record, node, prefix_refs, items
 
1546
 
 
1547
    def _read_old_roots(self):
 
1548
        old_chks_to_enqueue = []
 
1549
        all_old_chks = self._all_old_chks
 
1550
        for record, node, prefix_refs, items in \
 
1551
                self._read_nodes_from_store(self._old_root_keys):
 
1552
            # Uninteresting node
 
1553
            prefix_refs = [p_r for p_r in prefix_refs
 
1554
                                if p_r[1] not in all_old_chks]
 
1555
            new_refs = [p_r[1] for p_r in prefix_refs]
 
1556
            all_old_chks.update(new_refs)
 
1557
            # TODO: This might be a good time to turn items into StaticTuple
 
1558
            #       instances and possibly intern them. However, this does not
 
1559
            #       impact 'initial branch' performance, so I'm not worrying
 
1560
            #       about this yet
 
1561
            self._all_old_items.update(items)
 
1562
            # Queue up the uninteresting references
 
1563
            # Don't actually put them in the 'to-read' queue until we have
 
1564
            # finished checking the interesting references
 
1565
            old_chks_to_enqueue.extend(prefix_refs)
 
1566
        return old_chks_to_enqueue
 
1567
 
 
1568
    def _enqueue_old(self, new_prefixes, old_chks_to_enqueue):
 
1569
        # At this point, we have read all the uninteresting and interesting
 
1570
        # items, so we can queue up the uninteresting stuff, knowing that we've
 
1571
        # handled the interesting ones
 
1572
        for prefix, ref in old_chks_to_enqueue:
 
1573
            not_interesting = True
 
1574
            for i in xrange(len(prefix), 0, -1):
 
1575
                if prefix[:i] in new_prefixes:
 
1576
                    not_interesting = False
 
1577
                    break
 
1578
            if not_interesting:
 
1579
                # This prefix is not part of the remaining 'interesting set'
 
1580
                continue
 
1581
            self._old_queue.append(ref)
 
1582
 
 
1583
    def _read_all_roots(self):
 
1584
        """Read the root pages.
 
1585
 
 
1586
        This is structured as a generator, so that the root records can be
 
1587
        yielded up to whoever needs them without any buffering.
 
1588
        """
 
1589
        # This is the bootstrap phase
 
1590
        if not self._old_root_keys:
 
1591
            # With no old_root_keys we can just shortcut and be ready
 
1592
            # for _flush_new_queue
 
1593
            self._new_queue = list(self._new_root_keys)
 
1594
            return
 
1595
        old_chks_to_enqueue = self._read_old_roots()
 
1596
        # filter out any root keys that are already known to be uninteresting
 
1597
        new_keys = set(self._new_root_keys).difference(self._all_old_chks)
 
1598
        # These are prefixes that are present in new_keys that we are
 
1599
        # thinking to yield
 
1600
        new_prefixes = set()
 
1601
        # We are about to yield all of these, so we don't want them getting
 
1602
        # added a second time
 
1603
        processed_new_refs = self._processed_new_refs
 
1604
        processed_new_refs.update(new_keys)
 
1605
        for record, node, prefix_refs, items in \
 
1606
                self._read_nodes_from_store(new_keys):
 
1607
            # At this level, we now know all the uninteresting references
 
1608
            # So we filter and queue up whatever is remaining
 
1609
            prefix_refs = [p_r for p_r in prefix_refs
 
1610
                           if p_r[1] not in self._all_old_chks
 
1611
                              and p_r[1] not in processed_new_refs]
 
1612
            refs = [p_r[1] for p_r in prefix_refs]
 
1613
            new_prefixes.update([p_r[0] for p_r in prefix_refs])
 
1614
            self._new_queue.extend(refs)
 
1615
            # TODO: We can potentially get multiple items here, however the
 
1616
            #       current design allows for this, as callers will do the work
 
1617
            #       to make the results unique. We might profile whether we
 
1618
            #       gain anything by ensuring unique return values for items
 
1619
            # TODO: This might be a good time to cast to StaticTuple, as
 
1620
            #       self._new_item_queue will hold the contents of multiple
 
1621
            #       records for an extended lifetime
 
1622
            new_items = [item for item in items
 
1623
                               if item not in self._all_old_items]
 
1624
            self._new_item_queue.extend(new_items)
 
1625
            new_prefixes.update([self._search_key_func(item[0])
 
1626
                                 for item in new_items])
 
1627
            processed_new_refs.update(refs)
 
1628
            yield record
 
1629
        # For new_prefixes we have the full length prefixes queued up.
 
1630
        # However, we also need possible prefixes. (If we have a known ref to
 
1631
        # 'ab', then we also need to include 'a'.) So expand the
 
1632
        # new_prefixes to include all shorter prefixes
 
1633
        for prefix in list(new_prefixes):
 
1634
            new_prefixes.update([prefix[:i] for i in xrange(1, len(prefix))])
 
1635
        self._enqueue_old(new_prefixes, old_chks_to_enqueue)
 
1636
 
 
1637
    def _flush_new_queue(self):
 
1638
        # No need to maintain the heap invariant anymore, just pull things out
 
1639
        # and process them
 
1640
        refs = set(self._new_queue)
 
1641
        self._new_queue = []
 
1642
        # First pass, flush all interesting items and convert to using direct refs
 
1643
        all_old_chks = self._all_old_chks
 
1644
        processed_new_refs = self._processed_new_refs
 
1645
        all_old_items = self._all_old_items
 
1646
        new_items = [item for item in self._new_item_queue
 
1647
                           if item not in all_old_items]
 
1648
        self._new_item_queue = []
 
1649
        if new_items:
 
1650
            yield None, new_items
 
1651
        refs = refs.difference(all_old_chks)
 
1652
        processed_new_refs.update(refs)
 
1653
        while refs:
 
1654
            # TODO: Using a SimpleSet for self._processed_new_refs and
 
1655
            #       saved as much as 10MB of peak memory. However, it requires
 
1656
            #       implementing a non-pyrex version.
 
1657
            next_refs = set()
 
1658
            next_refs_update = next_refs.update
 
1659
            # Inlining _read_nodes_from_store improves 'bzr branch bzr.dev'
 
1660
            # from 1m54s to 1m51s. Consider it.
 
1661
            for record, _, p_refs, items in self._read_nodes_from_store(refs):
 
1662
                if all_old_items:
 
1663
                    # using the 'if' check saves about 145s => 141s, when
 
1664
                    # streaming initial branch of Launchpad data.
 
1665
                    items = [item for item in items
 
1666
                             if item not in all_old_items]
 
1667
                yield record, items
 
1668
                next_refs_update([p_r[1] for p_r in p_refs])
 
1669
                del p_refs
 
1670
            # set1.difference(set/dict) walks all of set1, and checks if it
 
1671
            # exists in 'other'.
 
1672
            # set1.difference(iterable) walks all of iterable, and does a
 
1673
            # 'difference_update' on a clone of set1. Pick wisely based on the
 
1674
            # expected sizes of objects.
 
1675
            # in our case it is expected that 'new_refs' will always be quite
 
1676
            # small.
 
1677
            next_refs = next_refs.difference(all_old_chks)
 
1678
            next_refs = next_refs.difference(processed_new_refs)
 
1679
            processed_new_refs.update(next_refs)
 
1680
            refs = next_refs
 
1681
 
 
1682
    def _process_next_old(self):
 
1683
        # Since we don't filter uninteresting any further than during
 
1684
        # _read_all_roots, process the whole queue in a single pass.
 
1685
        refs = self._old_queue
 
1686
        self._old_queue = []
 
1687
        all_old_chks = self._all_old_chks
 
1688
        for record, _, prefix_refs, items in self._read_nodes_from_store(refs):
 
1689
            # TODO: Use StaticTuple here?
 
1690
            self._all_old_items.update(items)
 
1691
            refs = [r for _,r in prefix_refs if r not in all_old_chks]
 
1692
            self._old_queue.extend(refs)
 
1693
            all_old_chks.update(refs)
 
1694
 
 
1695
    def _process_queues(self):
 
1696
        while self._old_queue:
 
1697
            self._process_next_old()
 
1698
        return self._flush_new_queue()
 
1699
 
 
1700
    def process(self):
 
1701
        for record in self._read_all_roots():
 
1702
            yield record, []
 
1703
        for record, items in self._process_queues():
 
1704
            yield record, items
1442
1705
 
1443
1706
 
1444
1707
def iter_interesting_nodes(store, interesting_root_keys,
1455
1718
    :return: Yield
1456
1719
        (interesting record, {interesting key:values})
1457
1720
    """
1458
 
    # TODO: consider that it may be more memory efficient to use the 20-byte
1459
 
    #       sha1 string, rather than tuples of hexidecimal sha1 strings.
1460
 
    # TODO: Try to factor out a lot of the get_record_stream() calls into a
1461
 
    #       helper function similar to _read_bytes. This function should be
1462
 
    #       able to use nodes from the _page_cache as well as actually
1463
 
    #       requesting bytes from the store.
1464
 
 
1465
 
    (all_uninteresting_chks, all_uninteresting_items, interesting_keys,
1466
 
     interesting_to_yield, interesting_items) = _find_all_uninteresting(store,
1467
 
        interesting_root_keys, uninteresting_root_keys, pb)
1468
 
 
1469
 
    # Now that we know everything uninteresting, we can yield information from
1470
 
    # our first request
1471
 
    interesting_items.difference_update(all_uninteresting_items)
1472
 
    interesting_to_yield = set(interesting_to_yield) - all_uninteresting_chks
1473
 
    if interesting_items:
1474
 
        yield None, interesting_items
1475
 
    if interesting_to_yield:
1476
 
        # We request these records again, rather than buffering the root
1477
 
        # records, most likely they are still in the _group_cache anyway.
1478
 
        for record in store.get_record_stream(interesting_to_yield,
1479
 
                                              'unordered', False):
1480
 
            yield record, []
1481
 
    all_uninteresting_chks.update(interesting_to_yield)
1482
 
    interesting_keys.difference_update(all_uninteresting_chks)
1483
 
 
1484
 
    chks_to_read = interesting_keys
1485
 
    counter = 0
1486
 
    while chks_to_read:
1487
 
        next_chks = set()
1488
 
        for record in store.get_record_stream(chks_to_read, 'unordered', False):
1489
 
            counter += 1
1490
 
            if pb is not None:
1491
 
                pb.update('find chk pages', counter)
1492
 
            # TODO: Handle 'absent'?
1493
 
            bytes = record.get_bytes_as('fulltext')
1494
 
            # We don't care about search_key_func for this code, because we
1495
 
            # only care about external references.
1496
 
            node = _deserialise(bytes, record.key, search_key_func=None)
1497
 
            if type(node) is InternalNode:
1498
 
                # all_uninteresting_chks grows large, as it lists all nodes we
1499
 
                # don't want to process (including already seen interesting
1500
 
                # nodes).
1501
 
                # small.difference_update(large) scales O(large), but
1502
 
                # small.difference(large) scales O(small).
1503
 
                # Also, we know we just _deserialised this node, so we can
1504
 
                # access the dict directly.
1505
 
                chks = set(node._items.itervalues()).difference(
1506
 
                            all_uninteresting_chks)
1507
 
                # Is set() and .difference_update better than:
1508
 
                # chks = [chk for chk in node.refs()
1509
 
                #              if chk not in all_uninteresting_chks]
1510
 
                next_chks.update(chks)
1511
 
                # These are now uninteresting everywhere else
1512
 
                all_uninteresting_chks.update(chks)
1513
 
                interesting_items = []
1514
 
            else:
1515
 
                interesting_items = [item for item in node._items.iteritems()
1516
 
                                     if item not in all_uninteresting_items]
1517
 
                # TODO: Do we need to filter out items that we have already
1518
 
                #       seen on other pages? We don't really want to buffer the
1519
 
                #       whole thing, but it does mean that callers need to
1520
 
                #       understand they may get duplicate values.
1521
 
                # all_uninteresting_items.update(interesting_items)
1522
 
            yield record, interesting_items
1523
 
        chks_to_read = next_chks
 
1721
    iterator = CHKMapDifference(store, interesting_root_keys,
 
1722
                                uninteresting_root_keys,
 
1723
                                search_key_func=store._search_key_func,
 
1724
                                pb=pb)
 
1725
    return iterator.process()
1524
1726
 
1525
1727
 
1526
1728
try:
1527
1729
    from bzrlib._chk_map_pyx import (
 
1730
        _bytes_to_text_key,
1528
1731
        _search_key_16,
1529
1732
        _search_key_255,
1530
1733
        _deserialise_leaf_node,
1531
1734
        _deserialise_internal_node,
1532
1735
        )
1533
 
except ImportError:
 
1736
except ImportError, e:
 
1737
    osutils.failed_to_load_extension(e)
1534
1738
    from bzrlib._chk_map_py import (
 
1739
        _bytes_to_text_key,
1535
1740
        _search_key_16,
1536
1741
        _search_key_255,
1537
1742
        _deserialise_leaf_node,
1539
1744
        )
1540
1745
search_key_registry.register('hash-16-way', _search_key_16)
1541
1746
search_key_registry.register('hash-255-way', _search_key_255)
 
1747
 
 
1748
 
 
1749
def _check_key(key):
 
1750
    """Helper function to assert that a key is properly formatted.
 
1751
 
 
1752
    This generally shouldn't be used in production code, but it can be helpful
 
1753
    to debug problems.
 
1754
    """
 
1755
    if type(key) is not StaticTuple:
 
1756
        raise TypeError('key %r is not StaticTuple but %s' % (key, type(key)))
 
1757
    if len(key) != 1:
 
1758
        raise ValueError('key %r should have length 1, not %d' % (key, len(key),))
 
1759
    if type(key[0]) is not str:
 
1760
        raise TypeError('key %r should hold a str, not %r'
 
1761
                        % (key, type(key[0])))
 
1762
    if not key[0].startswith('sha1:'):
 
1763
        raise ValueError('key %r should point to a sha1:' % (key,))
 
1764
 
 
1765