Reduce heap fragmentation when allocating bytes

jkczyz · jkczyz · commit 6d4b4b2ba1eb · 2024-12-18T11:37:54.000-06:00
When allocating Vec&lt;u8&gt; 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.
diff --git a/lightning/src/offers/alloc.rs 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 <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);
+	}
+}
diff --git a/lightning/src/offers/invoice.rs b/lightning/src/offers/invoice.rs
@@ -118,11 +118,12 @@ 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;
 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, SIGNATURE_TYPES};
+use crate::offers::merkle::{SignError, SignFn, SignatureTlvStream, SignatureTlvStreamRef, TaggedHash, TlvStream, self, SIGNATURE_TLV_RECORD_SIZE};
 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};
@@ -502,32 +503,17 @@ impl UnsignedBolt12Invoice {
 		let (_, _, _, invoice_tlv_stream, _, _, experimental_invoice_tlv_stream) =
 			contents.as_tlv_stream();
 
-		let mut signature_tlv_stream = TlvStream::new(invreq_bytes)
-			.range(SIGNATURE_TYPES)
-			.peekable();
-		let signature_len = signature_tlv_stream
-			.peek()
-			.map_or(SIGNATURE_TLV_RECORD_SIZE, |record| record.end - record.start);
-		let signature_tlv_stream_start = signature_tlv_stream
-			.peek()
-			.map_or(0, |first_record| first_record.start);
-		let signature_tlv_stream_end = signature_tlv_stream
-			.last()
-			.map_or(0, |last_record| last_record.end);
-		let signature_tlv_stream_len = signature_tlv_stream_end - signature_tlv_stream_start;
-
 		// 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 the invoice will only have one signature using the same number of bytes as
-		// the first (and probably only) signature from the invoice request.
-		let mut bytes = Vec::with_capacity(
+		// the invoice request's signature.
+		let mut bytes = Vec::with_rounded_capacity(
 			invreq_bytes.len()
-				- signature_tlv_stream_len
 				+ invoice_tlv_stream.serialized_length()
-				+ signature_len
+				+ if contents.is_for_offer() { 0 } else { SIGNATURE_TLV_RECORD_SIZE }
 				+ experimental_invoice_tlv_stream.serialized_length(),
 		);
 
@@ -545,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()
@@ -558,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);
@@ -589,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
@@ -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
@@ -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/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 })
