Skip to content

Commit ad39c2a

Browse files
BeniaminDrasoveanBeniaminDrasovean
committed
fix after review
1 parent 3d00f95 commit ad39c2a

2 files changed

Lines changed: 161 additions & 28 deletions

File tree

‎state/trackableDataTrie/trackableDataTrie.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -259,9 +259,12 @@ func (tdt *trackableDataTrie) SaveDirtyData(mainTrie common.Trie) ([]*stateChang
259259
}
260260

261261
func (tdt *trackableDataTrie) rollbackAppliedUpdates(dtr state.DataTrie, oldValues []core.TrieData) {
262+
var trieUpdate core.TrieData
263+
var err error
264+
262265
for i := len(oldValues) - 1; i >= 0; i-- {
263-
trieUpdate := oldValues[i]
264-
err := dtr.UpdateWithVersion(trieUpdate.Key, trieUpdate.Value, trieUpdate.Version)
266+
trieUpdate = oldValues[i]
267+
err = dtr.UpdateWithVersion(trieUpdate.Key, trieUpdate.Value, trieUpdate.Version)
265268
if err != nil {
266269
log.Error("could not apply rollback updates", "err", err, "key", trieUpdate.Key, "account", tdt.identifier)
267270
}
@@ -314,6 +317,10 @@ func (tdt *trackableDataTrie) updateTrie(dtr state.DataTrie) ([]*stateChange.Dat
314317

315318
isFirstMigration := oldVal.Version == core.NotSpecified && dataEntry.newVersion == core.AutoBalanceEnabled
316319
if isFirstMigration && len(dataTrieKey) != 0 {
320+
if !wasDeleted {
321+
trieUpdates = append(trieUpdates, oldVal)
322+
}
323+
317324
deletedData := core.TrieData{
318325
Key: dataTrieKey,
319326
Value: nil,

‎state/trackableDataTrie/trackableDataTrie_test.go‎

Lines changed: 152 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1192,8 +1192,8 @@ func TestTrackableDataTrie_SaveDirtyDataShouldRollbackPreviousUpdateWhenLaterGet
11921192
expectedErr := errors.New("expected get error")
11931193
firstKey := []byte("key1")
11941194
secondKey := []byte("key2")
1195-
firstOldValue := []byte("old1")
1196-
firstNewValue := []byte("new1")
1195+
oldValue := []byte("old1")
1196+
newValue := []byte("new1")
11971197
identifier := []byte("identifier")
11981198

11991199
getCalls := 0
@@ -1203,27 +1203,52 @@ func TestTrackableDataTrie_SaveDirtyDataShouldRollbackPreviousUpdateWhenLaterGet
12031203
trie := &trieMock.TrieStub{
12041204
GetCalled: func(key []byte) ([]byte, uint32, error) {
12051205
getCalls++
1206+
// second loop iteration always fail
1207+
if getCalls == 2 {
1208+
return nil, 0, expectedErr
1209+
}
12061210
if bytes.Equal(key, firstKey) {
1207-
return append(firstOldValue, append(firstKey, identifier...)...), 0, nil
1211+
return append(oldValue, append(firstKey, identifier...)...), 0, nil
12081212
}
12091213
if bytes.Equal(key, secondKey) {
1210-
return nil, 0, expectedErr
1214+
return append(oldValue, append(secondKey, identifier...)...), 0, nil
12111215
}
12121216
return nil, 0, nil
12131217
},
12141218
UpdateWithVersionCalled: func(key, value []byte, version core.TrieNodeVersion) error {
12151219
updateCalls++
1220+
// first time is the update
12161221
if updateCalls == 1 {
1217-
assert.Equal(t, firstKey, key)
1218-
assert.Equal(t, append(firstNewValue, append(firstKey, identifier...)...), value)
1219-
assert.Equal(t, core.NotSpecified, version)
1220-
return nil
1222+
if bytes.Equal(key, firstKey) {
1223+
assert.Equal(t, append(newValue, append(firstKey, identifier...)...), value)
1224+
assert.Equal(t, core.NotSpecified, version)
1225+
return nil
1226+
}
1227+
if bytes.Equal(key, secondKey) {
1228+
assert.Equal(t, append(newValue, append(secondKey, identifier...)...), value)
1229+
assert.Equal(t, core.NotSpecified, version)
1230+
return nil
1231+
}
1232+
assert.Fail(t, "this should not happen")
1233+
}
1234+
1235+
// second time is the rollback
1236+
if updateCalls == 2 {
1237+
rollbackCalled = true
1238+
1239+
if bytes.Equal(key, firstKey) {
1240+
assert.Equal(t, append(oldValue, append(firstKey, identifier...)...), value)
1241+
assert.Equal(t, core.NotSpecified, version)
1242+
return nil
1243+
}
1244+
if bytes.Equal(key, secondKey) {
1245+
assert.Equal(t, append(oldValue, append(secondKey, identifier...)...), value)
1246+
assert.Equal(t, core.NotSpecified, version)
1247+
return nil
1248+
}
1249+
assert.Fail(t, "this should not happen")
12211250
}
12221251

1223-
assert.Equal(t, firstKey, key)
1224-
assert.Equal(t, append(firstOldValue, append(firstKey, identifier...)...), value)
1225-
assert.Equal(t, core.NotSpecified, version)
1226-
rollbackCalled = true
12271252
return nil
12281253
},
12291254
}
@@ -1237,8 +1262,8 @@ func TestTrackableDataTrie_SaveDirtyDataShouldRollbackPreviousUpdateWhenLaterGet
12371262
)
12381263
tdt.SetDataTrie(trie)
12391264

1240-
_ = tdt.SaveKeyValue(firstKey, firstNewValue)
1241-
_ = tdt.SaveKeyValue(secondKey, []byte("new2"))
1265+
_ = tdt.SaveKeyValue(firstKey, newValue)
1266+
_ = tdt.SaveKeyValue(secondKey, newValue)
12421267

12431268
_, _, err := tdt.SaveDirtyData(trie)
12441269
require.ErrorIs(t, err, expectedErr)
@@ -1252,8 +1277,8 @@ func TestTrackableDataTrie_SaveDirtyDataShouldRollbackPreviousUpdateWhenLaterUpd
12521277
expectedErr := errors.New("expected update error")
12531278
firstKey := []byte("key1")
12541279
secondKey := []byte("key2")
1255-
firstOldValue := []byte("old1")
1256-
firstNewValue := []byte("new1")
1280+
oldValue := []byte("old1")
1281+
newValue := []byte("new1")
12571282
identifier := []byte("identifier")
12581283

12591284
updateCalls := 0
@@ -1262,26 +1287,44 @@ func TestTrackableDataTrie_SaveDirtyDataShouldRollbackPreviousUpdateWhenLaterUpd
12621287
trie := &trieMock.TrieStub{
12631288
GetCalled: func(key []byte) ([]byte, uint32, error) {
12641289
if bytes.Equal(key, firstKey) {
1265-
return append(firstOldValue, append(firstKey, identifier...)...), 0, nil
1290+
return append(oldValue, append(firstKey, identifier...)...), 0, nil
1291+
}
1292+
if bytes.Equal(key, secondKey) {
1293+
return append(oldValue, append(secondKey, identifier...)...), 0, nil
12661294
}
12671295
return nil, 0, nil
12681296
},
12691297
UpdateWithVersionCalled: func(key, value []byte, version core.TrieNodeVersion) error {
12701298
updateCalls++
12711299
switch updateCalls {
12721300
case 1:
1273-
assert.Equal(t, firstKey, key)
1274-
assert.Equal(t, append(firstNewValue, append(firstKey, identifier...)...), value)
1275-
assert.Equal(t, core.NotSpecified, version)
1301+
if bytes.Equal(key, firstKey) {
1302+
assert.Equal(t, append(newValue, append(firstKey, identifier...)...), value)
1303+
assert.Equal(t, core.NotSpecified, version)
1304+
return nil
1305+
}
1306+
if bytes.Equal(key, secondKey) {
1307+
assert.Equal(t, append(newValue, append(secondKey, identifier...)...), value)
1308+
assert.Equal(t, core.NotSpecified, version)
1309+
return nil
1310+
}
1311+
assert.Fail(t, "this should not happen")
12761312
return nil
12771313
case 2:
1278-
assert.Equal(t, secondKey, key)
12791314
return expectedErr
12801315
case 3:
1281-
assert.Equal(t, firstKey, key)
1282-
assert.Equal(t, append(firstOldValue, append(firstKey, identifier...)...), value)
1283-
assert.Equal(t, core.NotSpecified, version)
12841316
rollbackCalled = true
1317+
if bytes.Equal(key, firstKey) {
1318+
assert.Equal(t, append(oldValue, append(firstKey, identifier...)...), value)
1319+
assert.Equal(t, core.NotSpecified, version)
1320+
return nil
1321+
}
1322+
if bytes.Equal(key, secondKey) {
1323+
assert.Equal(t, append(oldValue, append(secondKey, identifier...)...), value)
1324+
assert.Equal(t, core.NotSpecified, version)
1325+
return nil
1326+
}
1327+
assert.Fail(t, "this should not happen")
12851328
return nil
12861329
default:
12871330
require.Fail(t, "unexpected update call")
@@ -1299,8 +1342,8 @@ func TestTrackableDataTrie_SaveDirtyDataShouldRollbackPreviousUpdateWhenLaterUpd
12991342
)
13001343
tdt.SetDataTrie(trie)
13011344

1302-
_ = tdt.SaveKeyValue(firstKey, firstNewValue)
1303-
_ = tdt.SaveKeyValue(secondKey, []byte("new2"))
1345+
_ = tdt.SaveKeyValue(firstKey, newValue)
1346+
_ = tdt.SaveKeyValue(secondKey, newValue)
13041347

13051348
_, _, err := tdt.SaveDirtyData(trie)
13061349
require.ErrorIs(t, err, expectedErr)
@@ -1500,3 +1543,86 @@ func TestTrackableDataTrie_SaveDirtyDataShouldRollbackDeleteWhenMetadataDecodeFa
15001543
assert.True(t, deleteCalled)
15011544
assert.True(t, rollbackCalled)
15021545
}
1546+
1547+
func TestTrackableDataTrie_SaveDirtyDataShouldRollbackDeleteOfANonMigratedKey(t *testing.T) {
1548+
t.Parallel()
1549+
1550+
expectedErr := errors.New("expected delete error")
1551+
identifier := []byte("identifier")
1552+
hasher := &hashingMocks.HasherMock{}
1553+
1554+
firstKey := []byte("key1")
1555+
firstOldTrieValue := append([]byte("old value 1"), append(firstKey, identifier...)...)
1556+
1557+
secondKey := []byte("key2")
1558+
secondOldTrieValue := append([]byte("old value 2"), append(secondKey, identifier...)...)
1559+
1560+
originalTrieValues := map[string][]byte{
1561+
string(firstKey): append([]byte(nil), firstOldTrieValue...),
1562+
string(secondKey): append([]byte(nil), secondOldTrieValue...),
1563+
}
1564+
1565+
trieValues := map[string][]byte{
1566+
string(firstKey): append([]byte(nil), firstOldTrieValue...),
1567+
string(secondKey): append([]byte(nil), secondOldTrieValue...),
1568+
}
1569+
1570+
successfulDeletes := make(map[string]struct{})
1571+
1572+
trie := &trieMock.TrieStub{
1573+
GetCalled: func(key []byte) ([]byte, uint32, error) {
1574+
// With AutoBalance enabled, retrieveValueFromTrie first checks the
1575+
// hashed key. These entries are intentionally old-format only.
1576+
if bytes.Equal(key, hasher.Compute(string(firstKey))) ||
1577+
bytes.Equal(key, hasher.Compute(string(secondKey))) {
1578+
return nil, 0, nil
1579+
}
1580+
1581+
return trieValues[string(key)], 0, nil
1582+
},
1583+
DeleteCalled: func(key []byte) error {
1584+
if len(successfulDeletes) == 1 {
1585+
return expectedErr
1586+
}
1587+
1588+
successfulDeletes[string(key)] = struct{}{}
1589+
delete(trieValues, string(key))
1590+
return nil
1591+
},
1592+
UpdateWithVersionCalled: func(key, value []byte, _ core.TrieNodeVersion) error {
1593+
if len(value) == 0 {
1594+
delete(trieValues, string(key))
1595+
return nil
1596+
}
1597+
1598+
trieValues[string(key)] = append([]byte(nil), value...)
1599+
return nil
1600+
},
1601+
}
1602+
1603+
tdt, err := trackableDataTrie.NewTrackableDataTrie(
1604+
identifier,
1605+
hasher,
1606+
&marshallerMock.MarshalizerMock{},
1607+
&enableEpochsHandlerMock.EnableEpochsHandlerStub{
1608+
IsFlagEnabledCalled: func(flag core.EnableEpochFlag) bool {
1609+
return flag == common.AutoBalanceDataTriesFlag
1610+
},
1611+
},
1612+
&stateMock.StateAccessesCollectorStub{},
1613+
)
1614+
require.NoError(t, err)
1615+
1616+
tdt.SetDataTrie(trie)
1617+
1618+
require.NoError(t, tdt.SaveKeyValue(firstKey, nil))
1619+
require.NoError(t, tdt.SaveKeyValue(secondKey, nil))
1620+
1621+
_, _, err = tdt.SaveDirtyData(trie)
1622+
require.ErrorIs(t, err, expectedErr)
1623+
require.Len(t, successfulDeletes, 1)
1624+
1625+
for key, expectedValue := range originalTrieValues {
1626+
assert.Equal(t, expectedValue, trieValues[key])
1627+
}
1628+
}

0 commit comments

Comments
 (0)