Merge changes to use btrees in StoreTree more sensibly.

7 files changed, 116 insertions, 128 deletions

diff --git a/obnamlib/checksumtree.py b/obnamlib/checksumtree.py
index 3e2a6daa..43e27ed4 100644
--- a/obnamlib/checksumtree.py
+++ b/obnamlib/checksumtree.py
@@ -34,6 +34,7 @@ class ChecksumTree(obnamlib.StoreTree):
         key_bytes = len(self.key('', 0, 0))
         obnamlib.StoreTree.__init__(self, fs, name, key_bytes, node_size, 
                                     upload_queue_size, lru_size)
+        self.keep_just_one_tree = True
 
     def key(self, checksum, chunk_id, client_id):
         return struct.pack(self.fmt, checksum, chunk_id, client_id)
@@ -42,37 +43,31 @@ class ChecksumTree(obnamlib.StoreTree):
         return struct.unpack(self.fmt, key)
 
     def add(self, checksum, chunk_id, client_id):
-        self.require_forest()
+        self.start_changes()
         key = self.key(checksum, chunk_id, client_id)
-        if self.forest.trees:
-            t = self.forest.trees[-1]
-        else:
-            t = self.forest.new_tree()
-        t.insert(key, '')
+        self.tree.insert(key, '')
 
     def find(self, checksum):
-        if self.init_forest() and self.forest.trees:
-            t = self.forest.trees[-1]
+        if self.init_forest():
             minkey = self.key(checksum, 0, 0)
             maxkey = self.key(checksum, obnamlib.MAX_ID, obnamlib.MAX_ID)
+            t = self.forest.trees[-1]
             pairs = t.lookup_range(minkey, maxkey)
             return [self.unkey(key)[1] for key, value in pairs]
         else:
             return []
 
     def remove(self, checksum, chunk_id, client_id):
-        self.require_forest()
-        if self.forest.trees:
-            t = self.forest.trees[-1]
-            key = self.key(checksum, chunk_id, client_id)
-            t.remove_range(key, key)
+        self.start_changes()
+        key = self.key(checksum, chunk_id, client_id)
+        self.tree.remove_range(key, key)
 
     def chunk_is_used(self, checksum, chunk_id):
         '''Is a given chunk used by anyone?'''
-        if self.init_forest() and self.forest.trees:
-            t = self.forest.trees[-1]
+        if self.init_forest():
             minkey = self.key(checksum, chunk_id, 0)
             maxkey = self.key(checksum, chunk_id, obnamlib.MAX_ID)
+            t = self.forest.trees[-1]
             return len(t.lookup_range(minkey, maxkey)) > 0
         else:
             return False
diff --git a/obnamlib/chunklist.py b/obnamlib/chunklist.py
index 60859283..8c651bfa 100644
--- a/obnamlib/chunklist.py
+++ b/obnamlib/chunklist.py
@@ -36,27 +36,23 @@ class ChunkList(obnamlib.StoreTree):
         self.key_bytes = len(self.key(0))
         obnamlib.StoreTree.__init__(self, fs, 'chunklist', self.key_bytes, 
                                     node_size, upload_queue_size, lru_size)
+        self.keep_just_one_tree = True
 
     def key(self, chunk_id):
         return struct.pack('!Q', chunk_id)
 
     def add(self, chunk_id, checksum):
-        self.require_forest()
-        if not self.forest.trees:
-            t = self.forest.new_tree()
-        else:
-            t = self.forest.trees[-1]
-        t.insert(self.key(chunk_id), checksum)
+        self.start_changes()
+        self.tree.insert(self.key(chunk_id), checksum)
         
     def get_checksum(self, chunk_id):
-        if self.init_forest() and self.forest.trees:
-            return self.forest.trees[-1].lookup(self.key(chunk_id))
+        if self.init_forest():
+            t = self.forest.trees[-1]
+            return t.lookup(self.key(chunk_id))
         raise KeyError(chunk_id)
         
     def remove(self, chunk_id):
