~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/chk_map.py

  • Committer: Gordon Tyler
  • Date: 2010-01-14 15:24:04 UTC
  • mto: (5037.3.1 integration)
  • mto: This revision was merged to the branch mainline in revision 5046.
  • Revision ID: gordon@doxxx.net-20100114152404-d64ik2afl96lcml0
Reverted changes to test_rules since the original should work now.

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2008 Canonical Ltd
 
1
# Copyright (C) 2008, 2009 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
42
41
 
43
42
from bzrlib import lazy_import
44
43
lazy_import.lazy_import(globals(), """
45
 
from bzrlib import versionedfile
 
44
from bzrlib import (
 
45
    errors,
 
46
    versionedfile,
 
47
    )
46
48
""")
47
49
from bzrlib import (
48
 
    errors,
49
50
    lru_cache,
50
51
    osutils,
51
52
    registry,
 
53
    static_tuple,
52
54
    trace,
53
55
    )
 
56
from bzrlib.static_tuple import StaticTuple
54
57
 
55
58
# approx 4MB
56
59
# If each line is 50 bytes, and you have 255 internal pages, with 255-way fan
59
62
# We are caching bytes so len(value) is perfectly accurate
60
63
_page_cache = lru_cache.LRUSizeCache(_PAGE_CACHE_SIZE)
61
64
 
 
65
def clear_cache():
 
66
    _page_cache.clear()
 
67
 
62
68
# If a ChildNode falls below this many bytes, we check for a remap
63
69
_INTERESTING_NEW_SIZE = 50
64
70
# If a ChildNode shrinks by more than this amount, we check for a remap
79
85
class CHKMap(object):
80
86
    """A persistent map from string to string backed by a CHK store."""
81
87
 
 
88
    __slots__ = ('_store', '_root_node', '_search_key_func')
 
89
 
82
90
    def __init__(self, store, root_key, search_key_func=None):
83
91
        """Create a CHKMap object.
84
92
 
107
115
            of old_key is removed.
108
116
        """
109
117
        delete_count = 0
 
118
        # Check preconditions first.
 
119
        as_st = StaticTuple.from_sequence
 
120
        new_items = set([as_st(key) for (old, key, value) in delta
 
121
                         if key is not None and old is None])
 
122
        existing_new = list(self.iteritems(key_filter=new_items))
 
123
        if existing_new:
 
124
            raise errors.InconsistentDeltaDelta(delta,
 
125
                "New items are already in the map %r." % existing_new)
 
126
        # Now apply changes.
110
127
        for old, new, value in delta:
111
128
            if old is not None and old != new:
112
129
                self.unmap(old, check_remap=False)
121
138
 
122
139
    def _ensure_root(self):
123
140
        """Ensure that the root node is an object not a key."""
124
 
        if type(self._root_node) is tuple:
 
141
        if type(self._root_node) is StaticTuple:
125
142
            # Demand-load the root
126
143
            self._root_node = self._get_node(self._root_node)
127
144
 
135
152
        :param node: A tuple key or node object.
136
153
        :return: A node object.
137
154
        """
138
 
        if type(node) is tuple:
 
155
        if type(node) is StaticTuple:
139
156
            bytes = self._read_bytes(node)
140
157
            return _deserialise(bytes, node,
141
158
                search_key_func=self._search_key_func)
182
199
            for key, value in sorted(node._items.iteritems()):
183
200
                # Don't use prefix nor indent here to line up when used in
184
201
                # tests in conjunction with assertEqualDiff
185
 
                result.append('      %r %r' % (key, value))
 
202
                result.append('      %r %r' % (tuple(key), value))
186
203
        return result
187
204
 
188
205
    @classmethod
203
220
            multiple pages.
204
221
        :return: The root chk of the resulting CHKMap.
205
222
        """
206
 
        result = CHKMap(store, None, search_key_func=search_key_func)
 
223
        root_key = klass._create_directly(store, initial_value,
 
224
            maximum_size=maximum_size, key_width=key_width,
 
225
            search_key_func=search_key_func)
 
226
        if type(root_key) is not StaticTuple:
 
227
            raise AssertionError('we got a %s instead of a StaticTuple'
 
228
                                 % (type(root_key),))
 
