Skip to content

Commit

Permalink
Merge pull request #121 from cardano-foundation/fix_119
Browse files Browse the repository at this point in the history
Only rollback address token balance if needed
  • Loading branch information
Sotatek-HuyLe3a authored Jan 29, 2024
2 parents ee2af49 + 60c1210 commit 853641d
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 25 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import java.sql.Timestamp;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

import static org.cardanofoundation.ledgersync.jooq.Tables.*;
Expand All @@ -30,6 +31,10 @@ public class CustomAddressTokenBalanceRepository {
@Transactional(readOnly = true, propagation = Propagation.REQUIRES_NEW)
public List<AddressTokenBalance> findAllByAddressFingerprintPairIn(Collection<Pair<String, String>> addressFingerprintPairs) {

if (addressFingerprintPairs.isEmpty()) {
return Collections.emptyList();
}

Condition condition = null;
for (Pair<String, String> addressFingerprintPair : addressFingerprintPairs) {
String address = addressFingerprintPair.getFirst();
Expand All @@ -50,6 +55,10 @@ public List<AddressTokenBalance> findAllByAddressFingerprintPairIn(Collection<Pa

@Transactional(readOnly = true, propagation = Propagation.REQUIRES_NEW)
public List<AddressTokenBalance> findAllByAddressMultiAssetIdPairIn(Collection<Pair<Long, Long>> addressMultiAssetIdPairs) {
if (addressMultiAssetIdPairs.isEmpty()) {
return Collections.emptyList();
}

Condition condition = null;
for (Pair<Long, Long> addressMultiAssetIdPair : addressMultiAssetIdPairs) {
Long addressId = addressMultiAssetIdPair.getFirst();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -150,9 +150,27 @@ public void rollbackAddressBalances(Collection<Tx> txs) {
* to subtract an address's total token balance record
*/
var addressTokens = addressTokenRepository.findAllByTxIn(txs);
var addressMultiAssetIdPairs = addressTokens.stream()
.map(addressToken -> Pair.of(addressToken.getAddressId(), addressToken.getMultiAssetId()))
.collect(Collectors.toSet());

if (!addressTokens.isEmpty()) {
var addressMultiAssetIdPairs = addressTokens.stream()
.map(addressToken -> Pair.of(addressToken.getAddressId(), addressToken.getMultiAssetId()))
.collect(Collectors.toSet());

// Find all address token balance records
var addressTokenBalanceMap = customAddressTokenBalanceRepository
.findAllByAddressMultiAssetIdPairIn(addressMultiAssetIdPairs)
.stream()
.collect(Collectors.toMap(
addressTokenBalance ->
Pair.of(addressTokenBalance.getAddressId(),
addressTokenBalance.getMultiAssetId()),
Function.identity()
));

// Rollback address token balances
rollbackAddressTokenBalances(addressTokens, addressTokenBalanceMap);
addressTokenBalanceRepository.saveAll(addressTokenBalanceMap.values());
}

// Find all address records
var addressMap = addressRepository.findAllByAddressIn(addressSet)
Expand All @@ -164,25 +182,10 @@ public void rollbackAddressBalances(Collection<Tx> txs) {
.stream()
.collect(Collectors.toMap(BaseEntity::getId, Function.identity()));

// Find all address token balance records
var addressTokenBalanceMap = customAddressTokenBalanceRepository
.findAllByAddressMultiAssetIdPairIn(addressMultiAssetIdPairs)
.stream()
.collect(Collectors.toMap(
addressTokenBalance ->
Pair.of(addressTokenBalance.getAddressId(),
addressTokenBalance.getMultiAssetId()),
Function.identity()
));

// Rollback address native balances
rollbackAddressNativeBalances(addressTxBalances, addressMap, stakeAddressMap);

// Rollback address token balances
rollbackAddressTokenBalances(addressTokens, addressTokenBalanceMap);

addressRepository.saveAll(addressMap.values());
addressTokenBalanceRepository.saveAll(addressTokenBalanceMap.values());
stakeAddressRepository.saveAll(stakeAddressMap.values());
log.info("Address balances rollback finished");
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -681,14 +681,11 @@ void shouldRollbackTxsWithNoTokensSuccessfullyTest() {
.thenReturn(new ArrayList<>(addressMap.values()));
Mockito.when(addressRepository.findAllStakeAddressByAddressIn(Mockito.anyCollection()))
.thenReturn(new ArrayList<>(stakeAddressMap.values()));
Mockito.when(customAddressTokenBalanceRepository
.findAllByAddressMultiAssetIdPairIn(Mockito.anyCollection()))
.thenReturn(Collections.emptyList());

Assertions.assertDoesNotThrow(() -> victim.rollbackAddressBalances(txMap.values()));

Mockito.verify(addressRepository, Mockito.times(1)).saveAll(addressesCaptor.capture());
Mockito.verify(addressTokenBalanceRepository, Mockito.times(1)).saveAll(addressTokenBalancesCaptor.capture());
Mockito.verify(addressTokenBalanceRepository, Mockito.never()).saveAll(Mockito.anyCollection());
Mockito.verify(stakeAddressRepository, Mockito.times(1)).saveAll(stakeAddressesCaptor.capture());
Mockito.verify(aggregatedDataCachingService, Mockito.times(4))
.subtractAccountTxCountAtEpoch(Mockito.anyInt(), Mockito.anyString(), Mockito.anyInt());
Expand All @@ -700,9 +697,6 @@ void shouldRollbackTxsWithNoTokensSuccessfullyTest() {
Assertions.assertEquals(0L, address.getTxCount());
});

Collection<AddressTokenBalance> addressTokenBalances = addressTokenBalancesCaptor.getValue();
Assertions.assertTrue(CollectionUtils.isEmpty(addressTokenBalances));

Collection<StakeAddress> stakeAddresses = stakeAddressesCaptor.getValue();
stakeAddresses.forEach(stakeAddress ->
Assertions.assertEquals(BigInteger.ZERO, stakeAddress.getBalance()));
Expand Down

0 comments on commit 853641d

Please sign in to comment.