diff --git a/lightning/src/offers/alloc.rs b/lightning/src/offers/alloc.rs new file mode 100644 index 00000000000..bc968cbb4ba --- /dev/null +++ b/lightning/src/offers/alloc.rs @@ -0,0 +1,49 @@ +// This file is Copyright its original authors, visible in version control +// history. +// +// This file is licensed under the Apache License, Version 2.0 or the MIT license +// , at your option. +// You may not use this file except in accordance with one or both of these +// licenses. + +/// Trait for rounding the size of memory allocations in order to reduce heap fragmentation. +pub(super) trait WithRoundedCapacity { + fn with_rounded_capacity(capacity: usize) -> Self; +} + +const MIN_ALLOCATION: usize = 512; +const MAX_ALLOCATION: usize = 1 << 16; + +impl WithRoundedCapacity for Vec { + fn with_rounded_capacity(capacity: usize) -> Self { + let capacity = if capacity == 0 { + 0 + } else if capacity <= MIN_ALLOCATION { + MIN_ALLOCATION + } else if capacity >= MAX_ALLOCATION { + MAX_ALLOCATION + } else { + capacity.next_power_of_two() + }; + + Vec::with_capacity(capacity) + } +} + +#[cfg(test)] +mod tests { + use super::{WithRoundedCapacity, MAX_ALLOCATION, MIN_ALLOCATION}; + + #[test] + fn rounds_capacity_to_power_of_two() { + assert_eq!(Vec::with_rounded_capacity(0).capacity(), 0); + assert_eq!(Vec::with_rounded_capacity(1).capacity(), MIN_ALLOCATION); + assert_eq!(Vec::with_rounded_capacity(512).capacity(), MIN_ALLOCATION); + assert_eq!(Vec::with_rounded_capacity(4095).capacity(), 4096); + assert_eq!(Vec::with_rounded_capacity(4096).capacity(), 4096); + assert_eq!(Vec::with_rounded_capacity(4097).capacity(), 8192); + assert_eq!(Vec::with_rounded_capacity(65537).capacity(), MAX_ALLOCATION); + assert_eq!(Vec::with_rounded_capacity(usize::MAX).capacity(), MAX_ALLOCATION); + } +} diff --git a/lightning/src/offers/invoice.rs b/lightning/src/offers/invoice.rs index 136ea2625de..aff69ff8b56 100644 --- a/lightning/src/offers/invoice.rs +++ b/lightning/src/offers/invoice.rs @@ -118,6 +118,7 @@ use crate::ln::channelmanager::PaymentId; use crate::types::features::{Bolt12InvoiceFeatures, InvoiceRequestFeatures, OfferFeatures}; use crate::ln::inbound_payment::{ExpandedKey, IV_LEN}; use crate::ln::msgs::DecodeError; +use crate::offers::alloc::WithRoundedCapacity; use crate::offers::invoice_macros::{invoice_accessors_common, invoice_builder_methods_common}; #[cfg(test)] use crate::offers::invoice_macros::invoice_builder_methods_test; @@ -507,9 +508,9 @@ impl UnsignedBolt12Invoice { // - 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( + // This assumes the invoice will only have one signature using the same number of bytes as + // the invoice request's signature. + let mut bytes = Vec::with_rounded_capacity( invreq_bytes.len() + invoice_tlv_stream.serialized_length() + if contents.is_for_offer() { 0 } else { SIGNATURE_TLV_RECORD_SIZE } @@ -530,7 +531,7 @@ impl UnsignedBolt12Invoice { let mut experimental_tlv_stream = TlvStream::new(remaining_bytes) .range(EXPERIMENTAL_TYPES) .peekable(); - let mut experimental_bytes = Vec::with_capacity( + let mut experimental_bytes = Vec::with_rounded_capacity( remaining_bytes.len() - experimental_tlv_stream .peek() @@ -543,7 +544,6 @@ impl UnsignedBolt12Invoice { } 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); @@ -574,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 { diff --git a/lightning/src/offers/invoice_request.rs b/lightning/src/offers/invoice_request.rs index 74bbdb8a0bf..20bcaf30a79 100644 --- a/lightning/src/offers/invoice_request.rs +++ b/lightning/src/offers/invoice_request.rs @@ -77,6 +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::alloc::WithRoundedCapacity; use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE}; use crate::offers::nonce::Nonce; use crate::offers::offer::{EXPERIMENTAL_OFFER_TYPES, ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OFFER_TYPES, Offer, OfferContents, OfferId, OfferTlvStream, OfferTlvStreamRef}; @@ -477,7 +478,7 @@ impl UnsignedInvoiceRequest { // - 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( + let mut bytes = Vec::with_rounded_capacity( offer.bytes.len() + payer_tlv_stream.serialized_length() + invoice_request_tlv_stream.serialized_length() @@ -498,7 +499,7 @@ impl UnsignedInvoiceRequest { let mut experimental_tlv_stream = TlvStream::new(remaining_bytes) .range(EXPERIMENTAL_OFFER_TYPES) .peekable(); - let mut experimental_bytes = Vec::with_capacity( + let mut experimental_bytes = Vec::with_rounded_capacity( remaining_bytes.len() - experimental_tlv_stream .peek() @@ -511,7 +512,6 @@ impl UnsignedInvoiceRequest { } 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); @@ -544,12 +544,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 { diff --git a/lightning/src/offers/mod.rs b/lightning/src/offers/mod.rs index e4fe7d789db..a1cd9e171b5 100644 --- a/lightning/src/offers/mod.rs +++ b/lightning/src/offers/mod.rs @@ -15,6 +15,7 @@ #[macro_use] pub mod offer; +pub(super) mod alloc; pub mod invoice; pub mod invoice_error; mod invoice_macros; diff --git a/lightning/src/offers/offer.rs b/lightning/src/offers/offer.rs index 613f9accd47..0ebe939244c 100644 --- a/lightning/src/offers/offer.rs +++ b/lightning/src/offers/offer.rs @@ -91,6 +91,7 @@ use crate::ln::channelmanager::PaymentId; use crate::types::features::OfferFeatures; use crate::ln::inbound_payment::{ExpandedKey, IV_LEN}; use crate::ln::msgs::{DecodeError, MAX_VALUE_MSAT}; +use crate::offers::alloc::WithRoundedCapacity; use crate::offers::merkle::{TaggedHash, TlvRecord, TlvStream}; use crate::offers::nonce::Nonce; use crate::offers::parse::{Bech32Encode, Bolt12ParseError, Bolt12SemanticError, ParsedMessage}; @@ -438,7 +439,7 @@ macro_rules! offer_builder_methods { ( } } - let mut bytes = Vec::new(); + let mut bytes = Vec::with_rounded_capacity($self.offer.serialized_length()); $self.offer.write(&mut bytes).unwrap(); let id = OfferId::from_valid_offer_tlv_stream(&bytes); diff --git a/lightning/src/offers/refund.rs b/lightning/src/offers/refund.rs index a68d0eb658e..a8a5941fde5 100644 --- a/lightning/src/offers/refund.rs +++ b/lightning/src/offers/refund.rs @@ -98,6 +98,7 @@ use crate::ln::channelmanager::PaymentId; use crate::types::features::InvoiceRequestFeatures; use crate::ln::inbound_payment::{ExpandedKey, IV_LEN}; use crate::ln::msgs::{DecodeError, MAX_VALUE_MSAT}; +use crate::offers::alloc::WithRoundedCapacity; use crate::offers::invoice_request::{ExperimentalInvoiceRequestTlvStream, ExperimentalInvoiceRequestTlvStreamRef, InvoiceRequestTlvStream, InvoiceRequestTlvStreamRef}; use crate::offers::nonce::Nonce; use crate::offers::offer::{ExperimentalOfferTlvStream, ExperimentalOfferTlvStreamRef, OfferTlvStream, OfferTlvStreamRef}; @@ -338,7 +339,7 @@ macro_rules! refund_builder_methods { ( $self.refund.payer.0 = metadata; } - let mut bytes = Vec::new(); + let mut bytes = Vec::with_rounded_capacity($self.refund.serialized_length()); $self.refund.write(&mut bytes).unwrap(); Ok(Refund { diff --git a/lightning/src/offers/static_invoice.rs b/lightning/src/offers/static_invoice.rs index 39c17eb3bcc..9a7b924f2d0 100644 --- a/lightning/src/offers/static_invoice.rs +++ b/lightning/src/offers/static_invoice.rs @@ -14,6 +14,7 @@ use crate::blinded_path::payment::BlindedPaymentPath; use crate::io; use crate::ln::inbound_payment::ExpandedKey; use crate::ln::msgs::DecodeError; +use crate::offers::alloc::WithRoundedCapacity; use crate::offers::invoice::{ check_invoice_signing_pubkey, construct_payment_paths, filter_fallbacks, ExperimentalInvoiceTlvStream, ExperimentalInvoiceTlvStreamRef, FallbackAddress, @@ -292,7 +293,7 @@ impl UnsignedStaticInvoice { // - 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( + let mut bytes = Vec::with_rounded_capacity( offer_bytes.len() + invoice_tlv_stream.serialized_length() + SIGNATURE_TLV_RECORD_SIZE @@ -311,7 +312,7 @@ impl UnsignedStaticInvoice { let mut experimental_tlv_stream = TlvStream::new(remaining_bytes).range(EXPERIMENTAL_OFFER_TYPES).peekable(); - let mut experimental_bytes = Vec::with_capacity( + let mut experimental_bytes = Vec::with_rounded_capacity( remaining_bytes.len() - experimental_tlv_stream .peek() @@ -324,7 +325,6 @@ impl UnsignedStaticInvoice { } 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); @@ -344,12 +344,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 })