Don't over-allocate invoice bytes

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 <LICENSE-APACHE
+// or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, 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<u8> {
+	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);
+	}
+}

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 {

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 {

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;

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);

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 {

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 })