Improve CalculateFromBlockchain (#47)

Author: jeffyanta

pkg/code/balance/calculator.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/code-payments/code-server/pkg/code/common"
 	code_data "github.com/code-payments/code-server/pkg/code/data"
+	"github.com/code-payments/code-server/pkg/code/data/balance"
 	"github.com/code-payments/code-server/pkg/code/data/timelock"
 	"github.com/code-payments/code-server/pkg/metrics"
 	"github.com/code-payments/code-server/pkg/solana"
@@ -109,14 +110,41 @@ func CalculateFromCache(ctx context.Context, data code_data.Provider, tokenAccou
 //
 // todo: add a batching variant
 func CalculateFromBlockchain(ctx context.Context, data code_data.Provider, tokenAccount *common.Account) (uint64, error) {
+	var cachedQuarks uint64
+	var cachedSlot uint64
+	checkpointRecord, err := data.GetBalanceCheckpoint(ctx, tokenAccount.PublicKey().ToBase58())
+	if err == nil {
+		cachedQuarks = checkpointRecord.Quarks
+		cachedSlot = checkpointRecord.SlotCheckpoint
+	} else if err != balance.ErrCheckpointNotFound {
+		return 0, err
+	}
+
 	// todo: we may need something that's more resistant to RPC nodes with stale account state
-	balance, err := data.GetBlockchainBalance(ctx, tokenAccount.PublicKey().ToBase58())
+	quarks, slot, err := data.GetBlockchainBalance(ctx, tokenAccount.PublicKey().ToBase58())
 	if err == solana.ErrNoBalance {
 		return 0, nil
 	} else if err != nil {
-		return 0, err
+		// RPC node threw an error. Return the cached balance
+		return cachedQuarks, nil
 	}
-	return balance, nil
+
+	// RPC node is behind, use cached balance
+	if cachedSlot > slot {
+		return cachedQuarks, nil
+	}
+
+	// Observed a balance that's more recent. Best-effort update the checkpoint.
+	if cachedSlot == 0 || (slot > cachedSlot && quarks != cachedQuarks) {
+		newCheckpointRecord := &balance.Record{
+			TokenAccount:   tokenAccount.PublicKey().ToBase58(),
+			Quarks:         quarks,
+			SlotCheckpoint: slot,
+		}
+		data.SaveBalanceCheckpoint(ctx, newCheckpointRecord)
+	}
+
+	return quarks, nil
 }
 
 // Calculate calculates a token account's balance using a starting point and a set

pkg/code/data/blockchain.go

@@ -23,7 +23,7 @@ type BlockchainData interface {
 
 	GetBlockchainAccountInfo(ctx context.Context, account string, commitment solana.Commitment) (*solana.AccountInfo, error)
 	GetBlockchainAccountDataAfterBlock(ctx context.Context, account string, slot uint64) ([]byte, uint64, error)
-	GetBlockchainBalance(ctx context.Context, account string) (uint64, error)
+	GetBlockchainBalance(ctx context.Context, account string) (uint64, uint64, error)
 	GetBlockchainBlock(ctx context.Context, slot uint64) (*solana.Block, error)
 	GetBlockchainBlockSignatures(ctx context.Context, slot uint64) ([]string, error)
 	GetBlockchainBlocksWithLimit(ctx context.Context, start uint64, limit uint64) ([]uint64, error)
@@ -284,21 +284,21 @@ func (dp *BlockchainProvider) GetBlockchainMinimumBalanceForRentExemption(ctx co
 	return res, err
 }
 
-func (dp *BlockchainProvider) GetBlockchainBalance(ctx context.Context, account string) (uint64, error) {
+func (dp *BlockchainProvider) GetBlockchainBalance(ctx context.Context, account string) (uint64, uint64, error) {
 	tracer := metrics.TraceMethodCall(ctx, blockchainProviderMetricsName, "GetBlockchainBalance")
 	defer tracer.End()
 
 	accountId, err := base58.Decode(account)
 	if err != nil {
-		return 0, err
+		return 0, 0, err
 	}
 
-	res, err := dp.sc.GetTokenAccountBalance(accountId)
+	quarks, slot, err := dp.sc.GetTokenAccountBalance(accountId)
 
 	if err != nil {
 		tracer.OnError(err)
 	}
-	return res, err
+	return quarks, slot, err
 }
 
 func (dp *BlockchainProvider) GetBlockchainTransactionTokenBalances(ctx context.Context, sig string) (*solana.TransactionTokenBalances, error) {

pkg/solana/client.go

@@ -184,7 +184,7 @@ type Client interface {
 	GetSignatureStatuses([]Signature) ([]*SignatureStatus, error)
 	GetSignaturesForAddress(owner ed25519.PublicKey, commitment Commitment, limit uint64, before, until string) ([]*TransactionSignature, error)
 	GetSlot(Commitment) (uint64, error)
-	GetTokenAccountBalance(ed25519.PublicKey) (uint64, error)
+	GetTokenAccountBalance(ed25519.PublicKey) (uint64, uint64, error)
 	GetTokenAccountsByOwner(owner, mint ed25519.PublicKey) ([]ed25519.PublicKey, error)
 	GetTransaction(Signature, Commitment) (ConfirmedTransaction, error)
 	GetTransactionTokenBalances(Signature) (TransactionTokenBalances, error)
@@ -691,7 +691,7 @@ func (c *client) GetBalance(account ed25519.PublicKey) (uint64, error) {
 	return 0, errors.Errorf("invalid value in response")
 }
 
-func (c *client) GetTokenAccountBalance(account ed25519.PublicKey) (uint64, error) {
+func (c *client) GetTokenAccountBalance(account ed25519.PublicKey) (uint64, uint64, error) {
 	var resp struct {
 		Context struct {
 			Slot int64 `json:"slot"`
@@ -701,22 +701,22 @@ func (c *client) GetTokenAccountBalance(account ed25519.PublicKey) (uint64, erro
 	if err := c.call(&resp, "getTokenAccountBalance", base58.Encode(account[:]), CommitmentProcessed); err != nil {
 		jsonRPCErr, ok := err.(*jsonrpc.RPCError)
 		if !ok {
-			return 0, errors.Wrapf(err, "getTokenAccountBalance() failed to send request")
+			return 0, 0, errors.Wrapf(err, "getTokenAccountBalance() failed to send request")
 		}
 
 		if jsonRPCErr.Code == invalidParamCode {
-			return 0, ErrNoBalance
+			return 0, 0, ErrNoBalance
 		}
 
-		return 0, errors.Wrapf(err, "getTokenAccountBalance() failed to send request")
+		return 0, 0, errors.Wrapf(err, "getTokenAccountBalance() failed to send request")
 	}
 
 	quarks, err := strconv.ParseUint(resp.Value.Amount, 10, 64)
 	if err != nil {
-		return 0, errors.Errorf("invalid value in response")
+		return 0, 0, errors.Errorf("invalid value in response")
 	}
 
-	return quarks, nil
+	return quarks, uint64(resp.Context.Slot), nil
 }
 
 func (c *client) SubmitTransaction(txn Transaction, commitment Commitment) (Signature, error) {