229
        return root_key
 
230
 
 
231
    @classmethod
 
232
    def _create_via_map(klass, store, initial_value, maximum_size=0,
 
233
                        key_width=1, search_key_func=None):
 
234
        result = klass(store, None, search_key_func=search_key_func)
207
235
        result._root_node.set_maximum_size(maximum_size)
208
236
        result._root_node._key_width = key_width
209
237
        delta = []
210
238
        for key, value in initial_value.items():
211
239
            delta.append((None, key, value))
212
 
        return result.apply_delta(delta)
 
240
        root_key = result.apply_delta(delta)
 
241
        return root_key
 
242
 
 
243
    @classmethod
 
244
    def _create_directly(klass, store, initial_value, maximum_size=0,
 
245
                         key_width=1, search_key_func=None):
 
246
        node = LeafNode(search_key_func=search_key_func)
 
247
        node.set_maximum_size(maximum_size)
 
248
        node._key_width = key_width
 
249
        as_st = StaticTuple.from_sequence
 
250
        node._items = dict([(as_st(key), val) for key, val
 
251
                                               in initial_value.iteritems()])
 
252
        node._raw_size = sum([node._key_value_len(key, value)
 
253
                              for key,value in node._items.iteritems()])
 
254
        node._len = len(node._items)
 
255
        node._compute_search_prefix()
 
256
        node._compute_serialised_prefix()
 
257
        if (node._len > 1
 
258
            and maximum_size
 
259
            and node._current_size() > maximum_size):
 
260
            prefix, node_details = node._split(store)
 
261
            if len(node_details) == 1:
 
262
                raise AssertionError('Failed to split using node._split')
 
263
            node = InternalNode(prefix, search_key_func=search_key_func)
 
264
            node.set_maximum_size(maximum_size)
 
265
            node._key_width = key_width
 
266
            for split, subnode in node_details:
 
267
                node.add_node(split, subnode)
 
268
        keys = list(node.serialise(store))
 
269
        return keys[-1]
213
270
 
214
271
    def iter_changes(self, basis):
215
272
        """Iterate over the changes between basis and self.
435
492
    def iteritems(self, key_filter=None):
436
493
        """Iterate over the entire CHKMap's contents."""
437
494
        self._ensure_root()
 
495
        if key_filter is not None:
 
496
            as_st = StaticTuple.from_sequence
 
497
            key_filter = [as_st(key) for key in key_filter]
438
498
        return self._root_node.iteritems(self._store, key_filter=key_filter)
439
499
 
440
500
    def key(self):
441
501
        """Return the key for this map."""
442
 
        if type(self._root_node) is tuple:
 
502
        if type(self._root_node) is StaticTuple:
443
503
            return self._root_node
444
504
        else:
445
505
            return self._root_node._key
449
509
        return len(self._root_node)
450
510
 
451
511
    def map(self, key, value):
452
 
        """Map a key tuple to value."""
 
512
        """Map a key tuple to value.
 
513
        
 
514
        :param key: A key to map.
 
515
        :param value: The value to assign to key.
 
516
        """
 
517
        key = StaticTuple.from_sequence(key)
453
518
        # Need a root object.
454
519
        self._ensure_root()
455
520
        prefix, node_details = self._root_node.map(self._store, key, value)
466
531
    def _node_key(self, node):
467
532
        """Get the key for a node whether it's a tuple or node."""
468
533
        if type(node) is tuple:
 
534
            node = StaticTuple.from_sequence(node)
 
535
        if type(node) is StaticTuple:
469
536
            return node
470
537
        else:
471
538
            return node._key
472
539
 
473
540
    def unmap(self, key, check_remap=True):
474
541
        """remove key from the map."""
 
542
        key = StaticTuple.from_sequence(key)
475
543
        self._ensure_root()
476
544
        if type(self._root_node) is InternalNode:
