From 577b1a2f825a4b9b631330d94ab2ddd22abe3aff Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Fri, 11 Sep 2026 12:07:41 -0400 Subject: [PATCH] GetBalance now returns per-mint balances with the ability to filter --- go.mod | 2 +- go.sum | 4 +- ocp/rpc/balance/server.go | 77 +++++++++++++++++++++++---------- ocp/rpc/balance/server_test.go | 79 ++++++++++++++++++++++++++++++++++ 4 files changed, 136 insertions(+), 26 deletions(-) diff --git a/go.mod b/go.mod index 0f39fb4..45610b2 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.0 + github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54 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 a7fe8e0..faddc84 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.0 h1:j8Z5TRsNd8b02mCjJSq8+e1CoRh0JCnRVDldTlb9+Ds= -github.com/code-payments/ocp-protobuf-api v1.16.0/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8= +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/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/rpc/balance/server.go b/ocp/rpc/balance/server.go index b383252..9f78964 100644 --- a/ocp/rpc/balance/server.go +++ b/ocp/rpc/balance/server.go @@ -45,6 +45,19 @@ 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{}{} + } + } + ownerMetadata, err := common.GetOwnerMetadata(ctx, s.data, owner) if err == common.ErrOwnerNotFound { return &balancepb.GetBalanceResponse{ @@ -61,61 +74,79 @@ func (s *server) GetBalance(ctx context.Context, req *balancepb.GetBalanceReques }, nil } - coreMintValue, err := s.calculateCoreMintValue(ctx, owner) + balancesByMint, err := s.calculateCoreMintValueByMint(ctx, owner, mintFilter) if err != nil { log.With(zap.Error(err)).Warn("failure calculating core mint value") return nil, status.Error(codes.Internal, "") } + var totalCoreMintValue uint64 + for _, mintBalance := range balancesByMint { + totalCoreMintValue += mintBalance.CoreMintValue + } + return &balancepb.GetBalanceResponse{ - Result: balancepb.GetBalanceResponse_OK, - CoreMintValue: coreMintValue, + Result: balancepb.GetBalanceResponse_OK, + CoreMintValue: totalCoreMintValue, + BalancesByMint: balancesByMint, }, nil } -func (s *server) calculateCoreMintValue(ctx context.Context, owner *common.Account) (uint64, error) { +// 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) if err != nil { - return 0, err + return nil, err } quarksByMint := make(map[string]uint64) for _, cached := range balanceByTokenAccount { + if mintFilter != nil { + if _, ok := mintFilter[cached.MintAccount]; !ok { + continue + } + } quarksByMint[cached.MintAccount] += cached.Quarks } - var coreMintValue uint64 + balancesByMint := make(map[string]*balancepb.MintBalance) for mint, quarks := range quarksByMint { if quarks == 0 { continue } - if mint == common.CoreMintAccount.PublicKey().ToBase58() { - coreMintValue += quarks - continue - } - mintAccount, err := common.NewAccountFromPublicKeyString(mint) if err != nil { - return 0, err + return nil, err } - reserveState, err := s.mintDataProvider.GetLiveReserveState(ctx, mintAccount) - if err != nil { - return 0, err + var coreMintValue uint64 + if mint == common.CoreMintAccount.PublicKey().ToBase58() { + coreMintValue = quarks + } else { + reserveState, err := s.mintDataProvider.GetLiveReserveState(ctx, mintAccount) + if err != nil { + return nil, err + } + + coreMintValue, _ = currencycreator.EstimateSell(¤cycreator.EstimateSellArgs{ + CurrentSupplyInQuarks: reserveState.SupplyFromBonding, + SellAmountInQuarks: quarks, + ValueMintDecimals: uint8(common.CoreMintDecimals), + SellFeeBps: 0, + }) } - sellValue, _ := currencycreator.EstimateSell(¤cycreator.EstimateSellArgs{ - CurrentSupplyInQuarks: reserveState.SupplyFromBonding, - SellAmountInQuarks: quarks, - ValueMintDecimals: uint8(common.CoreMintDecimals), - SellFeeBps: 0, - }) - coreMintValue += sellValue + balancesByMint[mint] = &balancepb.MintBalance{ + Mint: mintAccount.ToProto(), + CoreMintValue: coreMintValue, + } } - return coreMintValue, nil + return balancesByMint, nil } diff --git a/ocp/rpc/balance/server_test.go b/ocp/rpc/balance/server_test.go index 26abe2a..d1e77c9 100644 --- a/ocp/rpc/balance/server_test.go +++ b/ocp/rpc/balance/server_test.go @@ -97,6 +97,7 @@ func TestGetBalance_HappyPath(t *testing.T) { require.NoError(t, err) assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) assert.EqualValues(t, 0, resp.CoreMintValue) + assert.Empty(t, resp.BalancesByMint) setupCachedBalance(t, env, primaryCoreMintAccountRecords, common.ToCoreMintQuarks(42)) setupCachedBalance(t, env, primaryLaunchpadMintAccountRecords, currencycreator.ToQuarks(100)) @@ -115,6 +116,77 @@ func TestGetBalance_HappyPath(t *testing.T) { require.NoError(t, err) 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) +} + +func TestGetBalance_MintFilter(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) + + ownerAccount := testutil.NewRandomAccount(t) + + primaryCoreMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, coreVmConfig, 0, commonpb.AccountType_PRIMARY) + primaryLaunchpadMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, launchpadVmConfig, 0, commonpb.AccountType_PRIMARY) + + setupCachedBalance(t, env, primaryCoreMintAccountRecords, common.ToCoreMintQuarks(42)) + setupCachedBalance(t, env, primaryLaunchpadMintAccountRecords, currencycreator.ToQuarks(100)) + + expectedLaunchpadMintValue, _ := currencycreator.EstimateSell(¤cycreator.EstimateSellArgs{ + CurrentSupplyInQuarks: currencycreator.ToQuarks(1_000), + SellAmountInQuarks: currencycreator.ToQuarks(100), + ValueMintDecimals: uint8(common.CoreMintDecimals), + SellFeeBps: 0, + }) + require.NotZero(t, expectedLaunchpadMintValue) + + // Filtering to the core mint excludes the launchpad currency from both the + // total and the per-mint breakdown + resp, err := env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{ + Owner: ownerAccount.ToProto(), + Mints: []*commonpb.SolanaAccountId{common.CoreMintAccount.ToProto()}, + }) + require.NoError(t, err) + 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)) + + resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{ + Owner: ownerAccount.ToProto(), + Mints: []*commonpb.SolanaAccountId{launchpadMint.ToProto()}, + }) + require.NoError(t, err) + 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) + + // Duplicate mints in the filter don't double count + resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{ + Owner: ownerAccount.ToProto(), + Mints: []*commonpb.SolanaAccountId{common.CoreMintAccount.ToProto(), launchpadMint.ToProto(), common.CoreMintAccount.ToProto()}, + }) + require.NoError(t, err) + assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) + assert.EqualValues(t, common.ToCoreMintQuarks(42)+expectedLaunchpadMintValue, resp.CoreMintValue) + require.Len(t, resp.BalancesByMint, 2) + + // A mint the owner doesn't hold results in an OK response with no balances + resp, err = env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{ + Owner: ownerAccount.ToProto(), + Mints: []*commonpb.SolanaAccountId{testutil.NewRandomAccount(t).ToProto()}, + }) + require.NoError(t, err) + assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result) + assert.EqualValues(t, 0, resp.CoreMintValue) + assert.Empty(t, resp.BalancesByMint) } func TestGetBalance_UnmanagedAccountsExcluded(t *testing.T) { @@ -206,6 +278,13 @@ 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()] + require.True(t, ok) + assert.Equal(t, mint.PublicKey().ToBytes(), mintBalance.Mint.Value) + assert.EqualValues(t, expectedCoreMintValue, mintBalance.CoreMintValue) +} + func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, quarks uint64) { require.NoError(t, balance_util.ApplyDeltasInTx(env.ctx, env.data, &balance.Delta{ TokenAccount: accountRecords.General.TokenAccount,