-        self.require_forest()
-        if self.forest.trees:
-            t = self.forest.trees[-1]
-            key = self.key(chunk_id)
-            t.remove_range(key, key)
+        self.start_changes()
+        key = self.key(chunk_id)
+        self.tree.remove_range(key, key)
 
diff --git a/obnamlib/clientlist.py b/obnamlib/clientlist.py
index e67c9220..ff924388 100644
--- a/obnamlib/clientlist.py
+++ b/obnamlib/clientlist.py
@@ -44,6 +44,7 @@ class ClientList(obnamlib.StoreTree):
         self.maxkey = self.hashkey('\xff' * self.hash_len, obnamlib.MAX_ID)
         obnamlib.StoreTree.__init__(self, fs, 'clientlist', self.key_bytes, 
                                     node_size, upload_queue_size, lru_size)
+        self.keep_just_one_tree = True
 
     def hashfunc(self, string):
         return hashlib.new('md5', string).digest()
@@ -84,27 +85,20 @@ class ClientList(obnamlib.StoreTree):
         return self.find_client_id(t, client_name)
 
     def add_client(self, client_name):
-        self.require_forest()
-        if not self.forest.trees:
-            t = self.forest.new_tree()
-        else:
-            t = self.forest.new_tree(old=self.forest.trees[-1])
-
-        if self.find_client_id(t, client_name) is None:
+        self.start_changes()
+        if self.find_client_id(self.tree, client_name) is None:
             while True:
                 candidate_id = self.random_id()
                 key = self.key(client_name, candidate_id)
                 try:
-                    t.lookup(key)
+                    self.tree.lookup(key)
                 except KeyError:
                     break
-            t.insert(self.key(client_name, candidate_id), client_name)
+            self.tree.insert(self.key(client_name, candidate_id), client_name)
         
     def remove_client(self, client_name):
-        self.require_forest()
-        if self.forest.trees:
-            t = self.forest.new_tree(old=self.forest.trees[-1])
-            client_id = self.find_client_id(t, client_name)
-            if client_id is not None:
-                t.remove(self.key(client_name, client_id))
+        self.start_changes()
+        client_id = self.find_client_id(self.tree, client_name)
+        if client_id is not None:
+            self.tree.remove(self.key(client_name, client_id))
 
diff --git a/obnamlib/clientmetadatatree.py b/obnamlib/clientmetadatatree.py
index 087d80fe..b3d633cd 100644
--- a/obnamlib/clientmetadatatree.py
+++ b/obnamlib/clientmetadatatree.py
@@ -68,7 +68,6 @@ class ClientMetadataTree(obnamlib.StoreTree):
         obnamlib.StoreTree.__init__(self, fs, client_dir, key_bytes, 
                                     node_size, upload_queue_size, lru_size)
         self.genhash = self.hash_name('generation')
-        self.curgen = None
         self.known_generations = dict()
 
     def hash_name(self, filename):
@@ -143,12 +142,10 @@ class ClientMetadataTree(obnamlib.StoreTree):
         return tree.insert(key, struct.pack('!Q', value))
 
     def commit(self, current_time=time.time):
-        if self.forest:
-            if self.curgen:
-                now = int(current_time())
-                self._insert_int(self.curgen, self.genkey(self.GEN_ENDED), now)
-                self.curgen = None
-            self.forest.commit()
+        if self.tree:
+            now = int(current_time())
+            self._insert_int(self.tree, self.genkey(self.GEN_ENDED), now)
+        obnamlib.StoreTree.commit(self)
 
     def find_generation(self, genid):
         if genid in self.known_generations:
@@ -170,21 +167,16 @@ class ClientMetadataTree(obnamlib.StoreTree):
             return []
 
     def start_generation(self, current_time=time.time):
-        assert self.curgen is None
-        if self.forest.trees:
-            old = self.forest.trees[-1]
-        else:
-            old = None
-        self.curgen = self.forest.new_tree(old=old)
+        self.start_changes()
         gen_id = self.forest.new_id()
         now = int(current_time())
