storage: correctly restore create and ver

Add a restore func to correctly restore create reversion and
version of keys for the index.

storage/index.go

@@ -11,6 +11,7 @@ type index interface {
 	Get(key []byte, atRev int64) (rev, created reversion, ver int64, err error)
 	Range(key, end []byte, atRev int64) ([][]byte, []reversion)
 	Put(key []byte, rev reversion)
+	Restore(key []byte, created, modified reversion, ver int64)
 	Tombstone(key []byte, rev reversion) error
 	Compact(rev int64) map[reversion]struct{}
 	Equal(b index) bool
@@ -42,6 +43,21 @@ func (ti *treeIndex) Put(key []byte, rev reversion) {
 	okeyi.put(rev.main, rev.sub)
 }
 
+func (ti *treeIndex) Restore(key []byte, created, modified reversion, ver int64) {
+	keyi := &keyIndex{key: key}
+
+	ti.Lock()
+	defer ti.Unlock()
+	item := ti.tree.Get(keyi)
+	if item == nil {
+		keyi.restore(created, modified, ver)
+		ti.tree.ReplaceOrInsert(keyi)
+		return
+	}
+	okeyi := item.(*keyIndex)
+	okeyi.put(modified.main, modified.sub)
+}
+
 func (ti *treeIndex) Get(key []byte, atRev int64) (modified, created reversion, ver int64, err error) {
 	keyi := &keyIndex{key: key}
 

storage/key_index.go

@@ -78,6 +78,16 @@ func (ki *keyIndex) put(main int64, sub int64) {
 	ki.modified = rev
 }
 
+func (ki *keyIndex) restore(created, modified reversion, ver int64) {
+	if len(ki.generations) != 0 {
+		log.Panicf("store.keyindex: cannot restore non-empty keyIndex")
+	}
+
+	ki.modified = modified
+	g := generation{created: created, ver: ver, revs: []reversion{modified}}
+	ki.generations = append(ki.generations, g)
+}
+
 // tombstone puts a reversion, pointing to a tombstone, to the keyIndex.
 // It also creates a new empty generation in the keyIndex.
 func (ki *keyIndex) tombstone(main int64, sub int64) {

storage/kvstore.go

@@ -205,7 +205,7 @@ func (s *store) Restore() error {
 		// restore index
 		switch e.Type {
 		case storagepb.PUT:
-			s.kvindex.Put(e.Kv.Key, rev)
+			s.kvindex.Restore(e.Kv.Key, reversion{e.Kv.CreateIndex, 0}, rev, e.Kv.Version)
 		case storagepb.DELETE:
 			s.kvindex.Tombstone(e.Kv.Key, rev)
 		default: