Skip to content

Commit

Permalink
Use fixed-size Vec allocations for BOLT12 messages
Browse files Browse the repository at this point in the history
Instead of using elaborate calculations to determine the exact amount of
bytes need for a BOLT12 message are allocated, use a fixed size amount.
This reduces the code complexity and potentially reduces heap
fragmentation in the normal case.
  • Loading branch information
jkczyz committed Jan 31, 2025
1 parent 569f906 commit ffaccc0
Show file tree
Hide file tree
Showing 6 changed files with 34 additions and 107 deletions.
53 changes: 12 additions & 41 deletions lightning/src/offers/invoice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ use crate::offers::invoice_macros::{invoice_accessors_common, invoice_builder_me
#[cfg(test)]
use crate::offers::invoice_macros::invoice_builder_methods_test_common;
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};
use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE};
use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self};
use crate::offers::nonce::Nonce;
use crate::offers::offer::{Amount, EXPERIMENTAL_OFFER_TYPES, ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OFFER_TYPES, OfferTlvStream, OfferTlvStreamRef, Quantity};
use crate::offers::parse::{Bolt12ParseError, Bolt12SemanticError, ParsedMessage};
Expand Down Expand Up @@ -520,19 +520,8 @@ impl UnsignedBolt12Invoice {
let (_, _, _, invoice_tlv_stream, _, _, experimental_invoice_tlv_stream) =
contents.as_tlv_stream();

// Allocate enough space for the invoice, which will include:
// - all TLV records from `invreq_bytes` except signatures,
// - all invoice-specific TLV records, and
// - a signature TLV record once the invoice is signed.
//
// This assumes both the invoice request and the invoice will each only have one signature
// using SIGNATURE_TYPES.start as the TLV record. Thus, it is accounted for by invreq_bytes.
let mut bytes = Vec::with_capacity(
invreq_bytes.len()
+ invoice_tlv_stream.serialized_length()
+ if contents.is_for_offer() { 0 } else { SIGNATURE_TLV_RECORD_SIZE }
+ experimental_invoice_tlv_stream.serialized_length(),
);
const INVOICE_ALLOCATION_SIZE: usize = 1024;
let mut bytes = Vec::with_capacity(INVOICE_ALLOCATION_SIZE);

// Use the invoice_request bytes instead of the invoice_request TLV stream as the latter may
// have contained unknown TLV records, which are not stored in `InvoiceRequestContents` or
Expand All @@ -545,23 +534,16 @@ impl UnsignedBolt12Invoice {

invoice_tlv_stream.write(&mut bytes).unwrap();

let mut experimental_tlv_stream = TlvStream::new(remaining_bytes)
.range(EXPERIMENTAL_TYPES)
.peekable();
let mut experimental_bytes = Vec::with_capacity(
remaining_bytes.len()
- experimental_tlv_stream
.peek()
.map_or(remaining_bytes.len(), |first_record| first_record.start)
+ experimental_invoice_tlv_stream.serialized_length(),
);
const EXPERIMENTAL_TLV_ALLOCATION_SIZE: usize = 0;
let mut experimental_bytes = Vec::with_capacity(EXPERIMENTAL_TLV_ALLOCATION_SIZE);

let experimental_tlv_stream = TlvStream::new(remaining_bytes)
.range(EXPERIMENTAL_TYPES);
for record in experimental_tlv_stream {
record.write(&mut experimental_bytes).unwrap();
}

experimental_invoice_tlv_stream.write(&mut experimental_bytes).unwrap();
debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());

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

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

Ok(Bolt12Invoice {
Expand Down Expand Up @@ -965,13 +939,6 @@ impl Hash for Bolt12Invoice {
}

impl InvoiceContents {
fn is_for_offer(&self) -> bool {
match self {
InvoiceContents::ForOffer { .. } => true,
InvoiceContents::ForRefund { .. } => false,
}
}

/// Whether the original offer or refund has expired.
#[cfg(feature = "std")]
fn is_offer_or_refund_expired(&self) -> bool {
Expand Down Expand Up @@ -1362,7 +1329,11 @@ pub(super) const EXPERIMENTAL_INVOICE_TYPES: core::ops::RangeFrom<u64> = 3_000_0

#[cfg(not(test))]
tlv_stream!(
ExperimentalInvoiceTlvStream, ExperimentalInvoiceTlvStreamRef, EXPERIMENTAL_INVOICE_TYPES, {}
ExperimentalInvoiceTlvStream, ExperimentalInvoiceTlvStreamRef, EXPERIMENTAL_INVOICE_TYPES, {
// When adding experimental TLVs, update EXPERIMENTAL_TLV_ALLOCATION_SIZE accordingly in
// both UnsignedBolt12Invoice:new and UnsignedStaticInvoice::new to avoid unnecessary
// allocations.
}
);

#[cfg(test)]
Expand Down
41 changes: 11 additions & 30 deletions lightning/src/offers/invoice_request.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ use crate::ln::channelmanager::PaymentId;
use crate::types::features::InvoiceRequestFeatures;
use crate::ln::inbound_payment::{ExpandedKey, IV_LEN};
use crate::ln::msgs::DecodeError;
use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE};
use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self};
use crate::offers::nonce::Nonce;
use crate::offers::offer::{Amount, EXPERIMENTAL_OFFER_TYPES, ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OFFER_TYPES, Offer, OfferContents, OfferId, OfferTlvStream, OfferTlvStreamRef};
use crate::offers::parse::{Bolt12ParseError, ParsedMessage, Bolt12SemanticError};
Expand Down Expand Up @@ -473,17 +473,8 @@ impl UnsignedInvoiceRequest {
_experimental_offer_tlv_stream, experimental_invoice_request_tlv_stream,
) = contents.as_tlv_stream();

// Allocate enough space for the invoice_request, which will include:
// - all TLV records from `offer.bytes`,
// - all invoice_request-specific TLV records, and
// - a signature TLV record once the invoice_request is signed.
let mut bytes = Vec::with_capacity(
offer.bytes.len()
+ payer_tlv_stream.serialized_length()
+ invoice_request_tlv_stream.serialized_length()
+ SIGNATURE_TLV_RECORD_SIZE
+ experimental_invoice_request_tlv_stream.serialized_length(),
);
const INVOICE_REQUEST_ALLOCATION_SIZE: usize = 512;
let mut bytes = Vec::with_capacity(INVOICE_REQUEST_ALLOCATION_SIZE);

payer_tlv_stream.write(&mut bytes).unwrap();

Expand All @@ -495,23 +486,16 @@ impl UnsignedInvoiceRequest {

invoice_request_tlv_stream.write(&mut bytes).unwrap();

let mut experimental_tlv_stream = TlvStream::new(remaining_bytes)
.range(EXPERIMENTAL_OFFER_TYPES)
.peekable();
let mut experimental_bytes = Vec::with_capacity(
remaining_bytes.len()
- experimental_tlv_stream
.peek()
.map_or(remaining_bytes.len(), |first_record| first_record.start)
+ experimental_invoice_request_tlv_stream.serialized_length(),
);
const EXPERIMENTAL_TLV_ALLOCATION_SIZE: usize = 0;
let mut experimental_bytes = Vec::with_capacity(EXPERIMENTAL_TLV_ALLOCATION_SIZE);

let experimental_tlv_stream = TlvStream::new(remaining_bytes)
.range(EXPERIMENTAL_OFFER_TYPES);
for record in experimental_tlv_stream {
record.write(&mut experimental_bytes).unwrap();
}

experimental_invoice_request_tlv_stream.write(&mut experimental_bytes).unwrap();
debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());

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

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

Ok(InvoiceRequest {
Expand Down Expand Up @@ -1127,7 +1105,10 @@ pub(super) const EXPERIMENTAL_INVOICE_REQUEST_TYPES: core::ops::Range<u64> =
#[cfg(not(test))]
tlv_stream!(
ExperimentalInvoiceRequestTlvStream, ExperimentalInvoiceRequestTlvStreamRef,
EXPERIMENTAL_INVOICE_REQUEST_TYPES, {}
EXPERIMENTAL_INVOICE_REQUEST_TYPES, {
// When adding experimental TLVs, update EXPERIMENTAL_TLV_ALLOCATION_SIZE accordingly in
// UnsignedInvoiceRequest::new to avoid unnecessary allocations.
}
);

#[cfg(test)]
Expand Down
8 changes: 1 addition & 7 deletions lightning/src/offers/merkle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
use bitcoin::hashes::{Hash, HashEngine, sha256};
use bitcoin::secp256k1::{Message, PublicKey, Secp256k1, self};
use bitcoin::secp256k1::constants::SCHNORR_SIGNATURE_SIZE;
use bitcoin::secp256k1::schnorr::Signature;
use crate::io;
use crate::util::ser::{BigSize, Readable, Writeable, Writer};
Expand All @@ -26,10 +25,6 @@ tlv_stream!(SignatureTlvStream, SignatureTlvStreamRef<'a>, SIGNATURE_TYPES, {
(240, signature: Signature),
});

/// Size of a TLV record in `SIGNATURE_TYPES` when the type is 1000. TLV types are encoded using
/// BigSize, so a TLV record with type 240 will use two less bytes.
pub(super) const SIGNATURE_TLV_RECORD_SIZE: usize = 3 + 1 + SCHNORR_SIGNATURE_SIZE;

/// A hash for use in a specific context by tweaking with a context-dependent tag as per [BIP 340]
/// and computed over the merkle root of a TLV stream to sign as defined in [BOLT 12].
///
Expand Down Expand Up @@ -253,7 +248,6 @@ pub(super) struct TlvRecord<'a> {
type_bytes: &'a [u8],
// The entire TLV record.
pub(super) record_bytes: &'a [u8],
pub(super) start: usize,
pub(super) end: usize,
}

Expand All @@ -278,7 +272,7 @@ impl<'a> Iterator for TlvStream<'a> {
self.data.set_position(end);

Some(TlvRecord {
r#type, type_bytes, record_bytes, start: start as usize, end: end as usize,
r#type, type_bytes, record_bytes, end: end as usize,
})
} else {
None
Expand Down
3 changes: 2 additions & 1 deletion lightning/src/offers/offer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -438,7 +438,8 @@ macro_rules! offer_builder_methods { (
}
}

let mut bytes = Vec::new();
const OFFER_ALLOCATION_SIZE: usize = 512;
let mut bytes = Vec::with_capacity(OFFER_ALLOCATION_SIZE);
$self.offer.write(&mut bytes).unwrap();

let id = OfferId::from_valid_offer_tlv_stream(&bytes);
Expand Down
3 changes: 2 additions & 1 deletion lightning/src/offers/refund.rs
Original file line number Diff line number Diff line change
Expand Up @@ -338,7 +338,8 @@ macro_rules! refund_builder_methods { (
$self.refund.payer.0 = metadata;
}

let mut bytes = Vec::new();
const REFUND_ALLOCATION_SIZE: usize = 512;
let mut bytes = Vec::with_capacity(REFUND_ALLOCATION_SIZE);
$self.refund.write(&mut bytes).unwrap();

Ok(Refund {
Expand Down
33 changes: 6 additions & 27 deletions lightning/src/offers/static_invoice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@ use crate::offers::invoice_macros::{invoice_accessors_common, invoice_builder_me
use crate::offers::invoice_request::InvoiceRequest;
use crate::offers::merkle::{
self, SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream,
SIGNATURE_TLV_RECORD_SIZE,
};
use crate::offers::nonce::Nonce;
use crate::offers::offer::{
Expand Down Expand Up @@ -288,16 +287,8 @@ impl UnsignedStaticInvoice {
fn new(offer_bytes: &Vec<u8>, contents: InvoiceContents) -> Self {
let (_, invoice_tlv_stream, _, experimental_invoice_tlv_stream) = contents.as_tlv_stream();

// Allocate enough space for the invoice, which will include:
// - all TLV records from `offer_bytes`,
// - all invoice-specific TLV records, and
// - a signature TLV record once the invoice is signed.
let mut bytes = Vec::with_capacity(
offer_bytes.len()
+ invoice_tlv_stream.serialized_length()
+ SIGNATURE_TLV_RECORD_SIZE
+ experimental_invoice_tlv_stream.serialized_length(),
);
const INVOICE_ALLOCATION_SIZE: usize = 1024;
let mut bytes = Vec::with_capacity(INVOICE_ALLOCATION_SIZE);

// Use the offer bytes instead of the offer TLV stream as the latter may have contained
// unknown TLV records, which are not stored in `InvoiceContents`.
Expand All @@ -309,22 +300,16 @@ impl UnsignedStaticInvoice {

invoice_tlv_stream.write(&mut bytes).unwrap();

let mut experimental_tlv_stream =
TlvStream::new(remaining_bytes).range(EXPERIMENTAL_OFFER_TYPES).peekable();
let mut experimental_bytes = Vec::with_capacity(
remaining_bytes.len()
- experimental_tlv_stream
.peek()
.map_or(remaining_bytes.len(), |first_record| first_record.start)
+ experimental_invoice_tlv_stream.serialized_length(),
);
const EXPERIMENTAL_TLV_ALLOCATION_SIZE: usize = 0;
let mut experimental_bytes = Vec::with_capacity(EXPERIMENTAL_TLV_ALLOCATION_SIZE);

let experimental_tlv_stream =
TlvStream::new(remaining_bytes).range(EXPERIMENTAL_OFFER_TYPES);
for record in experimental_tlv_stream {
record.write(&mut experimental_bytes).unwrap();
}

experimental_invoice_tlv_stream.write(&mut experimental_bytes).unwrap();
debug_assert_eq!(experimental_bytes.len(), experimental_bytes.capacity());

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

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

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

0 comments on commit ffaccc0

Please sign in to comment.