-        self._insert_int(self.curgen, self.genkey(self.GEN_ID), gen_id)
-        self._insert_int(self.curgen, self.genkey(self.GEN_STARTED), now)
+        self._insert_int(self.tree, self.genkey(self.GEN_ID), gen_id)
+        self._insert_int(self.tree, self.genkey(self.GEN_STARTED), now)
 
     def set_current_generation_is_checkpoint(self, is_checkpoint):
         value = 1 if is_checkpoint else 0
         key = self.genkey(self.GEN_IS_CHECKPOINT)
-        self._insert_int(self.curgen, key, value)
+        self._insert_int(self.tree, key, value)
 
     def get_is_checkpoint(self, genid):
         tree = self.find_generation(genid)
@@ -196,8 +188,8 @@ class ClientMetadataTree(obnamlib.StoreTree):
 
     def remove_generation(self, genid):
         tree = self.find_generation(genid)
-        if tree == self.curgen:
-            self.curgen = None
+        if tree == self.tree:
+            self.tree = None
         self.forest.remove_tree(tree)
 
     def get_generation_id(self, tree):
@@ -215,8 +207,8 @@ class ClientMetadataTree(obnamlib.StoreTree):
                 self._lookup_time(tree, self.GEN_ENDED))
 
     def create(self, filename, encoded_metadata):
-        file_id = self.get_file_id(self.curgen, filename)
-        gen_id = self.get_generation_id(self.curgen)
+        file_id = self.get_file_id(self.tree, filename)
+        gen_id = self.get_generation_id(self.tree)
         try:
             old_metadata = self.get_metadata(gen_id, filename)
         except KeyError:
@@ -228,14 +220,14 @@ class ClientMetadataTree(obnamlib.StoreTree):
         parent = os.path.dirname(filename)
         if parent != filename: # root dir is its own parent
             basename = os.path.basename(filename)
-            parent_id = self.get_file_id(self.curgen, parent)
+            parent_id = self.get_file_id(self.tree, parent)
             key = self.fskey(parent_id, self.DIR_CONTENTS, file_id)
             # We could just insert, but that would cause unnecessary
             # churn in the tree if nothing changes.
             try:
-                self.curgen.lookup(key)
+                self.tree.lookup(key)
             except KeyError:
-                self.curgen.insert(key, basename)
+                self.tree.insert(key, basename)
 
     def get_metadata(self, genid, filename):
         tree = self.find_generation(genid)
@@ -245,41 +237,41 @@ class ClientMetadataTree(obnamlib.StoreTree):
         return tree.lookup(key)
 
     def set_metadata(self, filename, encoded_metadata):
-        file_id = self.get_file_id(self.curgen, filename)
+        file_id = self.get_file_id(self.tree, filename)
         key1 = self.fskey(file_id, self.FILE_NAME, file_id)
-        self.curgen.insert(key1, filename)
+        self.tree.insert(key1, filename)
         
         key2 = self.fskey(file_id, self.FILE_METADATA, 
                           self.FILE_METADATA_ENCODED)
-        self.curgen.insert(key2, encoded_metadata)
+        self.tree.insert(key2, encoded_metadata)
 
     def remove(self, filename):
-        file_id = self.get_file_id(self.curgen, filename)
-        genid = self.get_generation_id(self.curgen)
+        file_id = self.get_file_id(self.tree, filename)
+        genid = self.get_generation_id(self.tree)
 
         # Remove any children.
         minkey = self.fskey(file_id, self.DIR_CONTENTS, 0)
         maxkey = self.fskey(file_id, self.DIR_CONTENTS, obnamlib.MAX_ID)
-        for key, basename in self.curgen.lookup_range(minkey, maxkey):
+        for key, basename in self.tree.lookup_range(minkey, maxkey):
             self.remove(os.path.join(filename, basename))
 
         # Remove chunk refs.
         for chunkid in self.get_file_chunks(genid, filename):
             key = self.chunk_key(chunkid, file_id)
-            self.curgen.remove_range(key, key)
+            self.tree.remove_range(key, key)
             
         # Remove this file's metadata.
         minkey = self.fskey(file_id, 0, 0)
         maxkey = self.fskey(file_id, self.TYPE_MAX, self.SUBKEY_MAX)
