Rev 14: Allow only profiling one index class. in http://people.ubuntu.com/~robertc/baz2.0/plugins/index2/trunk

[Robert Collins]

At http://people.ubuntu.com/~robertc/baz2.0/plugins/index2/trunk

------------------------------------------------------------
revno: 14
revision-id: robertc at robertcollins.net-20080702085603-pcd6iimxvyji6wtd
parent: robertc at robertcollins.net-20080702084338-45774breb3cit718
committer: Robert Collins <robertc at robertcollins.net>
branch nick: trunk
timestamp: Wed 2008-07-02 18:56:03 +1000
message:
  Allow only profiling one index class.
modified:
  indexbench.py                  indexbench.py-20080702083855-5tju02y79rw7kkzh-1
=== modified file 'indexbench.py'
--- a/indexbench.py	2008-07-02 08:43:38 +0000
+++ b/indexbench.py	2008-07-02 08:56:03 +0000
@@ -12,11 +12,13 @@
 from bzrlib.knit import _KnitGraphIndex, KnitVersionedFiles
 from bzrlib.graph import Graph
 from bzrlib.commands import Command
+from bzrlib.option import Option
 
 
 key_count = 0
 
-def drop_cache():
+def drop_caches():
+    """Call a script called drop-caches which should flush os caches."""
     os.system('drop-caches')
 
 def iter_indices(names, t, factory):
@@ -38,7 +40,11 @@
 
 
 def time_index(names, source, factory, builder, target, label, name_keys,
-    text_keys):
+    text_keys, use_drop_cache):
+    if use_drop_cache:
+        drop_cache = drop_caches
+    else:
+        drop_cache = lambda:None
     drop_cache()
 # Baseline time
     now = time.time()
@@ -135,38 +141,48 @@
 class cmd_indexbench(Command):
     """Benchmark index performance."""
 
+    takes_options = [
+        Option('graph', help='bench the GraphIndex class'),
+        Option('btree', help='bench the BTreeGraphIndex class'),
+        Option('drop-cache', help='drop caches between fixtures'),
+        ]
+
+
     takes_args = ['sample_repository']
 
-    def run(self, sample_repository):
+    def run(self, sample_repository, graph=True, btree=True, drop_cache=False):
+        source = get_transport(sample_repository)
+        source = source.clone('.bzr/repository/indices')
+        names = source.list_dir('.')
+        name_keys = {}
+        text_keys = set()
+        for name, index in iter_indices(names, source, GraphIndex):
+            name_keys[name] = []
+            for item in index.iter_all_entries():
+                name_keys[name].append(item[1])
+            # Randomise for later tests
+            shuffle(name_keys[name])
+            if name.endswith('.tix'):
+                text_keys.update(name_keys[name])
+        text_keys = list(text_keys)
+        # pick 2000 text keys to pretend to reconstruct
+        text_keys = sample(text_keys, 2000)
+        global key_count
+        key_count = sum(map(len, name_keys.itervalues()))
+
+        if btree:
+            self.test_class(names, source, BTreeGraphIndex, BTreeBuilder,
+                'BTreeIndex', name_keys, text_keys, drop_cache)
+        if graph:
+            self.test_class(names, source, GraphIndex, GraphIndexBuilder,
+                'GraphIndex', name_keys, text_keys, drop_cache)
+
+    def test_class(self, names, source, factory, builder, label, name_keys,
+        text_keys, drop_cache):
         workingdir = tempfile.mkdtemp()
         t = get_transport(workingdir)
         try:
-            source = get_transport(sample_repository)
-            source = source.clone('.bzr/repository/indices')
-# step 1:
-# load indices (gets keys for sizing)
-            names = source.list_dir('.')
-            name_keys = {}
-            text_keys = set()
-            for name, index in iter_indices(names, source, GraphIndex):
-                name_keys[name] = []
-                for item in index.iter_all_entries():
-                    name_keys[name].append(item[1])
-                # Randomise for later tests
-                shuffle(name_keys[name])
-                if name.endswith('.tix'):
-                    text_keys.update(name_keys[name])
-            text_keys = list(text_keys)
-            # pick 2000 text keys to pretend to reconstruct
-            text_keys = sample(text_keys, 2000)
-            global key_count
-            key_count = sum(map(len, name_keys.itervalues()))
-            t.mkdir('1')
-            t.mkdir('2')
-            time_index(names, source, GraphIndex, GraphIndexBuilder, t.clone('1'),
-                'GraphIndex', name_keys, text_keys)
-            time_index(names, source, BTreeGraphIndex, BTreeBuilder, t.clone('2'),
-                'BTreeIndex', name_keys, text_keys)
-            sys.exit(0)
+            time_index(names, source, factory, builder, t, label, name_keys,
+                text_keys, drop_cache)
         finally:
             t.delete_tree('.')