Skip to content

Commit a4546e2

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 ccfde8d commit a4546e2

File tree

4 files changed

+71
-5
lines changed

4 files changed

+71
-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
@@ -1997,9 +1997,10 @@ fn mockRpcServer(allocator: std.mem.Allocator, cfg: config.Cmd) !void {
19971997
defer manifest.deinit(allocator);
19981998
}
19991999

2000+
const snapshot_slot = if (snap_files.incremental_info) |inc| inc.slot else snap_files.full.slot;
20002001
var rpc_hooks = sig.rpc.Hooks{};
20012002
defer rpc_hooks.deinit(allocator);
2002-
try accountsdb.registerRPCHooks(&rpc_hooks);
2003+
try accountsdb.registerRPCHooks(&rpc_hooks, null, snapshot_slot);
20032004

20042005
var server_ctx = try sig.rpc.server.Context.init(.{
20052006
.allocator = allocator,

src/replay/trackers.zig

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ const EpochSchedule = sig.core.EpochSchedule;
1010
const Slot = sig.core.Slot;
1111
const SlotConstants = sig.core.SlotConstants;
1212
const SlotState = sig.core.SlotState;
13+
const Commitment = sig.rpc.methods.common.Commitment;
1314

1415
pub const ForkChoiceProcessedSlot = struct {
1516
slot: std.atomic.Value(Slot) = .init(0),
@@ -122,6 +123,14 @@ pub const SlotTracker = struct {
122123
return elem.toRef();
123124
}
124125

126+
pub fn getSlotForCommitment(self: *const SlotTracker, commitment: Commitment) Slot {
127+
return switch (commitment) {
128+
.processed => self.latest_processed_slot.get(),
129+
.confirmed => self.latest_confirmed_slot.get(),
130+
.finalized => self.root.load(.monotonic),
131+
};
132+
}
133+
125134
pub const GetOrPutResult = struct {
126135
found_existing: bool,
127136
reference: Reference,

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)