-        self.curgen.remove_range(minkey, maxkey)
+        self.tree.remove_range(minkey, maxkey)
 
         # Also remove from parent's contents.
         parent = os.path.dirname(filename)
         if parent != filename: # root dir is its own parent
-            parent_id = self.get_file_id(self.curgen, parent)
+            parent_id = self.get_file_id(self.tree, parent)
             key = self.fskey(parent_id, self.DIR_CONTENTS, file_id)
             # The range removal will work even if the key does not exist.
-            self.curgen.remove_range(key, key)
+            self.tree.remove_range(key, key)
             
     def listdir(self, genid, dirname):
         tree = self.find_generation(genid)
@@ -301,22 +293,22 @@ class ClientMetadataTree(obnamlib.StoreTree):
                 for key, value in pairs]
     
     def set_file_chunks(self, filename, chunkids):
-        file_id = self.get_file_id(self.curgen, filename)
+        file_id = self.get_file_id(self.tree, filename)
         minkey = self.fskey(file_id, self.FILE_CHUNKS, 0)
         maxkey = self.fskey(file_id, self.FILE_CHUNKS, self.SUBKEY_MAX)
         old_chunks = set(struct.unpack('!Q', v)[0]
-                         for k,v in self.curgen.lookup_range(minkey, maxkey))
-        self.curgen.remove_range(minkey, maxkey)
+                         for k,v in self.tree.lookup_range(minkey, maxkey))
+        self.tree.remove_range(minkey, maxkey)
         for i, chunkid in enumerate(chunkids):
             key = self.fskey(file_id, self.FILE_CHUNKS, i)
-            self.curgen.insert(key, struct.pack('!Q', chunkid))
+            self.tree.insert(key, struct.pack('!Q', chunkid))
             if chunkid not in old_chunks:
-                self.curgen.insert(self.chunk_key(chunkid, file_id), '')
+                self.tree.insert(self.chunk_key(chunkid, file_id), '')
             else:
                 old_chunks.remove(chunkid)
         for chunkid in old_chunks:
             key = self.chunk_key(chunkid, file_id)
-            self.curgen.remove_range(key, key)
+            self.tree.remove_range(key, key)
 
     def chunk_in_use(self, gen_id, chunk_id):
         '''Is a chunk used by a generation?'''
diff --git a/obnamlib/clientmetadatatree_tests.py b/obnamlib/clientmetadatatree_tests.py
index 11ec2612..8f7d990f 100644
--- a/obnamlib/clientmetadatatree_tests.py
+++ b/obnamlib/clientmetadatatree_tests.py
@@ -36,38 +36,37 @@ class ClientMetadataTreeTests(unittest.TestCase):
         shutil.rmtree(self.tempdir)
 
     def test_has_not_current_generation_initially(self):
-        self.assertEqual(self.client.curgen, None)
+        self.assertEqual(self.client.tree, None)
     
     def test_lists_no_generations_initially(self):
         self.assertEqual(self.client.list_generations(), [])
 
     def test_starts_generation(self):
-        self.client.require_forest()
         self.client.start_generation(current_time=lambda: 12765)
-        self.assertNotEqual(self.client.curgen, None)
+        self.assertNotEqual(self.client.tree, None)
         
         def lookup(x):
             key = self.client.genkey(x)
-            return self.client._lookup_int(self.client.curgen, key)
+            return self.client._lookup_int(self.client.tree, key)
 
-        genid = self.client.get_generation_id(self.client.curgen)
+        genid = self.client.get_generation_id(self.client.tree)
         self.assertEqual(lookup(self.client.GEN_ID), genid)
         self.assertEqual(lookup(self.client.GEN_STARTED), 12765)
         self.assertFalse(self.client.get_is_checkpoint(genid))
 
     def test_starts_second_generation(self):
-        self.client.require_forest()
         self.client.start_generation(current_time=lambda: 1)