477
545
            unmapped = self._root_node.unmap(self._store, key,
491
559
 
492
560
        :return: The key of the root node.
493
561
        """
494
 
        if type(self._root_node) is tuple:
 
562
        if type(self._root_node) is StaticTuple:
495
563
            # Already saved.
496
564
            return self._root_node
497
565
        keys = list(self._root_node.serialise(self._store))
505
573
        adding the header bytes, and without prefix compression.
506
574
    """
507
575
 
 
576
    __slots__ = ('_key', '_len', '_maximum_size', '_key_width',
 
577
                 '_raw_size', '_items', '_search_prefix', '_search_key_func'
 
578
                )
 
579
 
508
580
    def __init__(self, key_width=1):
509
581
        """Create a node.
510
582
 
599
671
        the key/value pairs.
600
672
    """
601
673
 
 
674
    __slots__ = ('_common_serialised_prefix', '_serialise_key')
 
675
 
602
676
    def __init__(self, search_key_func=None):
603
677
        Node.__init__(self)
604
678
        # All of the keys in this leaf node share this common prefix
647
721
        :param bytes: The bytes of the node.
648
722
        :param key: The key that the serialised node has.
649
723
        """
 
724
        key = static_tuple.expect_static_tuple(key)
650
725
        return _deserialise_leaf_node(bytes, key,
651
726
                                      search_key_func=search_key_func)
652
727
 
764
839
                result[prefix] = node
765
840
            else:
766
841
                node = result[prefix]
767
 
            node.map(store, key, value)
 
842
            sub_prefix, node_details = node.map(store, key, value)
 
843
            if len(node_details) > 1:
 
844
                if prefix != sub_prefix:
 
845
                    # This node has been split and is now found via a different
 
846
                    # path
 
847
                    result.pop(prefix)
 
848
                new_node = InternalNode(sub_prefix,
 
849
                    search_key_func=self._search_key_func)
 
850
                new_node.set_maximum_size(self._maximum_size)
 
851
                new_node._key_width = self._key_width
 
852
                for split, node in node_details:
 
853
                    new_node.add_node(split, node)
 
854
                result[prefix] = new_node
768
855
        return common_prefix, result.items()
769
856
 
770
857
    def map(self, store, key, value):
810
897
            lines.append(serialized[prefix_len:])
811
898
            lines.extend(value_lines)
812
899
        sha1, _, _ = store.add_lines((None,), (), lines)
813
 
        self._key = ("sha1:" + sha1,)
 
900
        self._key = StaticTuple("sha1:" + sha1,).intern()
814
901
        bytes = ''.join(lines)
815
902
        if len(bytes) != self._current_size():
816
903
            raise AssertionError('Invalid _current_size')
884
971
        LeafNode or InternalNode.
885
972
    """
886
973
 
 
974
    __slots__ = ('_node_width',)
 
975
 
887
976
    def __init__(self, prefix='', search_key_func=None):
888
977
        Node.__init__(self)
889
978
        # The size of an internalnode with default values and no children.
931
1020
        :param key: The key that the serialised node has.
932
1021
        :return: An InternalNode instance.
933
1022
        """
 
1023
        key = static_tuple.expect_static_tuple(key)
934
1024
        return _deserialise_internal_node(bytes, key,
935
1025
                                          search_key_func=search_key_func)
936
1026
 
961
1051
            # for whatever we are missing
962
1052
            shortcut = True
963
1053
            for prefix, node in self._items.iteritems():
964
 
                if node.__class__ is tuple:
 
1054
                if node.__class__ is StaticTuple:
965
1055
                    keys[node] = (prefix, None)
966
1056
                else:
967
1057
                    yield node, None
996
1086
                    # A given key can only match 1 child node, if it isn't
997
1087
                    # there, then we can just return nothing
998
1088
                    return
999
 
                if node.__class__ is tuple:
 
1089
                if node.__class__ is StaticTuple:
1000
1090
                    keys[node] = (search_prefix, [key])
1001
1091
                else:
1002
1092
                    # This is loaded, and the only thing that can match,
1029
1119
                        # We can ignore this one
1030
1120
                        continue
1031
1121
                    node_key_filter = prefix_to_keys[search_prefix]
1032
 
                    if node.__class__ is tuple:
 
1122
                    if node.__class__ is StaticTuple:
1033
1123
                        keys[node] = (search_prefix, node_key_filter)
1034
1124
                    else:
1035
1125
                        yield node, node_key_filter
1044
1134
                        if sub_prefix in length_filter:
1045
1135
                            node_key_filter.extend(prefix_to_keys[sub_prefix])
1046
1136
                    if node_key_filter: # this key matched something, yield it
1047
 
                        if node.__class__ is tuple:
 
1137
                        if node.__class__ is StaticTuple:
1048
1138
                            keys[node] = (prefix, node_key_filter)
1049
1139
                        else:
1050
1140
                            yield node, node_key_filter
1182
1272
        :return: An iterable of the keys inserted by this operation.
1183
1273
        """
1184
1274
        for node in self._items.itervalues():
1185
 
            if type(node) is tuple:
 
1275
            if type(node) is StaticTuple:
1186
1276
                # Never deserialised.
1187
1277
                continue
1188
1278
            if node._key is not None:
1199
1289
        lines.append('%s\n' % (self._search_prefix,))
1200
1290
        prefix_len = len(self._search_prefix)
1201
1291
        for prefix, node in sorted(self._items.items()):
1202
 
            if type(node) is tuple:
 
1292
            if type(node) is StaticTuple:
1203
1293
                key = node[0]
1204
1294
            else:
1205
1295
                key = node._key[0]
1209
1299
                    % (serialised, self._search_prefix))
