feat: improve lock

core/state/snapshot/difflayer.go

@@ -528,12 +528,14 @@ func (dl *diffLayer) flatten() snapshot {
 
 	// Before actually writing all our data to the parent, first ensure that the
 	// parent hasn't been 'corrupted' by someone else already flattening into it
-	if parent.stale.Swap(true) {
+	if parent.canLookup.Swap(true) {
 		panic("parent diff layer is stale") // we've flattened into the same parent from two children, boo
 	}
-	if parent.canLookup.Swap(true) {
+
+	if parent.stale.Swap(true) {
 		panic("parent diff layer is stale") // we've flattened into the same parent from two children, boo
 	}
+
 	//log.Info("Layer flattening stale", "layer", parent.Root(), "destructs", len(dl.destructSet))
 	for hash, data := range dl.accountData {
 		parent.accountData[hash] = data

core/state/snapshot/lookup.go

@@ -108,8 +108,6 @@ func (l *Lookup) addLayer(diff *diffLayer) {
 		lookupAddLayerTimer.UpdateSince(now)
 	}(time.Now())
 	layerIDCounter++
-	l.lock.Lock()
-	defer l.lock.Unlock()
 
 	for accountHash, _ := range diff.accountData {
 		l.state2LayerRoots[accountHash.String()] = append(l.state2LayerRoots[accountHash.String()], diff)
@@ -120,8 +118,6 @@ func (l *Lookup) addLayer(diff *diffLayer) {
 			l.state2LayerRoots[accountHash.String()+storageHash.String()] = append(l.state2LayerRoots[accountHash.String()+storageHash.String()], diff)
 		}
 	}
-
-	//l.addDescendant(diff)
 }
 
 // removeLayer traverses all the dirty state within the given diff layer and
@@ -131,8 +127,6 @@ func (l *Lookup) removeLayer(diff *diffLayer) error {
 		lookupRemoveLayerTimer.UpdateSince(now)
 	}(time.Now())
 	layerIDRemoveCounter++
-	l.lock.Lock()
-	defer l.lock.Unlock()
 
 	diffRoot := diff.Root()
 	for accountHash, _ := range diff.accountData {
@@ -195,8 +189,6 @@ func (l *Lookup) removeLayer(diff *diffLayer) error {
 		}
 	}
 
-	//l.removeDescendant(diff)
-
 	return nil
 }
 
@@ -222,8 +214,6 @@ func (l *Lookup) addDescendant(topDiffLayer Snapshot) {
 	defer func(now time.Time) {
 		lookupAddDescendantTimer.UpdateSince(now)
 	}(time.Now())
-	l.lock.Lock()
-	defer l.lock.Unlock()
 
 	// Link the new layer into the descendents set
 	for h := range diffAncestors(topDiffLayer) {
@@ -240,10 +230,22 @@ func (l *Lookup) removeDescendant(bottomDiffLayer Snapshot) {
 	defer func(now time.Time) {
 		lookupRemoveDescendantTimer.UpdateSince(now)
 	}(time.Now())
+
+	delete(l.descendants, bottomDiffLayer.Root())
+}
+
+func (l *Lookup) AddSnapshot(diff *diffLayer) {
 	l.lock.Lock()
 	defer l.lock.Unlock()
+	l.addLayer(diff)
+	l.addDescendant(diff)
+}
 
-	delete(l.descendants, bottomDiffLayer.Root())
+func (l *Lookup) RemoveSnapshot(diff *diffLayer) {
+	l.lock.Lock()
+	defer l.lock.Unlock()
+	l.removeLayer(diff)
+	l.removeDescendant(diff)
 }
 
 func (l *Lookup) LookupAccount(accountAddrHash common.Hash, head common.Hash) Snapshot {

core/state/snapshot/snapshot.go

@@ -409,8 +409,7 @@ func (t *Tree) Update(blockRoot common.Hash, parentRoot common.Hash, accounts ma
 
 	t.layers[snap.root] = snap
 	// update lookup, which in the tree lock guard.
-	t.lookup.addLayer(snap)
-	t.lookup.addDescendant(snap)
+	t.lookup.AddSnapshot(snap)
 	if t.baseDiff == nil || reflect.ValueOf(t.baseDiff).IsNil() {
 		t.baseDiff = snap
 	}
@@ -486,8 +485,7 @@ func (t *Tree) Cap(root common.Hash, layers int) error {
 		if !ok {
 			return
 		}
-		t.lookup.removeLayer(diff)
-		t.lookup.removeDescendant(diff)
+		t.lookup.RemoveSnapshot(diff)
 	}
 	var remove func(root common.Hash, snap snapshot)
 	remove = func(root common.Hash, snap snapshot) {
@@ -585,8 +583,7 @@ func (t *Tree) cap(diff *diffLayer, layers int) *diskLayer {
 		if !ok {
 			return
 		}
-		t.lookup.removeLayer(diff)
-		t.lookup.removeDescendant(diff)
+		t.lookup.RemoveSnapshot(diff)
 	}
 
 	//TODO:check it?