-        genid1 = self.client.get_generation_id(self.client.curgen)
+        genid1 = self.client.get_generation_id(self.client.tree)
         self.client.commit()
+        self.assertEqual(self.client.tree, None)
         self.client.start_generation(current_time=lambda: 2)
-        self.assertNotEqual(self.client.curgen, None)
+        self.assertNotEqual(self.client.tree, None)
         
         def lookup(x):
             key = self.client.genkey(x)
-            return self.client._lookup_int(self.client.curgen, key)
+            return self.client._lookup_int(self.client.tree, key)
 
-        genid2 = self.client.get_generation_id(self.client.curgen)
+        genid2 = self.client.get_generation_id(self.client.tree)
         self.assertEqual(lookup(self.client.GEN_ID), genid2)
         self.assertNotEqual(genid1, genid2)
         self.assertEqual(lookup(self.client.GEN_STARTED), 2)
@@ -75,59 +74,52 @@ class ClientMetadataTreeTests(unittest.TestCase):
         self.assertEqual(self.client.list_generations(), [genid1, genid2])
 
     def test_sets_is_checkpoint(self):
-        self.client.require_forest()
         self.client.start_generation()
-        genid = self.client.get_generation_id(self.client.curgen)
+        genid = self.client.get_generation_id(self.client.tree)
         self.client.set_current_generation_is_checkpoint(True)
         self.assert_(self.client.get_is_checkpoint(genid))
 
     def test_unsets_is_checkpoint(self):
-        self.client.require_forest()
         self.client.start_generation()
-        genid = self.client.get_generation_id(self.client.curgen)
+        genid = self.client.get_generation_id(self.client.tree)
         self.client.set_current_generation_is_checkpoint(True)
         self.client.set_current_generation_is_checkpoint(False)
         self.assertFalse(self.client.get_is_checkpoint(genid))
 
     def test_removes_generation(self):
-        self.client.require_forest()
         self.client.start_generation()
         self.client.commit()
-        self.client.remove_generation(self.client.list_generations()[0])
+        genid = self.client.list_generations()[0]
+        self.client.remove_generation(genid)
         self.assertEqual(self.client.list_generations(), [])
 
     def test_removes_started_generation(self):
-        self.client.require_forest()
         self.client.start_generation()
         self.client.remove_generation(self.client.list_generations()[0])
         self.assertEqual(self.client.list_generations(), [])
-        self.assertEqual(self.client.curgen, None)
+        self.assertEqual(self.client.tree, None)
 
     def test_started_generation_has_start_time(self):
-        self.client.require_forest()
         self.client.start_generation(current_time=lambda: 1)
-        genid = self.client.get_generation_id(self.client.curgen)
+        genid = self.client.get_generation_id(self.client.tree)
         self.assertEqual(self.client.get_generation_times(genid), (1, None))
 
     def test_committed_generation_has_times(self):
-        self.client.require_forest()
         self.client.start_generation(current_time=lambda: 1)
-        genid = self.client.get_generation_id(self.client.curgen)
+        genid = self.client.get_generation_id(self.client.tree)
         self.client.commit(current_time=lambda: 2)
         self.assertEqual(self.client.get_generation_times(genid), (1, 2))
 
     def test_finds_generation_the_first_time(self):
-        self.client.require_forest()
         self.client.start_generation()
-        tree = self.client.curgen
+        tree = self.client.tree
         genid = self.client.get_generation_id(tree)
         self.client.commit()
         self.assertEqual(self.client.find_generation(genid), tree)
 
     def test_finds_generation_the_second_time(self):
-        self.client.require_forest()
         self.client.start_generation()
-        tree = self.client.curgen
+        tree = self.client.tree
         genid = self.client.get_generation_id(tree)
         self.client.commit()
         self.client.find_generation(genid)
@@ -138,7 +130,6 @@ class ClientMetadataTreeTests(unittest.TestCase):
         self.assertRaises(KeyError, self.client.find_generation, 0)
 
     def test_find_generation_raises_keyerror_for_unknown_generation(self):