1210
1300
            lines.append(serialised[prefix_len:])
1211
1301
        sha1, _, _ = store.add_lines((None,), (), lines)
1212
 
        self._key = ("sha1:" + sha1,)
 
1302
        self._key = StaticTuple("sha1:" + sha1,).intern()
1213
1303
        _page_cache.add(self._key, ''.join(lines))
1214
1304
        yield self._key
1215
1305
 
1244
1334
            raise AssertionError("unserialised nodes have no refs.")
1245
1335
        refs = []
1246
1336
        for value in self._items.itervalues():
1247
 
            if type(value) is tuple:
 
1337
            if type(value) is StaticTuple:
1248
1338
                refs.append(value)
1249
1339
            else:
1250
1340
                refs.append(value.key())
1351
1441
    return node
1352
1442
 
1353
1443
 
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()
 
1444
class CHKMapDifference(object):
 
1445
    """Iterate the stored pages and key,value pairs for (new - old).
 
1446
 
 
1447
    This class provides a generator over the stored CHK pages and the
 
1448
    (key, value) pairs that are in any of the new maps and not in any of the
 
1449
    old maps.
 
1450
 
 
1451
    Note that it may yield chk pages that are common (especially root nodes),
 
1452
    but it won't yield (key,value) pairs that are common.
 
1453
    """
 
1454
 
 
1455
    def __init__(self, store, new_root_keys, old_root_keys,
 
1456
                 search_key_func, pb=None):
 
1457
        # TODO: Should we add a StaticTuple barrier here? It would be nice to
 
1458
        #       force callers to use StaticTuple, because there will often be
 
1459
        #       lots of keys passed in here. And even if we cast it locally,
 
1460
        #       that just meanst that we will have *both* a StaticTuple and a
 
1461
        #       tuple() in memory, referring to the same object. (so a net
 
1462
        #       increase in memory, not a decrease.)
 
1463
        self._store = store
 
1464
        self._new_root_keys = new_root_keys
 
1465
        self._old_root_keys = old_root_keys
 
1466
        self._pb = pb
 
1467
        # All uninteresting chks that we have seen. By the time they are added
 
1468
        # here, they should be either fully ignored, or queued up for
 
1469
        # processing
 
1470
        # TODO: This might grow to a large size if there are lots of merge
 
1471
        #       parents, etc. However, it probably doesn't scale to O(history)
 
1472
        #       like _processed_new_refs does.
 
1473
        self._all_old_chks = set(self._old_root_keys)
 
1474
        # All items that we have seen from the old_root_keys
 
1475
        self._all_old_items = set()
 
1476
        # These are interesting items which were either read, or already in the
 
1477
        # interesting queue (so we don't need to walk them again)
 
1478
        # TODO: processed_new_refs becomes O(all_chks), consider switching to
 
1479
        #       SimpleSet here.
 
1480
        self._processed_new_refs = set()
 
1481
        self._search_key_func = search_key_func
 
1482
 
 
1483
        # The uninteresting and interesting nodes to be searched
 
1484
        self._old_queue = []
 
1485
        self._new_queue = []
 
1486
        # Holds the (key, value) items found when processing the root nodes,
 
1487
        # waiting for the uninteresting nodes to be walked
 
1488
        self._new_item_queue = []
 
1489
        self._state = None
 
1490
 
 
1491
    def _read_nodes_from_store(self, keys):
 
1492
        # We chose not to use _page_cache, because we think in terms of records
 
1493
        # to be yielded. Also, we expect to touch each page only 1 time during
 
1494
        # this code. (We may want to evaluate saving the raw bytes into the
 
1495
        # page cache, which would allow a working tree update after the fetch
 
1496
        # to not have to read the bytes again.)
 
1497
        as_st = StaticTuple.from_sequence
 
1498
        stream = self._store.get_record_stream(keys, 'unordered', True)
 
1499
        for record in stream:
 
1500
            if self._pb is not None:
 
1501
                self._pb.tick()
 
1502
            if record.storage_kind == 'absent':
 
1503
                raise errors.NoSuchRevision(self._store, record.key)
1426
1504
            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)
 
1505
            node = _deserialise(bytes, record.key,
 
1506
                                search_key_func=self._search_key_func)
1430
1507
            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)
 
1508
                # Note we don't have to do node.refs() because we know that
 
1509
                # there are no children that have been pushed into this node
 
1510
                # Note: Using as_st() here seemed to save 1.2MB, which would
 
1511
                #       indicate that we keep 100k prefix_refs around while
 
1512
                #       processing. They *should* be shorter lived than that...
 
1513
                #       It does cost us ~10s of processing time
 
1514
                #prefix_refs = [as_st(item) for item in node._items.iteritems()]
 
1515
                prefix_refs = node._items.items()
 
1516
                items = []
1437
1517
            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)
 
1518
                prefix_refs = []
 
1519
                # Note: We don't use a StaticTuple here. Profiling showed a
 
1520
                #       minor memory improvement (0.8MB out of 335MB peak 0.2%)
 
1521
                #       But a significant slowdown (15s / 145s, or 10%)
 
1522
                items = node._items.items()
 
1523
            yield record, node, prefix_refs, items
 
1524
 
 
1525
    def _read_old_roots(self):
 
1526
        old_chks_to_enqueue = []
 
1527
        all_old_chks = self._all_old_chks
 
1528
        for record, node, prefix_refs, items in \
 
1529
                self._read_nodes_from_store(self._old_root_keys):
 
1530
            # Uninteresting node
 
1531
            prefix_refs = [p_r for p_r in prefix_refs
 
1532
                                if p_r[1] not in all_old_chks]
 
1533
            new_refs = [p_r[1] for p_r in prefix_refs]
 
1534
            all_old_chks.update(new_refs)
 
1535
            # TODO: This might be a good time to turn items into StaticTuple
 
1536
            #       instances and possibly intern them. However, this does not
 
1537
            #       impact 'initial branch' performance, so I'm not worrying
 
1538
            #       about this yet
 
1539
            self._all_old_items.update(items)
 
1540
            # Queue up the uninteresting references
 
1541
            # Don't actually put them in the 'to-read' queue until we have
 
1542
            # finished checking the interesting references
 
1543
            old_chks_to_enqueue.extend(prefix_refs)
 
1544
        return old_chks_to_enqueue
 
1545
 
 
1546
    def _enqueue_old(self, new_prefixes, old_chks_to_enqueue):
 
1547
        # At this point, we have read all the uninteresting and interesting
 
1548
        # items, so we can queue up the uninteresting stuff, knowing that we've
 
1549
        # handled the interesting ones
 
1550
        for prefix, ref in old_chks_to_enqueue:
 
1551
            not_interesting = True
 
1552
            for i in xrange(len(prefix), 0, -1):
 
1553
                if prefix[:i] in new_prefixes:
 
1554
                    not_interesting = False
 
1555
                    break
 
1556
            if not_interesting:
 
1557
                # This prefix is not part of the remaining 'interesting set'
 
1558
                continue
 
1559
            self._old_queue.append(ref)
 
1560
 
 
1561
    def _read_all_roots(self):
 
