Skip to content

Commit 6d4b4b2

Browse files
jkczyzjkczyz
committed
Reduce heap fragmentation when allocating bytes
When allocating Vec<u8> for BOLT12 messages, use allocations that are powers of two starting at 512 instead of using the exact number of bytes provided. This helps reduce the amount of heap fragmentation.
1 parent 49bccc2 commit 6d4b4b2

File tree

5 files changed

+62
-47
lines changed

5 files changed

+62
-47
lines changed

‎lightning/src/offers/alloc.rs

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
// This file is Copyright its original authors, visible in version control
2+
// history.
3+
//
4+
// This file is licensed under the Apache License, Version 2.0 <LICENSE-APACHE
5+
// or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
6+
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
7+
// You may not use this file except in accordance with one or both of these
8+
// licenses.
9+
10+
/// Trait for rounding the size of memory allocations in order to reduce heap fragmentation.
11+
pub(super) trait WithRoundedCapacity {
12+
fn with_rounded_capacity(capacity: usize) -> Self;
13+
}
14+
15+
const MIN_ALLOCATION: usize = 512;
16+
const MAX_ALLOCATION: usize = 1 << 16;
17+
18+
impl WithRoundedCapacity for Vec<u8> {
19+
fn with_rounded_capacity(capacity: usize) -> Self {
20+
let capacity = if capacity == 0 {
21+
0
22+
} else if capacity <= MIN_ALLOCATION {
23+
MIN_ALLOCATION
24+
} else if capacity >= MAX_ALLOCATION {
25+
MAX_ALLOCATION
26+
} else {
27+
capacity.next_power_of_two()
28+
};
29+
30+
Vec::with_capacity(capacity)
31+
}
32+
}
33+
34+
#[cfg(test)]
35+
mod tests {
36+
use super::{WithRoundedCapacity, MAX_ALLOCATION, MIN_ALLOCATION};
37+
38+
#[test]
39+
fn rounds_capacity_to_power_of_two() {
40+
assert_eq!(Vec::with_rounded_capacity(0).capacity(), 0);
41+
assert_eq!(Vec::with_rounded_capacity(1).capacity(), MIN_ALLOCATION);
42+
assert_eq!(Vec::with_rounded_capacity(512).capacity(), MIN_ALLOCATION);
43+
assert_eq!(Vec::with_rounded_capacity(4095).capacity(), 4096);
44+
assert_eq!(Vec::with_rounded_capacity(4096).capacity(), 4096);
45+
assert_eq!(Vec::with_rounded_capacity(4097).capacity(), 8192);
46+
assert_eq!(Vec::with_rounded_capacity(65537).capacity(), MAX_ALLOCATION);
47+
assert_eq!(Vec::with_rounded_capacity(usize::MAX).capacity(), MAX_ALLOCATION);
48+
}
49+
}

‎lightning/src/offers/invoice.rs