-        self.client.require_forest()
         self.assertRaises(KeyError, self.client.find_generation, 0)
 
 
@@ -151,9 +142,8 @@ class ClientMetadataTreeFileOpsTests(unittest.TestCase):
                                             obnamlib.DEFAULT_NODE_SIZE,
                                             obnamlib.DEFAULT_UPLOAD_QUEUE_SIZE,
                                             obnamlib.DEFAULT_LRU_SIZE)
-        self.client.require_forest()
         self.client.start_generation()
-        self.clientid = self.client.get_generation_id(self.client.curgen)
+        self.clientid = self.client.get_generation_id(self.client.tree)
         self.file_metadata = obnamlib.Metadata(st_mode=stat.S_IFREG | 0666)
         self.file_encoded = obnamlib.store.encode_metadata(self.file_metadata)
         self.dir_metadata = obnamlib.Metadata(st_mode=stat.S_IFDIR | 0777)
@@ -246,24 +236,24 @@ class ClientMetadataTreeFileOpsTests(unittest.TestCase):
     def test_generation_has_no_chunk_refs_initially(self):
         minkey = self.client.chunk_key(0, 0)
         maxkey = self.client.chunk_key(obnamlib.MAX_ID, obnamlib.MAX_ID)
-        self.assertEqual(self.client.curgen.lookup_range(minkey, maxkey), [])
+        self.assertEqual(self.client.tree.lookup_range(minkey, maxkey), [])
 
     def test_set_file_chunks_adds_chunk_refs(self):
         self.client.set_file_chunks('/foo', [1, 2])
-        file_id = self.client.get_file_id(self.client.curgen, '/foo')
+        file_id = self.client.get_file_id(self.client.tree, '/foo')
         minkey = self.client.chunk_key(0, 0)
         maxkey = self.client.chunk_key(obnamlib.MAX_ID, obnamlib.MAX_ID)
