refactor(bls): allocations around VMAS (ChainSafe#395)

rework allocations around `verifyMultipleAggregateSignatures`.

- we're batching on average about ~30 signature sets per batch on our
highest load fleet (according to metrics), so a cap of about 32 makes
sense for stack allocations. Anything beyond that, use heap allocations.
- avoid copies for `msgs` which was unnecessary

Author: spiral-ladder

bindings/napi/blst.zig

@@ -27,6 +27,14 @@ const ThreadPool = bls.ThreadPool;
 const DST = bls.DST;
 const MAX_AGGREGATE_PER_JOB = bls.MAX_AGGREGATE_PER_JOB;
 
+/// In upstream lodestar we split batchable sets into chunks of minimum size 16.
+/// Cost savings after ~16 are not significant.
+/// In metrics, we can observe that sas fleet receives on average ~30 signature sets,
+/// so a safe bound is about 32.
+///
+/// See: packages/beacon-node/src/chain/bls/multithread/worker.ts
+const BATCH_VERIFY_SIZE = 32;
+
 /// Cached thread pool reference for parallel verification.
 /// Initialized lazily on first use, torn down via `deinitThreadPool`.
 var thread_pool: ?*ThreadPool = null;
@@ -429,17 +437,40 @@ pub fn verifyMultipleAggregateSignatures(sets: js.Array, pks_validate: ?js.Boole
     const n_elems = try sets.length();
     if (n_elems == 0) return js.Boolean.from(false);
 
-    const msgs = try allocator.alloc([32]u8, n_elems);
-    defer allocator.free(msgs);
-
-    const pks = try allocator.alloc(*NativePublicKey, n_elems);
-    defer allocator.free(pks);
-
-    const sigs = try allocator.alloc(*NativeSignature, n_elems);
-    defer allocator.free(sigs);
-
-    const rands = try allocator.alloc([32]u8, n_elems);
-    defer allocator.free(rands);
+    var msgs_stack: [BATCH_VERIFY_SIZE][]const u8 = undefined;
+    var pks_stack: [BATCH_VERIFY_SIZE]*NativePublicKey = undefined;
+    var sigs_stack: [BATCH_VERIFY_SIZE]*NativeSignature = undefined;
+    var rands_stack: [BATCH_VERIFY_SIZE][32]u8 = undefined;
+
+    var msgs_heap: ?[][]const u8 = null;
+    defer if (msgs_heap) |buf| allocator.free(buf);
+    var pks_heap: ?[]*NativePublicKey = null;
+    defer if (pks_heap) |buf| allocator.free(buf);
+    var sigs_heap: ?[]*NativeSignature = null;
+    defer if (sigs_heap) |buf| allocator.free(buf);
+    var rands_heap: ?[][32]u8 = null;
+    defer if (rands_heap) |buf| allocator.free(buf);
+
+    const msgs = if (n_elems <= BATCH_VERIFY_SIZE) msgs_stack[0..n_elems] else blk: {
+        const buf = try allocator.alloc([]const u8, n_elems);
+        msgs_heap = buf;
+        break :blk buf;
+    };
+    const pks = if (n_elems <= BATCH_VERIFY_SIZE) pks_stack[0..n_elems] else blk: {
+        const buf = try allocator.alloc(*NativePublicKey, n_elems);
+        pks_heap = buf;
+        break :blk buf;
+    };
+    const sigs = if (n_elems <= BATCH_VERIFY_SIZE) sigs_stack[0..n_elems] else blk: {
+        const buf = try allocator.alloc(*NativeSignature, n_elems);
+        sigs_heap = buf;
+        break :blk buf;
+    };
+    const rands = if (n_elems <= BATCH_VERIFY_SIZE) rands_stack[0..n_elems] else blk: {
+        const buf = try allocator.alloc([32]u8, n_elems);
+        rands_heap = buf;
+        break :blk buf;
+    };
 
     var seed_bytes: [8]u8 = undefined;
     const io = napi_io.get();
@@ -454,7 +485,7 @@ pub fn verifyMultipleAggregateSignatures(sets: js.Array, pks_validate: ?js.Boole
         const msg_napi = try set.getNamedProperty("msg");
         const msg_bytes = try uint8SliceFromValue(.{ .val = msg_napi });
         if (msg_bytes.len != 32) return error.InvalidMessageLength;
-        @memcpy(&msgs[i], msg_bytes[0..32]);
+        msgs[i] = msg_bytes;
 
         const pk_napi = try set.getNamedProperty("pk");
         const wrapped_pk = try e.unwrap(PublicKey, pk_napi);

src/bls/ThreadPool.zig

@@ -196,7 +196,7 @@ pub fn submitAndWait(pool: *ThreadPool, io: std.Io, items: []*WorkItem) (PoolErr
 const VerifyMultiJob = struct {
     pks: []const *PublicKey,
     sigs: []const *Signature,
-    msgs: []const [32]u8,
+    msgs: []const []const u8,
     rands: []const [32]u8,
     dst: []const u8,
     pks_validate: bool,
@@ -231,7 +231,7 @@ const VerifyMultiWorkItem = struct {
                 job.sigs_groupcheck,
                 &job.rands[i],
                 RAND_BITS,
-                &job.msgs[i],
+                job.msgs[i],
             ) catch {
                 job.err_flag.store(true, .release);
                 break;
@@ -251,7 +251,7 @@ pub fn verifyMultipleAggregateSignatures(
     pool: *ThreadPool,
     io: std.Io,
     n_elems: usize,
-    msgs: []const [32]u8,
+    msgs: []const []const u8,
     dst: []const u8,
     pks: []const *PublicKey,
     pks_validate: bool,
@@ -492,6 +492,7 @@ test "verifyMultipleAggregateSignatures multi-threaded" {
     const num_sigs = 16;
 
     var msgs: [num_sigs][32]u8 = undefined;
+    var msg_refs: [num_sigs][]const u8 = undefined;
     var pks: [num_sigs]PublicKey = undefined;
     var sigs: [num_sigs]Signature = undefined;
     var pk_ptrs: [num_sigs]*PublicKey = undefined;
@@ -511,6 +512,7 @@ test "verifyMultipleAggregateSignatures multi-threaded" {
         const sk = try SecretKey.keyGen(&ikm_i, null);
         pks[i] = sk.toPublicKey();
         sigs[i] = sk.sign(&msgs[i], blst.DST, null);
+        msg_refs[i] = &msgs[i];
         pk_ptrs[i] = &pks[i];
         sig_ptrs[i] = &sigs[i];
     }
@@ -521,7 +523,7 @@ test "verifyMultipleAggregateSignatures multi-threaded" {
     const result = try pool.verifyMultipleAggregateSignatures(
         std.testing.io,
         num_sigs,
-        &msgs,
+        &msg_refs,
         blst.DST,
         &pk_ptrs,
         true,

src/bls/fast_verify.zig

@@ -12,7 +12,7 @@ const RAND_BITS = 8 * RAND_BYTES;
 pub fn verifyMultipleAggregateSignatures(
     pairing_buf: *align(Pairing.buf_align) [Pairing.sizeOf()]u8,
     n_elems: usize,
-    msgs: []const [32]u8,
+    msgs: []const []const u8,
     dst: []const u8,
     pks: []const *PublicKey,
     pks_validate: bool,
@@ -38,7 +38,7 @@ pub fn verifyMultipleAggregateSignatures(
             sigs_groupcheck,
             &rands[i],
             RAND_BITS,
-            &msgs[i],
+            msgs[i],
         );
     }