Lines changed: 6 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -118,11 +118,12 @@ use crate::ln::channelmanager::PaymentId;
118118
use crate::types::features::{Bolt12InvoiceFeatures, InvoiceRequestFeatures, OfferFeatures};
119119
use crate::ln::inbound_payment::{ExpandedKey, IV_LEN};
120120
use crate::ln::msgs::DecodeError;
121+
use crate::offers::alloc::WithRoundedCapacity;
121122
use crate::offers::invoice_macros::{invoice_accessors_common, invoice_builder_methods_common};
122123
#[cfg(test)]
123124
use crate::offers::invoice_macros::invoice_builder_methods_test;
124125
use crate::offers::invoice_request::{EXPERIMENTAL_INVOICE_REQUEST_TYPES, ExperimentalInvoiceRequestTlvStream, ExperimentalInvoiceRequestTlvStreamRef, INVOICE_REQUEST_PAYER_ID_TYPE, INVOICE_REQUEST_TYPES, IV_BYTES as INVOICE_REQUEST_IV_BYTES, InvoiceRequest, InvoiceRequestContents, InvoiceRequestTlvStream, InvoiceRequestTlvStreamRef};
125-
use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE, SIGNATURE_TYPES};
126+
use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE};
126127
use crate::offers::nonce::Nonce;
127128
use crate::offers::offer::{Amount, EXPERIMENTAL_OFFER_TYPES, ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OFFER_TYPES, OfferTlvStream, OfferTlvStreamRef, Quantity};
128129
use crate::offers::parse::{Bolt12ParseError, Bolt12SemanticError, ParsedMessage};
@@ -502,32 +503,17 @@ impl UnsignedBolt12Invoice {
502503
let (_, _, _, invoice_tlv_stream, _, _, experimental_invoice_tlv_stream) =
503504
contents.as_tlv_stream();
504505

505-
let mut signature_tlv_stream = TlvStream::new(invreq_bytes)
506-
.range(SIGNATURE_TYPES)
507-
.peekable();
508-
let signature_len = signature_tlv_stream
509-
.peek()
510-
.map_or(SIGNATURE_TLV_RECORD_SIZE, |record| record.end - record.start);
511-
let signature_tlv_stream_start = signature_tlv_stream
512-
.peek()
513-
.map_or(0, |first_record| first_record.start);
514-
let signature_tlv_stream_end = signature_tlv_stream
515-
.last()
516-
.map_or(0, |last_record| last_record.end);
517-
let signature_tlv_stream_len = signature_tlv_stream_end - signature_tlv_stream_start;
518-
519506
// Allocate enough space for the invoice, which will include:
520507
// - all TLV records from `invreq_bytes` except signatures,
521508
// - all invoice-specific TLV records, and
522509
// - a signature TLV record once the invoice is signed.
523510
//
524511
// This assumes the invoice will only have one signature using the same number of bytes as
525-
// the first (and probably only) signature from the invoice request.
526-
let mut bytes = Vec::with_capacity(
512+
// the invoice request's signature.
513+
let mut bytes = Vec::with_rounded_capacity(
527514
invreq_bytes.len()
528-
- signature_tlv_stream_len
529515
+ invoice_tlv_stream.serialized_length()
530-
+ signature_len
516+
+ if contents.is_for_offer() { 0 } else { SIGNATURE_TLV_RECORD_SIZE }
531517
+ experimental_invoice_tlv_stream.serialized_length(),
532518
);
533519

@@ -545,7 +531,7 @@ impl UnsignedBolt12Invoice {
545531
let mut experimental_tlv_stream = TlvStream::new(remaining_bytes)
546532
.range(EXPERIMENTAL_TYPES)
547533
.peekable();
548-
let mut experimental_bytes = Vec::with_capacity(
534+
let mut experimental_bytes = Vec::with_rounded_capacity(
549535
remaining_bytes.len()
550536
- experimental_tlv_stream
551537
.peek()
@@ -558,7 +544,6 @@ impl UnsignedBolt12Invoice {
558544
}
559545

560546
experimental_invoice_tlv_stream.write(&mut experimental_bytes).unwrap();
561-
debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());
562547

563548
let tlv_stream = TlvStream::new(&bytes).chain(TlvStream::new(&experimental_bytes));
564549
let tagged_hash = TaggedHash::from_tlv_stream(SIGNATURE_TAG, tlv_stream);
@@ -589,14 +574,6 @@ macro_rules! unsigned_invoice_sign_method { ($self: ident, $self_type: ty $(, $s
589574
signature_tlv_stream.write(&mut $self.bytes).unwrap();
590575

591576
// Append the experimental bytes after the signature.
592-
debug_assert_eq!(
593-
// The two-byte overallocation results from SIGNATURE_TLV_RECORD_SIZE accommodating TLV
594-
// records with types >= 253.
595-
$self.bytes.len()
596-
+ $self.experimental_bytes.len()
597-
+ if $self.contents.is_for_offer() { 0 } else { 2 },
598-
$self.bytes.capacity(),
599-
);
600577
$self.bytes.extend_from_slice(&$self.experimental_bytes);
601578

602579
Ok(Bolt12Invoice {

‎lightning/src/offers/invoice_request.rs

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ use crate::ln::channelmanager::PaymentId;
7777
use crate::types::features::InvoiceRequestFeatures;
7878
use crate::ln::inbound_payment::{ExpandedKey, IV_LEN};
7979
use crate::ln::msgs::DecodeError;
80+
use crate::offers::alloc::WithRoundedCapacity;
8081
use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE};
8182
use crate::offers::nonce::Nonce;
8283
use crate::offers::offer::{EXPERIMENTAL_OFFER_TYPES, ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OFFER_TYPES, Offer, OfferContents, OfferId, OfferTlvStream, OfferTlvStreamRef};
@@ -477,7 +478,7 @@ impl UnsignedInvoiceRequest {
477478
// - all TLV records from `offer.bytes`,
478479
// - all invoice_request-specific TLV records, and
479480
// - a signature TLV record once the invoice_request is signed.
480-
let mut bytes = Vec::with_capacity(
481+
let mut bytes = Vec::with_rounded_capacity(
481482
offer.bytes.len()
482483
+ payer_tlv_stream.serialized_length()
483484
+ invoice_request_tlv_stream.serialized_length()
@@ -498,7 +499,7 @@ impl UnsignedInvoiceRequest {
498499
let mut experimental_tlv_stream = TlvStream::new(remaining_bytes)
499500
.range(EXPERIMENTAL_OFFER_TYPES)
500501
.peekable();
501-
let mut experimental_bytes = Vec::with_capacity(
502+
let mut experimental_bytes = Vec::with_rounded_capacity(
502503
remaining_bytes.len()
503504
- experimental_tlv_stream
504505
.peek()
@@ -511,7 +512,6 @@ impl UnsignedInvoiceRequest {
511512
}
512513

513514
experimental_invoice_request_tlv_stream.write(&mut experimental_bytes).unwrap();
514-
debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());
515515

516516
let tlv_stream = TlvStream::new(&bytes).chain(TlvStream::new(&experimental_bytes));
517517
let tagged_hash = TaggedHash::from_tlv_stream(SIGNATURE_TAG, tlv_stream);
@@ -544,12 +544,6 @@ macro_rules! unsigned_invoice_request_sign_method { (
544544
signature_tlv_stream.write(&mut $self.bytes).unwrap();
545545

546546
// Append the experimental bytes after the signature.
547-
debug_assert_eq!(
548-
// The two-byte overallocation results from SIGNATURE_TLV_RECORD_SIZE accommodating TLV
549-
// records with types >= 253.
550-
$self.bytes.len() + $self.experimental_bytes.len() + 2,
551-
$self.bytes.capacity(),
552-
);
553547
$self.bytes.extend_from_slice(&$self.experimental_bytes);
554548

555549
Ok(InvoiceRequest {

‎lightning/src/offers/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#[macro_use]
1616
pub mod offer;
1717

18+
pub(super) mod alloc;
1819
pub mod invoice;
1920
pub mod invoice_error;
2021
mod invoice_macros;

‎lightning/src/offers/static_invoice.rs

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use crate::blinded_path::payment::BlindedPaymentPath;
1414
use crate::io;
1515
use crate::ln::inbound_payment::ExpandedKey;
1616
use crate::ln::msgs::DecodeError;
17+
use crate::offers::alloc::WithRoundedCapacity;
1718
use crate::offers::invoice::{
1819
check_invoice_signing_pubkey, construct_payment_paths, filter_fallbacks,
1920
ExperimentalInvoiceTlvStream, ExperimentalInvoiceTlvStreamRef, FallbackAddress,
@@ -292,7 +293,7 @@ impl UnsignedStaticInvoice {
292293
// - all TLV records from `offer_bytes`,
293294
// - all invoice-specific TLV records, and
294295
// - a signature TLV record once the invoice is signed.
295-
let mut bytes = Vec::with_capacity(
296+
let mut bytes = Vec::with_rounded_capacity(
296297
offer_bytes.len()
297298
+ invoice_tlv_stream.serialized_length()
298299
+ SIGNATURE_TLV_RECORD_SIZE
@@ -311,7 +312,7 @@ impl UnsignedStaticInvoice {
311312

312313
let mut experimental_tlv_stream =
313314
TlvStream::new(remaining_bytes).range(EXPERIMENTAL_OFFER_TYPES).peekable();
314-
let mut experimental_bytes = Vec::with_capacity(
315+
let mut experimental_bytes = Vec::with_rounded_capacity(
315316
remaining_bytes.len()
316317
- experimental_tlv_stream
317318
.peek()
@@ -324,7 +325,6 @@ impl UnsignedStaticInvoice {
324325
}
325326

326327
experimental_invoice_tlv_stream.write(&mut experimental_bytes).unwrap();
327-
debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());
328328

329329
let tlv_stream = TlvStream::new(&bytes).chain(TlvStream::new(&experimental_bytes));
330330
let tagged_hash = TaggedHash::from_tlv_stream(SIGNATURE_TAG, tlv_stream);
@@ -344,12 +344,6 @@ impl UnsignedStaticInvoice {
344344
signature_tlv_stream.write(&mut self.bytes).unwrap();
345345

346346
// Append the experimental bytes after the signature.
347-
debug_assert_eq!(
348-
// The two-byte overallocation results from SIGNATURE_TLV_RECORD_SIZE accommodating TLV
349-
// records with types >= 253.
350-
self.bytes.len() + self.experimental_bytes.len() + 2,
351-
self.bytes.capacity(),
352-
);
353347
self.bytes.extend_from_slice(&self.experimental_bytes);
354348

355349
Ok(StaticInvoice { bytes: self.bytes, contents: self.contents, signature })

0 commit comments

Comments
 (0)