-        self.assertEqual(set(self.client.curgen.lookup_range(minkey, maxkey)), 
+        self.assertEqual(set(self.client.tree.lookup_range(minkey, maxkey)), 
                          set([(self.client.chunk_key(1, file_id), ''),
                               (self.client.chunk_key(2, file_id), '')]))
 
     def test_set_file_chunks_removes_now_unused_chunk_refs(self):
         self.client.set_file_chunks('/foo', [1, 2])
         self.client.set_file_chunks('/foo', [1])
-        file_id = self.client.get_file_id(self.client.curgen, '/foo')
+        file_id = self.client.get_file_id(self.client.tree, '/foo')
         minkey = self.client.chunk_key(0, 0)
         maxkey = self.client.chunk_key(obnamlib.MAX_ID, obnamlib.MAX_ID)
-        self.assertEqual(self.client.curgen.lookup_range(minkey, maxkey), 
+        self.assertEqual(self.client.tree.lookup_range(minkey, maxkey), 
                          [(self.client.chunk_key(1, file_id), '')])
 
     def test_remove_removes_chunk_refs(self):
@@ -271,30 +261,30 @@ class ClientMetadataTreeFileOpsTests(unittest.TestCase):
         self.client.remove('/foo')
         minkey = self.client.chunk_key(0, 0)
         maxkey = self.client.chunk_key(obnamlib.MAX_ID, obnamlib.MAX_ID)
-        self.assertEqual(self.client.curgen.lookup_range(minkey, maxkey), [])
+        self.assertEqual(self.client.tree.lookup_range(minkey, maxkey), [])
         
     def test_report_chunk_not_in_use_initially(self):
-        gen_id = self.client.get_generation_id(self.client.curgen)
+        gen_id = self.client.get_generation_id(self.client.tree)
         self.assertFalse(self.client.chunk_in_use(gen_id, 0))
         
     def test_report_chunk_in_use_after_it_is(self):
-        gen_id = self.client.get_generation_id(self.client.curgen)
+        gen_id = self.client.get_generation_id(self.client.tree)
         self.client.set_file_chunks('/foo', [0])
         self.assertTrue(self.client.chunk_in_use(gen_id, 0))
 
     def test_lists_no_chunks_in_generation_initially(self):
-        gen_id = self.client.get_generation_id(self.client.curgen)
+        gen_id = self.client.get_generation_id(self.client.tree)
         self.assertEqual(self.client.list_chunks_in_generation(gen_id), [])
 
     def test_lists_used_chunks_in_generation(self):
-        gen_id = self.client.get_generation_id(self.client.curgen)
+        gen_id = self.client.get_generation_id(self.client.tree)
         self.client.set_file_chunks('/foo', [0])
         self.client.set_file_chunks('/bar', [1])
         self.assertEqual(set(self.client.list_chunks_in_generation(gen_id)), 
                          set([0, 1]))
 
     def test_lists_chunks_in_generation_only_once(self):
-        gen_id = self.client.get_generation_id(self.client.curgen)
+        gen_id = self.client.get_generation_id(self.client.tree)
         self.client.set_file_chunks('/foo', [0])
         self.client.set_file_chunks('/bar', [0])
         self.assertEqual(self.client.list_chunks_in_generation(gen_id), [0])
diff --git a/obnamlib/store.py b/obnamlib/store.py
index 4ad2fbec..b9eb0a9f 100644
--- a/obnamlib/store.py
+++ b/obnamlib/store.py
@@ -301,7 +301,7 @@ class Store(object):
                                                   self.node_size,
                                                   self.upload_queue_size, 
                                                   self.lru_size)
-        self.client.require_forest()
+        self.client.init_forest()
 
     @require_client_lock
     def unlock_client(self):
@@ -365,10 +365,9 @@ class Store(object):
         '''
         if self.new_generation is not None:
             raise obnamlib.Error('Cannot start two new generations')
-        self.client.require_forest()
         self.client.start_generation()
         self.new_generation = \
-            self.client.get_generation_id(self.client.curgen)
+            self.client.get_generation_id(self.client.tree)
         self.added_generations.append(self.new_generation)
         return self.new_generation
 
diff --git a/obnamlib/store_tree.py b/obnamlib/store_tree.py
index 4be15cb6..03e2d21d 100644
--- a/obnamlib/store_tree.py
+++ b/obnamlib/store_tree.py
@@ -21,7 +21,18 @@ import obnamlib
 
 class StoreTree(object):
 
-    '''A B-tree within a Store.'''
+    '''A B-tree within a Store.
+    
+    For read-only operation, call init_forest before doing anything.
+    
+    For read-write operation, call start_changes before doing anything,
+    and commit afterwards. In between, self.tree is the new tree to be 
+    modified. Note that self.tree is NOT available after init_forest.
+    
+    After init_forest or start_changes, self.forest is the opened forest.
+    Unlike self.tree, it will not go away after commit.
+    
+    '''
 
     def __init__(self, fs, dirname, key_bytes, node_size, upload_queue_size,
                  lru_size):
@@ -32,9 +43,12 @@ class StoreTree(object):
         self.upload_queue_size = upload_queue_size
         self.lru_size = lru_size
         self.forest = None
+        self.tree = None
+        self.keep_just_one_tree = False
 
     def init_forest(self):
         if self.forest is None:
+            assert self.tree is None
             if not self.fs.exists(self.dirname):
                 return False
             codec = btree.NodeCodec(self.key_bytes)
@@ -44,14 +58,22 @@ class StoreTree(object):
             self.forest = btree.Forest(ns)
         return True
 
-    def require_forest(self):
+    def start_changes(self):
         if not self.fs.exists(self.dirname):
             self.fs.mkdir(self.dirname)
         self.init_forest()
         assert self.forest is not None
+        if self.tree is None:
+            if self.forest.trees:
+                self.tree = self.forest.new_tree(self.forest.trees[-1])
+            else:
+                self.tree = self.forest.new_tree()
 
     def commit(self):
         if self.forest:
-            self.require_forest()
+            if self.keep_just_one_tree:
+                while len(self.forest.trees) > 1:
+                    self.forest.remove_tree(self.forest.trees[0])
             self.forest.commit()
+            self.tree = None