Skip to content

Commit 9d4321a

Browse files
committed
feat(rpc): implement getBalance RPC method
- Add getBalance handler with commitment level support in AccountsDB - Add getSlotForCommitment to SlotTracker for commitment-based slot resolution - Update registerRPCHooks signature to accept slot_tracker and snapshot_slot - Use snapshot slot as fallback when no live slot tracker is available
1 parent cd39802 commit 9d4321a

File tree

3 files changed

+62
-5
lines changed

3 files changed

+62
-5
lines changed

src/accountsdb/db.zig

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ const tracy = @import("tracy");
88

99
const sysvar = sig.runtime.sysvar;
1010
const snapgen = sig.accounts_db.snapshot.data.generate;
11+
const methods = sig.rpc.methods;
1112

1213
const Resolution = @import("../benchmarks.zig").Resolution;
1314

@@ -2873,9 +2874,64 @@ pub const AccountsDB = struct {
28732874
return slotSatisfiesMax(slot, max_slot) and slotSatisfiesMin(slot, min_slot);
28742875
}
28752876

2876-
pub fn registerRPCHooks(self: *AccountsDB, rpc_hooks: *sig.rpc.Hooks) !void {
2877+
pub fn registerRPCHooks(
2878+
self: *AccountsDB,
2879+
rpc_hooks: *sig.rpc.Hooks,
2880+
slot_tracker: ?*const sig.replay.trackers.SlotTracker,
2881+
snapshot_slot: Slot,
2882+
) !void {
28772883
try rpc_hooks.set(self.allocator, struct {
28782884
accountsdb: *AccountsDB,
2885+
slot_tracker: ?*const sig.replay.trackers.SlotTracker,
2886+
snapshot_slot: Slot,
2887+
2888+
pub fn getBalance(
2889+
this: @This(),
2890+
_: std.mem.Allocator,
2891+
params: methods.GetBalance,
2892+
) !methods.GetBalance.Response {
2893+
const config = params.config orelse methods.common.CommitmentSlotConfig{};
2894+
2895+
const commitment = config.commitment orelse .finalized;
2896+
const commitment_slot: Slot = if (this.slot_tracker) |tracker|
2897+
tracker.getSlotForCommitment(commitment)
2898+
else
2899+
this.snapshot_slot;
2900+
const max_slot: ?Slot = if (commitment == .processed) null else commitment_slot;
2901+
2902+
// Check minContextSlot constraint.
2903+
if (config.minContextSlot) |min_slot| {
2904+
if (commitment_slot < min_slot) {
2905+
return error.RpcMinContextSlotNotMet;
2906+
}
2907+
}
2908+
2909+
// Look up account
2910+
const result = this.accountsdb.getSlotAndAccountInSlotRangeWithReadLock(
2911+
&params.pubkey,
2912+
config.minContextSlot,
2913+
max_slot,
2914+
) catch return error.AccountsDbError;
2915+
2916+
const lamports: u64 = if (result) |r| blk: {
2917+
const account, _, var account_lock = r;
2918+
defer account_lock.unlock();
2919+
2920+
break :blk switch (account) {
2921+
.file => |aif| aif.account_info.lamports,
2922+
.unrooted_map => |um| um.lamports,
2923+
};
2924+
// TODO: is defaulting to 0 the correct behavior here?
2925+
} else 0;
2926+
2927+
return .{
2928+
.context = .{
2929+
.slot = commitment_slot,
2930+
.apiVersion = "2.0.15",
2931+
},
2932+
.value = lamports,
2933+
};
2934+
}
28792935

28802936
pub fn getAccountInfo(
28812937
this: @This(),
@@ -3098,7 +3154,7 @@ pub const AccountsDB = struct {
30983154

30993155
return null;
31003156
}
3101-
}{ .accountsdb = self });
3157+
}{ .accountsdb = self, .slot_tracker = slot_tracker, .snapshot_slot = snapshot_slot });
31023158
}
31033159
};
31043160

src/cmd.zig

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2064,9 +2064,10 @@ fn mockRpcServer(allocator: std.mem.Allocator, cfg: config.Cmd) !void {
20642064
defer manifest.deinit(allocator);
20652065
}
20662066

2067+
const snapshot_slot = if (snap_files.incremental_info) |inc| inc.slot else snap_files.full.slot;
20672068
var rpc_hooks = sig.rpc.Hooks{};
20682069
defer rpc_hooks.deinit(allocator);
2069-
try accountsdb.registerRPCHooks(&rpc_hooks);
2070+
try accountsdb.registerRPCHooks(&rpc_hooks, null, snapshot_slot);
20702071

20712072
var server_ctx = try sig.rpc.server.Context.init(.{
20722073
.allocator = allocator,

src/rpc/server/server.zig

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ test "serveSpawn getSnapshot" {
285285
var rpc_hooks = sig.rpc.Hooks{};
286286
defer rpc_hooks.deinit(accountsdb.allocator);
287287

288-
try accountsdb.registerRPCHooks(&rpc_hooks);
288+
try accountsdb.registerRPCHooks(&rpc_hooks, null, 0);
289289

290290
const sock_addr = std.net.Address.initIp4(.{ 0, 0, 0, 0 }, 0);
291291
var server_ctx = try Context.init(.{
@@ -440,7 +440,7 @@ test "serveSpawn getAccountInfo" {
440440

441441
var rpc_hooks = sig.rpc.Hooks{};
442442
defer rpc_hooks.deinit(allocator);
443-
try accountsdb.registerRPCHooks(&rpc_hooks);
443+
try accountsdb.registerRPCHooks(&rpc_hooks, null, expected_slot);
444444

445445
const test_sock_addr = std.net.Address.initIp4(.{ 0, 0, 0, 0 }, 0);
446446
var server_ctx = try Context.init(.{

0 commit comments

Comments
 (0)