1562
        """Read the root pages.
 
1563
 
 
1564
        This is structured as a generator, so that the root records can be
 
1565
        yielded up to whoever needs them without any buffering.
 
1566
        """
 
1567
        # This is the bootstrap phase
 
1568
        if not self._old_root_keys:
 
1569
            # With no old_root_keys we can just shortcut and be ready
 
1570
            # for _flush_new_queue
 
1571
            self._new_queue = list(self._new_root_keys)
 
1572
            return
 
1573
        old_chks_to_enqueue = self._read_old_roots()
 
1574
        # filter out any root keys that are already known to be uninteresting
 
1575
        new_keys = set(self._new_root_keys).difference(self._all_old_chks)
 
1576
        # These are prefixes that are present in new_keys that we are
 
1577
        # thinking to yield
 
1578
        new_prefixes = set()
 
1579
        # We are about to yield all of these, so we don't want them getting
 
1580
        # added a second time
 
1581
        processed_new_refs = self._processed_new_refs
 
1582
        processed_new_refs.update(new_keys)
 
1583
        for record, node, prefix_refs, items in \
 
1584
                self._read_nodes_from_store(new_keys):
 
1585
            # At this level, we now know all the uninteresting references
 
1586
            # So we filter and queue up whatever is remaining
 
1587
            prefix_refs = [p_r for p_r in prefix_refs
 
1588
                           if p_r[1] not in self._all_old_chks
 
1589
                              and p_r[1] not in processed_new_refs]
 
1590
            refs = [p_r[1] for p_r in prefix_refs]
 
1591
            new_prefixes.update([p_r[0] for p_r in prefix_refs])
 
1592
            self._new_queue.extend(refs)
 
1593
            # TODO: We can potentially get multiple items here, however the
 
1594
            #       current design allows for this, as callers will do the work
 
1595
            #       to make the results unique. We might profile whether we
 
1596
            #       gain anything by ensuring unique return values for items
 
1597
            # TODO: This might be a good time to cast to StaticTuple, as
 
1598
            #       self._new_item_queue will hold the contents of multiple
 
1599
            #       records for an extended lifetime
 
1600
            new_items = [item for item in items
 
1601
                               if item not in self._all_old_items]
 
1602
            self._new_item_queue.extend(new_items)
 
1603
            new_prefixes.update([self._search_key_func(item[0])
 
1604
                                 for item in new_items])
 
1605
            processed_new_refs.update(refs)
 
1606
            yield record
 
1607
        # For new_prefixes we have the full length prefixes queued up.
 
1608
        # However, we also need possible prefixes. (If we have a known ref to
 
1609
        # 'ab', then we also need to include 'a'.) So expand the
 
1610
        # new_prefixes to include all shorter prefixes
 
1611
        for prefix in list(new_prefixes):
 
1612
            new_prefixes.update([prefix[:i] for i in xrange(1, len(prefix))])
 
1613
        self._enqueue_old(new_prefixes, old_chks_to_enqueue)
 
1614
 
 
1615
    def _flush_new_queue(self):
 
1616
        # No need to maintain the heap invariant anymore, just pull things out
 
1617
        # and process them
 
1618
        refs = set(self._new_queue)
 
1619
        self._new_queue = []
 
1620
        # First pass, flush all interesting items and convert to using direct refs
 
1621
        all_old_chks = self._all_old_chks
 
1622
        processed_new_refs = self._processed_new_refs
 
1623
        all_old_items = self._all_old_items
 
1624
        new_items = [item for item in self._new_item_queue
 
1625
                           if item not in all_old_items]
 
1626
        self._new_item_queue = []
 
1627
        if new_items:
 
1628
            yield None, new_items
 
1629
        refs = refs.difference(all_old_chks)
 
1630
        processed_new_refs.update(refs)
 
1631
        while refs:
 
1632
            # TODO: Using a SimpleSet for self._processed_new_refs and
 
1633
            #       saved as much as 10MB of peak memory. However, it requires
 
1634
            #       implementing a non-pyrex version.
 
1635
            next_refs = set()
 
1636
            next_refs_update = next_refs.update
 
1637
            # Inlining _read_nodes_from_store improves 'bzr branch bzr.dev'
 
1638
            # from 1m54s to 1m51s. Consider it.
 
