From 7f4bb96c57f94aa672f59c5e26d795dd95bb45af Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Mon, 14 Sep 2026 10:49:02 -0400 Subject: [PATCH] Implement the GetBalances RPC --- go.mod | 2 +- go.sum | 4 +- ocp/balance/calculator.go | 53 +++++++-- ocp/balance/calculator_test.go | 56 +++++++++ ocp/data/balance/memory/store.go | 34 ++++++ ocp/data/balance/postgres/model.go | 24 ++++ ocp/data/balance/postgres/store.go | 14 +++ ocp/data/balance/store.go | 6 + ocp/data/balance/tests/tests.go | 23 ++++ ocp/data/internal.go | 4 + ocp/rpc/balance/server.go | 175 ++++++++++++++++++++++------- ocp/rpc/balance/server_test.go | 127 +++++++++++++++++++-- 12 files changed, 460 insertions(+), 62 deletions(-) diff --git a/go.mod b/go.mod index 45610b2..a980a77 100644 --- a/go.mod +++ b/go.mod @@ -11,7 +11,7 @@ require ( github.com/aws/aws-sdk-go-v2/service/dynamodb v1.59.0 github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2 github.com/code-payments/code-vm-indexer v1.2.0 - github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54 + github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260914142118-62b4738e3521 github.com/emirpasic/gods v1.12.0 github.com/envoyproxy/protoc-gen-validate v1.3.3 github.com/golang/protobuf v1.5.4 diff --git a/go.sum b/go.sum index faddc84..c4dcbbc 100644 --- a/go.sum +++ b/go.sum @@ -80,8 +80,8 @@ github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ= github.com/code-payments/code-vm-indexer v1.2.0 h1:rSHpBMiT9BKgmKcXg/VIoi/h0t7jNxGx07Qz59m+6Q0= github.com/code-payments/code-vm-indexer v1.2.0/go.mod h1:vn91YN2qNqb+gGJeZe2+l+TNxVmEEiRHXXnIn2Y40h8= -github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54 h1:idVNVzNMZWunI9RlDid6GNto1+z8HVApQhHqmfRIAPc= -github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8= +github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260914142118-62b4738e3521 h1:b6uE5r53BiWXhjJ9UquZzNYsrvLl6nERUdXIuuqYGwQ= +github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260914142118-62b4738e3521/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8= github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 h1:NmTXa/uVnDyp0TY5MKi197+3HWcnYWfnHGyaFthlnGw= github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6/go.mod h1:GL3xCUCBDV3CZiTSEKksMWbLE66hEyuu9qyDOOqM47Y= github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk= diff --git a/ocp/balance/calculator.go b/ocp/balance/calculator.go index 0be8314..61c8acc 100644 --- a/ocp/balance/calculator.go +++ b/ocp/balance/calculator.go @@ -185,21 +185,56 @@ func BatchCalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenA // // Note: Use this method when calculating balances for accounts that are managed by // Code (ie. Timelock account) and operate within the L2 system. -func BatchCalculateFromCacheByOwner(ctx context.Context, data ocp_data.Provider, owner *common.Account) (map[string]*Balance, error) { - tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCacheByOwner") - tracer.AddAttribute("owner", owner.PublicKey().ToBase58()) +func BatchCalculateFromCacheByOwner(ctx context.Context, data ocp_data.Provider, owner *common.Account, mints ...*common.Account) (map[string]*Balance, error) { + balancesByOwner, err := BatchCalculateFromCacheByOwners(ctx, data, []*common.Account{owner}, mints) + if err != nil { + return nil, err + } + + res, ok := balancesByOwner[owner.PublicKey().ToBase58()] + if !ok { + return make(map[string]*Balance), nil + } + return res, nil +} + +// BatchCalculateFromCacheByOwners is BatchCalculateFromCacheByOwner for a set +// of owners in a single ledger read. The result is keyed by owner, then by +// token account. Owners the ledger holds nothing for are omitted. An empty +// mints slice includes every mint, otherwise the read is limited to the +// provided mints. +func BatchCalculateFromCacheByOwners(ctx context.Context, data ocp_data.Provider, owners, mints []*common.Account) (map[string]map[string]*Balance, error) { + tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCacheByOwners") + tracer.AddAttribute("owner_count", len(owners)) + tracer.AddAttribute("mint_count", len(mints)) defer tracer.End() - balanceRecords, err := data.GetAllBalancesByOwner(ctx, owner.PublicKey().ToBase58()) - if err != nil && err != balance.ErrRecordNotFound { + ownerAddresses := make([]string, len(owners)) + for i, owner := range owners { + ownerAddresses[i] = owner.PublicKey().ToBase58() + } + + mintAddresses := make([]string, len(mints)) + for i, mint := range mints { + mintAddresses[i] = mint.PublicKey().ToBase58() + } + + balanceRecordsByOwner, err := data.GetAllBalancesByOwnerBatch(ctx, ownerAddresses, mintAddresses) + if err != nil { tracer.OnError(err) return nil, err } - res := make(map[string]*Balance, len(balanceRecords)) - for _, balanceRecord := range balanceRecords { - if cached, ok := balanceFromRecord(balanceRecord); ok { - res[balanceRecord.TokenAccount] = cached + res := make(map[string]map[string]*Balance, len(balanceRecordsByOwner)) + for owner, balanceRecords := range balanceRecordsByOwner { + balancesByTokenAccount := make(map[string]*Balance, len(balanceRecords)) + for _, balanceRecord := range balanceRecords { + if cached, ok := balanceFromRecord(balanceRecord); ok { + balancesByTokenAccount[balanceRecord.TokenAccount] = cached + } + } + if len(balancesByTokenAccount) > 0 { + res[owner] = balancesByTokenAccount } } return res, nil diff --git a/ocp/balance/calculator_test.go b/ocp/balance/calculator_test.go index 54ad38d..dcd2efc 100644 --- a/ocp/balance/calculator_test.go +++ b/ocp/balance/calculator_test.go @@ -120,6 +120,62 @@ func TestDefaultCalculationMethods_ByOwner(t *testing.T) { assert.Empty(t, balanceByAccount) } +func TestDefaultCalculationMethods_ByOwners(t *testing.T) { + env := setupBalanceTestEnv(t) + + coreVmConfig := testutil.NewRandomVmConfig(t, true) + otherVmConfig := testutil.NewRandomVmConfig(t, false) + + owner1 := testutil.NewRandomAccount(t) + owner2 := testutil.NewRandomAccount(t) + unknownOwner := testutil.NewRandomAccount(t) + + owner1CoreMint := newBalanceTestAccountForOwner(t, env, owner1, coreVmConfig) + owner1OtherMint := newBalanceTestAccountForOwner(t, env, owner1, otherVmConfig) + owner2CoreMint := newBalanceTestAccountForOwner(t, env, owner2, coreVmConfig) + owner2Unlocked := newBalanceTestAccountForOwner(t, env, owner2, otherVmConfig) + + saveBalanceTestRecord(t, env, owner1CoreMint, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, owner1OtherMint, &balance.Record{Quarks: 33, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, owner2CoreMint, &balance.Record{Quarks: 7, IsOpen: true, IsLocked: true}) + saveBalanceTestRecord(t, env, owner2Unlocked, &balance.Record{Quarks: 99, IsOpen: true}) + + // Every owner's managed accounts are reported in a single read. Unlocked + // accounts and owners the ledger holds nothing for are omitted. + balanceByOwnerAndAccount, err := BatchCalculateFromCacheByOwners(env.ctx, env.data, []*common.Account{owner1, owner2, unknownOwner}, nil) + require.NoError(t, err) + assert.Equal(t, map[string]map[string]*Balance{ + owner1.PublicKey().ToBase58(): { + owner1CoreMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner1CoreMint.mint(), Quarks: 42}, + owner1OtherMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner1OtherMint.mint(), Quarks: 33}, + }, + owner2.PublicKey().ToBase58(): { + owner2CoreMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner2CoreMint.mint(), Quarks: 7}, + }, + }, balanceByOwnerAndAccount) + + // A mint filter limits the read to the provided mints + balanceByOwnerAndAccount, err = BatchCalculateFromCacheByOwners(env.ctx, env.data, []*common.Account{owner1, owner2}, []*common.Account{otherVmConfig.Mint}) + require.NoError(t, err) + assert.Equal(t, map[string]map[string]*Balance{ + owner1.PublicKey().ToBase58(): { + owner1OtherMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner1OtherMint.mint(), Quarks: 33}, + }, + }, balanceByOwnerAndAccount) + + // The single owner variant applies the same filter + balanceByAccount, err := BatchCalculateFromCacheByOwner(env.ctx, env.data, owner1, otherVmConfig.Mint) + require.NoError(t, err) + assert.Equal(t, map[string]*Balance{ + owner1OtherMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner1OtherMint.mint(), Quarks: 33}, + }, balanceByAccount) + + // A filter naming a mint nobody holds is empty rather than an error + balanceByOwnerAndAccount, err = BatchCalculateFromCacheByOwners(env.ctx, env.data, []*common.Account{owner1, owner2}, []*common.Account{testutil.NewRandomAccount(t)}) + require.NoError(t, err) + assert.Empty(t, balanceByOwnerAndAccount) +} + func TestDefaultCalculation_ExternalAccount(t *testing.T) { env := setupBalanceTestEnv(t) externalAccount := testutil.NewRandomAccount(t) diff --git a/ocp/data/balance/memory/store.go b/ocp/data/balance/memory/store.go index c7c5794..aba9ca6 100644 --- a/ocp/data/balance/memory/store.go +++ b/ocp/data/balance/memory/store.go @@ -90,6 +90,40 @@ func (s *store) GetAllByOwner(_ context.Context, owner string) ([]*balance.Recor }) } +// GetAllByOwnerBatch implements balance.Store.GetAllByOwnerBatch +func (s *store) GetAllByOwnerBatch(_ context.Context, owners, mints []string) (map[string][]*balance.Record, error) { + s.mu.Lock() + defer s.mu.Unlock() + + ownerSet := make(map[string]struct{}, len(owners)) + for _, owner := range owners { + ownerSet[owner] = struct{}{} + } + + var mintSet map[string]struct{} + if len(mints) > 0 { + mintSet = make(map[string]struct{}, len(mints)) + for _, mint := range mints { + mintSet[mint] = struct{}{} + } + } + + res := make(map[string][]*balance.Record) + for _, item := range s.balanceRecords { + if _, ok := ownerSet[item.OwnerAccount]; !ok { + continue + } + if mintSet != nil { + if _, ok := mintSet[item.MintAccount]; !ok { + continue + } + } + cloned := item.Clone() + res[item.OwnerAccount] = append(res[item.OwnerAccount], &cloned) + } + return res, nil +} + // GetAllByOwnerAndMint implements balance.Store.GetAllByOwnerAndMint func (s *store) GetAllByOwnerAndMint(_ context.Context, owner, mint string) ([]*balance.Record, error) { s.mu.Lock() diff --git a/ocp/data/balance/postgres/model.go b/ocp/data/balance/postgres/model.go index eea77b8..eef39bf 100644 --- a/ocp/data/balance/postgres/model.go +++ b/ocp/data/balance/postgres/model.go @@ -136,6 +136,30 @@ func dbGetBatch(ctx context.Context, db *sqlx.DB, tokenAccounts ...string) ([]*m return res, nil } +func dbGetAllByOwnerBatch(ctx context.Context, db *sqlx.DB, owners, mints []string) ([]*model, error) { + res := []*model{} + if len(owners) == 0 { + return res, nil + } + + query := `SELECT ` + allColumns + ` FROM ` + tableName + ` + WHERE owner_account = ANY($1)` + args := []any{owners} + if len(mints) > 0 { + query += ` AND mint_account = ANY($2)` + args = append(args, mints) + } + query += ` ORDER BY id ASC` + + err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error { + return tx.SelectContext(ctx, &res, query, args...) + }) + if err != nil && !pgutil.IsNoRows(err) { + return nil, err + } + return res, nil +} + func dbGetAllByOwner(ctx context.Context, db *sqlx.DB, owner string, mint *string) ([]*model, error) { res := []*model{} diff --git a/ocp/data/balance/postgres/store.go b/ocp/data/balance/postgres/store.go index 00bbf3d..eb41ad2 100644 --- a/ocp/data/balance/postgres/store.go +++ b/ocp/data/balance/postgres/store.go @@ -68,6 +68,20 @@ func (s *store) GetAllByOwner(ctx context.Context, owner string) ([]*balance.Rec return fromModels(models), nil } +// GetAllByOwnerBatch implements balance.Store.GetAllByOwnerBatch +func (s *store) GetAllByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*balance.Record, error) { + models, err := dbGetAllByOwnerBatch(ctx, s.db, owners, mints) + if err != nil { + return nil, err + } + + res := make(map[string][]*balance.Record) + for _, model := range models { + res[model.OwnerAccount] = append(res[model.OwnerAccount], fromModel(model)) + } + return res, nil +} + // GetAllByOwnerAndMint implements balance.Store.GetAllByOwnerAndMint func (s *store) GetAllByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) { models, err := dbGetAllByOwner(ctx, s.db, owner, &mint) diff --git a/ocp/data/balance/store.go b/ocp/data/balance/store.go index 5b6a8db..c6ac11c 100644 --- a/ocp/data/balance/store.go +++ b/ocp/data/balance/store.go @@ -51,6 +51,12 @@ type Store interface { // ErrRecordNotFound is returned if no records exist. GetAllByOwner(ctx context.Context, owner string) ([]*Record, error) + // GetAllByOwnerBatch gets all balance records for a set of owners, keyed + // by owner. An empty mints slice includes every mint, otherwise only + // records for the provided mints are returned. Owners without records + // are omitted from the result. + GetAllByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*Record, error) + // GetAllByOwnerAndMint gets all balance records for an owner and mint. // // ErrRecordNotFound is returned if no records exist. diff --git a/ocp/data/balance/tests/tests.go b/ocp/data/balance/tests/tests.go index 0ee5985..7018880 100644 --- a/ocp/data/balance/tests/tests.go +++ b/ocp/data/balance/tests/tests.go @@ -45,6 +45,10 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { require.NoError(t, err) assert.Empty(t, batch) + byOwnerBatch, err := s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2"}, nil) + require.NoError(t, err) + assert.Empty(t, byOwnerBatch) + start := time.Now() expected := &balance.Record{ @@ -95,6 +99,25 @@ func testRecordHappyPath(t *testing.T, s balance.Store) { assert.Equal(t, "token_account_1", byOwner[0].TokenAccount) assert.Equal(t, "token_account_2", byOwner[1].TokenAccount) + byOwnerBatch, err = s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2", "owner_3"}, nil) + require.NoError(t, err) + require.Len(t, byOwnerBatch, 2) + require.Len(t, byOwnerBatch["owner_1"], 2) + assertEquivalentRecords(t, &cloned, byOwnerBatch["owner_1"][0]) + assert.Equal(t, "token_account_2", byOwnerBatch["owner_1"][1].TokenAccount) + require.Len(t, byOwnerBatch["owner_2"], 1) + assert.Equal(t, "token_account_3", byOwnerBatch["owner_2"][0].TokenAccount) + + byOwnerBatch, err = s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2"}, []string{"mint_2"}) + require.NoError(t, err) + require.Len(t, byOwnerBatch, 1) + require.Len(t, byOwnerBatch["owner_1"], 1) + assert.Equal(t, "token_account_2", byOwnerBatch["owner_1"][0].TokenAccount) + + byOwnerBatch, err = s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2"}, []string{"mint_3"}) + require.NoError(t, err) + assert.Empty(t, byOwnerBatch) + byOwnerAndMint, err := s.GetAllByOwnerAndMint(ctx, "owner_1", "mint_2") require.NoError(t, err) require.Len(t, byOwnerAndMint, 1) diff --git a/ocp/data/internal.go b/ocp/data/internal.go index 64bc59e..8edecf5 100644 --- a/ocp/data/internal.go +++ b/ocp/data/internal.go @@ -125,6 +125,7 @@ type DatabaseData interface { GetBalance(ctx context.Context, tokenAccount string) (*balance.Record, error) GetBalanceBatch(ctx context.Context, tokenAccounts ...string) (map[string]*balance.Record, error) GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) + GetAllBalancesByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*balance.Record, error) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error) CountLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64) (uint64, error) @@ -477,6 +478,9 @@ func (dp *DatabaseProvider) GetBalanceBatch(ctx context.Context, tokenAccounts . func (dp *DatabaseProvider) GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) { return dp.balance.GetAllByOwner(ctx, owner) } +func (dp *DatabaseProvider) GetAllBalancesByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*balance.Record, error) { + return dp.balance.GetAllByOwnerBatch(ctx, owners, mints) +} func (dp *DatabaseProvider) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) { return dp.balance.GetAllByOwnerAndMint(ctx, owner, mint) } diff --git a/ocp/rpc/balance/server.go b/ocp/rpc/balance/server.go index aba48c0..03ec78a 100644 --- a/ocp/rpc/balance/server.go +++ b/ocp/rpc/balance/server.go @@ -3,11 +3,13 @@ package balance import ( "context" + "github.com/pkg/errors" "go.uber.org/zap" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" balancepb "github.com/code-payments/ocp-protobuf-api/generated/go/balance/v1" + commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1" "github.com/code-payments/ocp-server/grpc/client" "github.com/code-payments/ocp-server/ocp/balance" @@ -47,72 +49,165 @@ func (s *server) GetBalance(ctx context.Context, req *balancepb.GetBalanceReques } log = log.With(zap.String("owner_account", owner.PublicKey().ToBase58())) - var mintFilter map[string]struct{} - if len(req.Mints) > 0 { - mintFilter = make(map[string]struct{}, len(req.Mints)) - for i, protoMint := range req.Mints { - mint, err := common.NewAccountFromProto(protoMint) - if err != nil { - log.With(zap.Error(err), zap.Int("index", i)).Warn("invalid mint account") - return nil, status.Error(codes.Internal, "") - } - mintFilter[mint.PublicKey().ToBase58()] = struct{}{} - } + mints, err := newMintFilter(req.Mints) + if err != nil { + log.With(zap.Error(err)).Warn("invalid mint account") + return nil, status.Error(codes.Internal, "") + } + + // The ledger holds a record for every account Code manages for the owner, + // and each carries the mint it holds. Accounts that have left the L2 system + // don't have a cached balance that can be trusted, so it omits them. The + // mint filter is applied at the ledger read. + balanceByTokenAccount, err := balance.BatchCalculateFromCacheByOwner(ctx, s.data, owner, mints...) + if err != nil { + log.With(zap.Error(err)).Warn("failure getting cached balances") + return nil, status.Error(codes.Internal, "") + } + + ownerBalance, err := s.valueOwnerBalance(ctx, owner, balanceByTokenAccount, newReserveStateCache()) + if err != nil { + log.With(zap.Error(err)).Warn("failure valuing owner balance") + return nil, status.Error(codes.Internal, "") } - ownerMetadata, err := common.GetOwnerMetadata(ctx, s.data, owner) - if err == common.ErrOwnerNotFound { - return &balancepb.GetBalanceResponse{ - Result: balancepb.GetBalanceResponse_NOT_FOUND, - }, nil - } else if err != nil { - log.With(zap.Error(err)).Warn("failure getting owner metadata") + return &balancepb.GetBalanceResponse{ + Result: balancepb.GetBalanceResponse_OK, + CoreMintValue: ownerBalance.CoreMintValue, + BalancesByMint: ownerBalance.BalancesByMint, + }, nil +} + +func (s *server) GetBalances(ctx context.Context, req *balancepb.GetBalancesRequest) (*balancepb.GetBalancesResponse, error) { + log := s.log.With(zap.String("method", "GetBalances")) + log = client.InjectLoggingMetadata(ctx, log, rpc.UserAgentName) + + mints, err := newMintFilter(req.Mints) + if err != nil { + log.With(zap.Error(err)).Warn("invalid mint account") return nil, status.Error(codes.Internal, "") } - if ownerMetadata.Type != common.OwnerTypeUser12Words { - return &balancepb.GetBalanceResponse{ - Result: balancepb.GetBalanceResponse_NOT_FOUND, - }, nil + // Duplicate owners collapse to a single entry + seenOwners := make(map[string]struct{}, len(req.Owners)) + owners := make([]*common.Account, 0, len(req.Owners)) + for i, protoOwner := range req.Owners { + owner, err := common.NewAccountFromProto(protoOwner) + if err != nil { + log.With(zap.Error(err), zap.Int("index", i)).Warn("invalid owner account") + return nil, status.Error(codes.Internal, "") + } + + if _, ok := seenOwners[owner.PublicKey().ToBase58()]; ok { + continue + } + seenOwners[owner.PublicKey().ToBase58()] = struct{}{} + owners = append(owners, owner) } - balancesByMint, err := s.calculateCoreMintValueByMint(ctx, owner, mintFilter) + // All owners' ledger records are read in a single batch, with the mint + // filter applied at the ledger read. Owners the ledger holds nothing for + // are absent from the result, and are reported with an empty balance. + balanceByOwnerAndTokenAccount, err := balance.BatchCalculateFromCacheByOwners(ctx, s.data, owners, mints) if err != nil { - log.With(zap.Error(err)).Warn("failure calculating core mint value") + log.With(zap.Error(err)).Warn("failure getting cached balances") return nil, status.Error(codes.Internal, "") } + // A single reserve state cache values every owner's holdings in a mint + // against the same supply, even if the mint data provider refreshes + // mid-request. + reserveStateCache := newReserveStateCache() + + balancesByOwner := make(map[string]*balancepb.OwnerBalance, len(owners)) + for _, owner := range owners { + ownerBalance, err := s.valueOwnerBalance(ctx, owner, balanceByOwnerAndTokenAccount[owner.PublicKey().ToBase58()], reserveStateCache) + if err != nil { + log.With(zap.Error(err), zap.String("owner_account", owner.PublicKey().ToBase58())).Warn("failure valuing owner balance") + return nil, status.Error(codes.Internal, "") + } + + balancesByOwner[owner.PublicKey().ToBase58()] = ownerBalance + } + + return &balancepb.GetBalancesResponse{ + Result: balancepb.GetBalancesResponse_OK, + BalancesByOwner: balancesByOwner, + }, nil +} + +// valueOwnerBalance values an owner's cached holdings. Owner metadata checks +// are intentionally skipped as an optimization, so an owner with no ledger +// records has an empty balance. +func (s *server) valueOwnerBalance(ctx context.Context, owner *common.Account, balanceByTokenAccount map[string]*balance.Balance, reserveStateCache reserveStateCache) (*balancepb.OwnerBalance, error) { + balancesByMint, err := s.calculateCoreMintValueByMint(ctx, balanceByTokenAccount, reserveStateCache) + if err != nil { + return nil, err + } + var totalCoreMintValue uint64 for _, mintBalance := range balancesByMint { totalCoreMintValue += mintBalance.CoreMintValue } - return &balancepb.GetBalanceResponse{ - Result: balancepb.GetBalanceResponse_OK, + return &balancepb.OwnerBalance{ + Owner: owner.ToProto(), CoreMintValue: totalCoreMintValue, BalancesByMint: balancesByMint, }, nil } -// calculateCoreMintValueByMint values the owner's holdings in each mint they -// hold a non-zero balance of. A nil mintFilter includes every mint, otherwise -// only mints in the filter are included. -func (s *server) calculateCoreMintValueByMint(ctx context.Context, owner *common.Account, mintFilter map[string]struct{}) (map[string]*balancepb.MintBalance, error) { - // The ledger holds a record for every account Code manages for the owner, - // and each carries the mint it holds. Accounts that have left the L2 system - // don't have a cached balance that can be trusted, so it omits them. - balanceByTokenAccount, err := balance.BatchCalculateFromCacheByOwner(ctx, s.data, owner) +// newMintFilter converts the request's mints into a deduplicated list for +// filtering the ledger read. A nil filter is returned when no mints are +// provided, which includes every mint. +func newMintFilter(protoMints []*commonpb.SolanaAccountId) ([]*common.Account, error) { + if len(protoMints) == 0 { + return nil, nil + } + + seen := make(map[string]struct{}, len(protoMints)) + mints := make([]*common.Account, 0, len(protoMints)) + for i, protoMint := range protoMints { + mint, err := common.NewAccountFromProto(protoMint) + if err != nil { + return nil, errors.Wrapf(err, "invalid mint account at index %d", i) + } + + if _, ok := seen[mint.PublicKey().ToBase58()]; ok { + continue + } + seen[mint.PublicKey().ToBase58()] = struct{}{} + mints = append(mints, mint) + } + return mints, nil +} + +// reserveStateCache pins the live reserve state observed for each launchpad +// mint so every valuation within a single RPC uses the same supply. +type reserveStateCache map[string]*currency_util.LiveReserveStateData + +func newReserveStateCache() reserveStateCache { + return make(reserveStateCache) +} + +func (s *server) getLiveReserveState(ctx context.Context, mint *common.Account, cache reserveStateCache) (*currency_util.LiveReserveStateData, error) { + if reserveState, ok := cache[mint.PublicKey().ToBase58()]; ok { + return reserveState, nil + } + + reserveState, err := s.mintDataProvider.GetLiveReserveState(ctx, mint) if err != nil { return nil, err } + cache[mint.PublicKey().ToBase58()] = reserveState + return reserveState, nil +} +// calculateCoreMintValueByMint values cached holdings in each mint with a +// non-zero balance. Each balance carries the mint it holds. +func (s *server) calculateCoreMintValueByMint(ctx context.Context, balanceByTokenAccount map[string]*balance.Balance, reserveStateCache reserveStateCache) (map[string]*balancepb.MintBalance, error) { quarksByMint := make(map[string]uint64) for _, cached := range balanceByTokenAccount { - if mintFilter != nil { - if _, ok := mintFilter[cached.MintAccount]; !ok { - continue - } - } quarksByMint[cached.MintAccount] += cached.Quarks } @@ -131,7 +226,7 @@ func (s *server) calculateCoreMintValueByMint(ctx context.Context, owner *common if mint == common.CoreMintAccount.PublicKey().ToBase58() { coreMintValue = quarks } else { - reserveState, err := s.mintDataProvider.GetLiveReserveState(ctx, mintAccount) + reserveState, err := s.getLiveReserveState(ctx, mintAccount, reserveStateCache) if err != nil { return nil, err } diff --git a/ocp/rpc/balance/server_test.go b/ocp/rpc/balance/server_test.go index 66712f6..496a3ed 100644 --- a/ocp/rpc/balance/server_test.go +++ b/ocp/rpc/balance/server_test.go @@ -87,8 +87,9 @@ func TestGetBalance_HappyPath(t *testing.T) { resp, err := env.client.GetBalance(env.ctx, req) require.NoError(t, err) - assert.Equal(t, balancepb.GetBalanceResponse_NOT_FOUND, resp.Result) + assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) assert.EqualValues(t, 0, resp.CoreMintValue) + assert.Empty(t, resp.BalancesByMint) primaryCoreMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, coreVmConfig, 0, commonpb.AccountType_PRIMARY) primaryLaunchpadMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, launchpadVmConfig, 0, commonpb.AccountType_PRIMARY) @@ -117,8 +118,8 @@ func TestGetBalance_HappyPath(t *testing.T) { assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) assert.EqualValues(t, common.ToCoreMintQuarks(42)+expectedLaunchpadMintValue, resp.CoreMintValue) require.Len(t, resp.BalancesByMint, 2) - assertMintBalance(t, resp, common.CoreMintAccount, common.ToCoreMintQuarks(42)) - assertMintBalance(t, resp, launchpadMint, expectedLaunchpadMintValue) + assertMintBalance(t, resp.BalancesByMint, common.CoreMintAccount, common.ToCoreMintQuarks(42)) + assertMintBalance(t, resp.BalancesByMint, launchpadMint, expectedLaunchpadMintValue) } func TestGetBalance_MintFilter(t *testing.T) { @@ -156,7 +157,7 @@ func TestGetBalance_MintFilter(t *testing.T) { assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) assert.EqualValues(t, common.ToCoreMintQuarks(42), resp.CoreMintValue) require.Len(t, resp.BalancesByMint, 1) - assertMintBalance(t, resp, common.CoreMintAccount, common.ToCoreMintQuarks(42)) + assertMintBalance(t, resp.BalancesByMint, common.CoreMintAccount, common.ToCoreMintQuarks(42)) resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{ Owner: ownerAccount.ToProto(), @@ -166,7 +167,7 @@ func TestGetBalance_MintFilter(t *testing.T) { assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) assert.EqualValues(t, expectedLaunchpadMintValue, resp.CoreMintValue) require.Len(t, resp.BalancesByMint, 1) - assertMintBalance(t, resp, launchpadMint, expectedLaunchpadMintValue) + assertMintBalance(t, resp.BalancesByMint, launchpadMint, expectedLaunchpadMintValue) // Duplicate mints in the filter don't double count resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{ @@ -233,8 +234,10 @@ func TestGetBalance_GiftCardOwnerAccount(t *testing.T) { Owner: giftCardOwnerAccount.ToProto(), }) require.NoError(t, err) - assert.Equal(t, balancepb.GetBalanceResponse_NOT_FOUND, resp.Result) - assert.EqualValues(t, 0, resp.CoreMintValue) + assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) + assert.EqualValues(t, common.ToCoreMintQuarks(42), resp.CoreMintValue) + require.Len(t, resp.BalancesByMint, 1) + assertMintBalance(t, resp.BalancesByMint, common.CoreMintAccount, common.ToCoreMintQuarks(42)) } func TestGetBalance_UnknownOwnerAccount(t *testing.T) { @@ -245,8 +248,104 @@ func TestGetBalance_UnknownOwnerAccount(t *testing.T) { Owner: testutil.NewRandomAccount(t).ToProto(), }) require.NoError(t, err) - assert.Equal(t, balancepb.GetBalanceResponse_NOT_FOUND, resp.Result) + assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) assert.EqualValues(t, 0, resp.CoreMintValue) + assert.Empty(t, resp.BalancesByMint) +} + +func TestGetBalances_HappyPath(t *testing.T) { + env, cleanup := setup(t) + defer cleanup() + + coreVmConfig := testutil.NewRandomVmConfig(t, true) + launchpadMint := testutil.SetupLaunchpadCurrency(t, env.data, env.reserveStore, env.holderStore) + launchpadVmConfig, err := common.GetVmConfigForMint(env.ctx, env.data, launchpadMint) + require.NoError(t, err) + + ownerAccount1 := testutil.NewRandomAccount(t) + ownerAccount2 := testutil.NewRandomAccount(t) + unknownOwnerAccount := testutil.NewRandomAccount(t) + giftCardOwnerAccount := testutil.NewRandomAccount(t) + + owner1CoreMintAccountRecords := setupAccountRecords(t, env, ownerAccount1, ownerAccount1, coreVmConfig, 0, commonpb.AccountType_PRIMARY) + owner1LaunchpadMintAccountRecords := setupAccountRecords(t, env, ownerAccount1, ownerAccount1, launchpadVmConfig, 0, commonpb.AccountType_PRIMARY) + owner2CoreMintAccountRecords := setupAccountRecords(t, env, ownerAccount2, ownerAccount2, coreVmConfig, 0, commonpb.AccountType_PRIMARY) + giftCardAccountRecords := setupAccountRecords(t, env, giftCardOwnerAccount, giftCardOwnerAccount, coreVmConfig, 0, commonpb.AccountType_REMOTE_SEND_GIFT_CARD) + + setupCachedBalance(t, env, owner1CoreMintAccountRecords, common.ToCoreMintQuarks(42)) + setupCachedBalance(t, env, owner1LaunchpadMintAccountRecords, currencycreator.ToQuarks(100)) + setupCachedBalance(t, env, owner2CoreMintAccountRecords, common.ToCoreMintQuarks(7)) + setupCachedBalance(t, env, giftCardAccountRecords, common.ToCoreMintQuarks(1)) + + expectedLaunchpadMintValue, _ := currencycreator.EstimateSell(¤cycreator.EstimateSellArgs{ + CurrentSupplyInQuarks: currencycreator.ToQuarks(1_000), + SellAmountInQuarks: currencycreator.ToQuarks(100), + ValueMintDecimals: uint8(common.CoreMintDecimals), + SellFeeBps: 0, + }) + require.NotZero(t, expectedLaunchpadMintValue) + + resp, err := env.client.GetBalances(env.ctx, &balancepb.GetBalancesRequest{ + Owners: []*commonpb.SolanaAccountId{ + ownerAccount1.ToProto(), + ownerAccount2.ToProto(), + unknownOwnerAccount.ToProto(), + giftCardOwnerAccount.ToProto(), + ownerAccount1.ToProto(), + }, + }) + require.NoError(t, err) + assert.Equal(t, balancepb.GetBalancesResponse_OK, resp.Result) + require.Len(t, resp.BalancesByOwner, 4) + + ownerBalance1 := assertOwnerBalance(t, resp, ownerAccount1, common.ToCoreMintQuarks(42)+expectedLaunchpadMintValue) + require.Len(t, ownerBalance1.BalancesByMint, 2) + assertMintBalance(t, ownerBalance1.BalancesByMint, common.CoreMintAccount, common.ToCoreMintQuarks(42)) + assertMintBalance(t, ownerBalance1.BalancesByMint, launchpadMint, expectedLaunchpadMintValue) + + ownerBalance2 := assertOwnerBalance(t, resp, ownerAccount2, common.ToCoreMintQuarks(7)) + require.Len(t, ownerBalance2.BalancesByMint, 1) + assertMintBalance(t, ownerBalance2.BalancesByMint, common.CoreMintAccount, common.ToCoreMintQuarks(7)) + + unknownOwnerBalance := assertOwnerBalance(t, resp, unknownOwnerAccount, 0) + assert.Empty(t, unknownOwnerBalance.BalancesByMint) + + giftCardOwnerBalance := assertOwnerBalance(t, resp, giftCardOwnerAccount, common.ToCoreMintQuarks(1)) + require.Len(t, giftCardOwnerBalance.BalancesByMint, 1) + assertMintBalance(t, giftCardOwnerBalance.BalancesByMint, common.CoreMintAccount, common.ToCoreMintQuarks(1)) + + // The mint filter applies to every owner + resp, err = env.client.GetBalances(env.ctx, &balancepb.GetBalancesRequest{ + Owners: []*commonpb.SolanaAccountId{ownerAccount1.ToProto(), ownerAccount2.ToProto()}, + Mints: []*commonpb.SolanaAccountId{launchpadMint.ToProto()}, + }) + require.NoError(t, err) + assert.Equal(t, balancepb.GetBalancesResponse_OK, resp.Result) + require.Len(t, resp.BalancesByOwner, 2) + + ownerBalance1 = assertOwnerBalance(t, resp, ownerAccount1, expectedLaunchpadMintValue) + require.Len(t, ownerBalance1.BalancesByMint, 1) + assertMintBalance(t, ownerBalance1.BalancesByMint, launchpadMint, expectedLaunchpadMintValue) + + ownerBalance2 = assertOwnerBalance(t, resp, ownerAccount2, 0) + assert.Empty(t, ownerBalance2.BalancesByMint) +} + +func TestGetBalances_UnknownOwnerAccount(t *testing.T) { + env, cleanup := setup(t) + defer cleanup() + + unknownOwnerAccount := testutil.NewRandomAccount(t) + + resp, err := env.client.GetBalances(env.ctx, &balancepb.GetBalancesRequest{ + Owners: []*commonpb.SolanaAccountId{unknownOwnerAccount.ToProto()}, + }) + require.NoError(t, err) + assert.Equal(t, balancepb.GetBalancesResponse_OK, resp.Result) + require.Len(t, resp.BalancesByOwner, 1) + + unknownOwnerBalance := assertOwnerBalance(t, resp, unknownOwnerAccount, 0) + assert.Empty(t, unknownOwnerBalance.BalancesByMint) } func setupAccountRecords(t *testing.T, env testEnv, ownerAccount, authorityAccount *common.Account, vmConfig *common.VmConfig, index uint64, accountType commonpb.AccountType) *common.AccountRecords { @@ -278,8 +377,16 @@ func setupAccountRecords(t *testing.T, env testEnv, ownerAccount, authorityAccou } } -func assertMintBalance(t *testing.T, resp *balancepb.GetBalanceResponse, mint *common.Account, expectedCoreMintValue uint64) { - mintBalance, ok := resp.BalancesByMint[mint.PublicKey().ToBase58()] +func assertOwnerBalance(t *testing.T, resp *balancepb.GetBalancesResponse, owner *common.Account, expectedCoreMintValue uint64) *balancepb.OwnerBalance { + ownerBalance, ok := resp.BalancesByOwner[owner.PublicKey().ToBase58()] + require.True(t, ok) + assert.Equal(t, owner.PublicKey().ToBytes(), ownerBalance.Owner.Value) + assert.EqualValues(t, expectedCoreMintValue, ownerBalance.CoreMintValue) + return ownerBalance +} + +func assertMintBalance(t *testing.T, balancesByMint map[string]*balancepb.MintBalance, mint *common.Account, expectedCoreMintValue uint64) { + mintBalance, ok := balancesByMint[mint.PublicKey().ToBase58()] require.True(t, ok) assert.Equal(t, mint.PublicKey().ToBytes(), mintBalance.Mint.Value) assert.EqualValues(t, expectedCoreMintValue, mintBalance.CoreMintValue)