1639
            for record, _, p_refs, items in self._read_nodes_from_store(refs):
 
1640
                if all_old_items:
 
1641
                    # using the 'if' check saves about 145s => 141s, when
 
1642
                    # streaming initial branch of Launchpad data.
 
1643
                    items = [item for item in items
 
1644
                             if item not in all_old_items]
 
1645
                yield record, items
 
1646
                next_refs_update([p_r[1] for p_r in p_refs])
 
1647
                del p_refs
 
1648
            # set1.difference(set/dict) walks all of set1, and checks if it
 
1649
            # exists in 'other'.
 
1650
            # set1.difference(iterable) walks all of iterable, and does a
 
1651
            # 'difference_update' on a clone of set1. Pick wisely based on the
 
1652
            # expected sizes of objects.
 
1653
            # in our case it is expected that 'new_refs' will always be quite
 
1654
            # small.
 
1655
            next_refs = next_refs.difference(all_old_chks)
 
1656
            next_refs = next_refs.difference(processed_new_refs)
 
1657
            processed_new_refs.update(next_refs)
 
1658
            refs = next_refs
 
1659
 
 
1660
    def _process_next_old(self):
 
1661
        # Since we don't filter uninteresting any further than during
 
1662
        # _read_all_roots, process the whole queue in a single pass.
 
1663
        refs = self._old_queue
 
1664
        self._old_queue = []
 
1665
        all_old_chks = self._all_old_chks
 
1666
        for record, _, prefix_refs, items in self._read_nodes_from_store(refs):
 
1667
            # TODO: Use StaticTuple here?
 
1668
            self._all_old_items.update(items)
 
1669
            refs = [r for _,r in prefix_refs if r not in all_old_chks]
 
1670
            self._old_queue.extend(refs)
 
1671
            all_old_chks.update(refs)
 
1672
 
 
1673
    def _process_queues(self):
 
1674
        while self._old_queue:
 
1675
            self._process_next_old()
 
1676
        return self._flush_new_queue()
 
1677
 
 
1678
    def process(self):
 
1679
        for record in self._read_all_roots():
 
1680
            yield record, []
 
1681
        for record, items in self._process_queues():
 
1682
            yield record, items
1442
1683
 
1443
1684
 
1444
1685
def iter_interesting_nodes(store, interesting_root_keys,
1455
1696
    :return: Yield
1456
1697
        (interesting record, {interesting key:values})
1457
1698
    """
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
 
1699
    iterator = CHKMapDifference(store, interesting_root_keys,
 
1700
                                uninteresting_root_keys,
 
1701
                                search_key_func=store._search_key_func,
 
1702
                                pb=pb)
 
1703
    return iterator.process()
1524
1704
 
1525
1705
 
1526
1706
try:
1530
1710
        _deserialise_leaf_node,
1531
1711
        _deserialise_internal_node,
1532
1712
        )
1533
 
except ImportError:
 
1713
except ImportError, e:
 
1714
    osutils.failed_to_load_extension(e)
1534
1715
    from bzrlib._chk_map_py import (
1535
1716
        _search_key_16,
1536
1717
        _search_key_255,
1539
1720
        )
1540
1721
search_key_registry.register('hash-16-way', _search_key_16)
1541
1722
search_key_registry.register('hash-255-way', _search_key_255)
 
1723
 
 
1724
 
 
1725
def _check_key(key):
 
1726
    """Helper function to assert that a key is properly formatted.
 
1727
 
 
1728
    This generally shouldn't be used in production code, but it can be helpful
 
1729
    to debug problems.
 
1730
    """
 
1731
    if type(key) is not StaticTuple:
 
1732
        raise TypeError('key %r is not StaticTuple but %s' % (key, type(key)))
 
1733
    if len(key) != 1:
 
1734
        raise ValueError('key %r should have length 1, not %d' % (key, len(key),))
 
1735
    if type(key[0]) is not str:
 
1736
        raise TypeError('key %r should hold a str, not %r'
 
1737
                        % (key, type(key[0])))
 
1738
    if not key[0].startswith('sha1:'):
 
1739
        raise ValueError('key %r should point to a sha1:' % (key,))
 
1740
 
 
1741