1use serde::Serialize;
2use std::collections::{BTreeMap, BTreeSet};
3
4use super::{invalid_resource, DynamicBackingPoolId, VNextError};
5
6#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize)]
7pub struct BackingChunkIdentity {
8 pool_id: DynamicBackingPoolId,
9 ordinal: u32,
10 generation: u64,
11}
12
13impl BackingChunkIdentity {
14 pub(super) fn from_parts(
15 pool_id: DynamicBackingPoolId,
16 ordinal: u32,
17 generation: u64,
18 ) -> Result<Self, VNextError> {
19 if ordinal == 0 || generation == 0 {
20 return Err(invalid_resource("backing chunk identity is invalid"));
21 }
22 Ok(Self {
23 pool_id,
24 ordinal,
25 generation,
26 })
27 }
28
29 pub fn pool_id(&self) -> &DynamicBackingPoolId {
30 &self.pool_id
31 }
32
33 pub const fn ordinal(&self) -> u32 {
34 self.ordinal
35 }
36
37 pub const fn generation(&self) -> u64 {
38 self.generation
39 }
40}
41
42#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
43pub struct BackingSegment {
44 chunk: BackingChunkIdentity,
45 offset_bytes: u64,
46 length_bytes: u64,
47}
48
49impl BackingSegment {
50 pub(crate) fn new(
51 chunk: BackingChunkIdentity,
52 offset_bytes: u64,
53 length_bytes: u64,
54 ) -> Result<Self, VNextError> {
55 Self::from_chunk(
56 &chunk.pool_id,
57 chunk.ordinal,
58 chunk.generation,
59 offset_bytes,
60 length_bytes,
61 )
62 }
63
64 pub(super) fn from_chunk(
65 pool_id: &DynamicBackingPoolId,
66 chunk_ordinal: u32,
67 chunk_generation: u64,
68 offset_bytes: u64,
69 length_bytes: u64,
70 ) -> Result<Self, VNextError> {
71 if chunk_ordinal == 0
72 || chunk_generation == 0
73 || length_bytes == 0
74 || offset_bytes.checked_add(length_bytes).is_none()
75 {
76 return Err(invalid_resource(
77 "backing segment has invalid chunk identity or physical range",
78 ));
79 }
80 Ok(Self {
81 chunk: BackingChunkIdentity::from_parts(
82 pool_id.clone(),
83 chunk_ordinal,
84 chunk_generation,
85 )?,
86 offset_bytes,
87 length_bytes,
88 })
89 }
90
91 pub fn chunk(&self) -> &BackingChunkIdentity {
92 &self.chunk
93 }
94
95 pub fn pool_id(&self) -> &DynamicBackingPoolId {
96 self.chunk.pool_id()
97 }
98
99 pub const fn chunk_ordinal(&self) -> u32 {
100 self.chunk.ordinal()
101 }
102
103 pub const fn chunk_generation(&self) -> u64 {
104 self.chunk.generation()
105 }
106
107 pub const fn offset_bytes(&self) -> u64 {
108 self.offset_bytes
109 }
110
111 pub const fn length_bytes(&self) -> u64 {
112 self.length_bytes
113 }
114}
115
116pub(super) fn backing_segment_range(
117 segments: &[BackingSegment],
118 physical_offset_bytes: u64,
119 size_bytes: u64,
120) -> Result<Vec<BackingSegment>, VNextError> {
121 let physical_end = physical_offset_bytes
122 .checked_add(size_bytes)
123 .ok_or_else(|| invalid_resource("logical backing projection range overflows u64"))?;
124 if size_bytes == 0 {
125 return Err(invalid_resource(
126 "logical backing projection must have non-zero size",
127 ));
128 }
129 let mut physical_cursor = 0_u64;
130 let mut covered = 0_u64;
131 let mut projection = Vec::new();
132 for segment in segments {
133 if physical_cursor >= physical_end {
134 break;
135 }
136 let segment_end = physical_cursor
137 .checked_add(segment.length_bytes())
138 .ok_or_else(|| invalid_resource("physical backing extent range overflows u64"))?;
139 let overlap_start = physical_cursor.max(physical_offset_bytes);
140 let overlap_end = segment_end.min(physical_end);
141 if overlap_start < overlap_end {
142 let within_segment = overlap_start - physical_cursor;
143 let translated_offset = segment
144 .offset_bytes()
145 .checked_add(within_segment)
146 .ok_or_else(|| invalid_resource("backing projection offset overflows u64"))?;
147 let length = overlap_end - overlap_start;
148 projection.push(BackingSegment::from_chunk(
149 segment.pool_id(),
150 segment.chunk_ordinal(),
151 segment.chunk_generation(),
152 translated_offset,
153 length,
154 )?);
155 covered = covered.checked_add(length).ok_or_else(|| {
156 invalid_resource("logical backing projection coverage overflows u64")
157 })?;
158 }
159 physical_cursor = segment_end;
160 }
161 if covered != size_bytes {
162 return Err(invalid_resource(
163 "logical backing projection exceeds its physical extent",
164 ));
165 }
166 Ok(projection)
167}
168
169#[derive(Debug, Clone, Copy)]
170pub(super) struct FreeExtent {
171 pub(super) chunk_generation: u64,
172 pub(super) length_bytes: u64,
173}
174
175#[derive(Debug, Clone, Default)]
176pub(super) struct FreeExtentIndex {
177 pub(super) by_offset: BTreeMap<(u32, u64), FreeExtent>,
178 pub(super) by_size: BTreeSet<(u64, u32, u64, u64)>,
179 pub(super) free_bytes: u64,
180 pub(super) search_probes: u64,
181}
182
183impl FreeExtentIndex {
184 fn rollback_segments(&mut self, segments: &[BackingSegment]) -> Result<(), VNextError> {
185 for segment in segments.iter().rev() {
186 self.release(segment)?;
187 }
188 Ok(())
189 }
190
191 fn with_rollback_context(
192 &mut self,
193 segments: &[BackingSegment],
194 error: VNextError,
195 ) -> VNextError {
196 match self.rollback_segments(segments) {
197 Ok(()) => error,
198 Err(rollback) => invalid_resource(format!(
199 "dynamic allocator failed and its journal rollback also failed: {error}; rollback: {rollback}"
200 )),
201 }
202 }
203
204 pub(super) fn insert_extent(
205 &mut self,
206 chunk_ordinal: u32,
207 chunk_generation: u64,
208 offset_bytes: u64,
209 length_bytes: u64,
210 ) -> Result<(), VNextError> {
211 if chunk_ordinal == 0
212 || chunk_generation == 0
213 || length_bytes == 0
214 || offset_bytes.checked_add(length_bytes).is_none()
215 || self.by_offset.contains_key(&(chunk_ordinal, offset_bytes))
216 {
217 return Err(invalid_resource("free extent identity or range is invalid"));
218 }
219 let end = offset_bytes + length_bytes;
220 if self
221 .by_offset
222 .range(..(chunk_ordinal, offset_bytes))
223 .next_back()
224 .is_some_and(|(&(ordinal, previous_offset), previous)| {
225 ordinal == chunk_ordinal && previous_offset + previous.length_bytes > offset_bytes
226 })
227 || self
228 .by_offset
229 .range((chunk_ordinal, offset_bytes)..)
230 .next()
231 .is_some_and(|(&(ordinal, next_offset), _)| {
232 ordinal == chunk_ordinal && next_offset < end
233 })
234 {
235 return Err(invalid_resource("free extent overlaps an existing extent"));
236 }
237 let next_free = self
238 .free_bytes
239 .checked_add(length_bytes)
240 .ok_or_else(|| invalid_resource("free extent bytes overflow u64"))?;
241 self.by_offset.insert(
242 (chunk_ordinal, offset_bytes),
243 FreeExtent {
244 chunk_generation,
245 length_bytes,
246 },
247 );
248 assert!(self
249 .by_size
250 .insert((length_bytes, chunk_ordinal, chunk_generation, offset_bytes,)));
251 self.free_bytes = next_free;
252 Ok(())
253 }
254
255 pub(super) fn remove_extent(
256 &mut self,
257 chunk_ordinal: u32,
258 offset_bytes: u64,
259 ) -> Result<FreeExtent, VNextError> {
260 let extent = *self
261 .by_offset
262 .get(&(chunk_ordinal, offset_bytes))
263 .ok_or_else(|| invalid_resource("free extent journal references a missing range"))?;
264 let size_key = (
265 extent.length_bytes,
266 chunk_ordinal,
267 extent.chunk_generation,
268 offset_bytes,
269 );
270 if !self.by_size.contains(&size_key) {
271 return Err(invalid_resource("free extent indexes diverged"));
272 }
273 let next_free_bytes = self
274 .free_bytes
275 .checked_sub(extent.length_bytes)
276 .ok_or_else(|| invalid_resource("free extent bytes underflowed"))?;
277 self.by_offset.remove(&(chunk_ordinal, offset_bytes));
278 assert!(self.by_size.remove(&size_key));
279 self.free_bytes = next_free_bytes;
280 Ok(extent)
281 }
282
283 pub(super) fn allocate_contiguous(
284 &mut self,
285 pool_id: &DynamicBackingPoolId,
286 size_bytes: u64,
287 ) -> Result<Option<BackingSegment>, VNextError> {
288 self.search_probes = self.search_probes.saturating_add(1);
289 let selected = self.by_size.range((size_bytes, 0, 0, 0)..).next().copied();
290 let Some((length_bytes, chunk_ordinal, chunk_generation, offset_bytes)) = selected else {
291 return Ok(None);
292 };
293 let segment = BackingSegment::from_chunk(
294 pool_id,
295 chunk_ordinal,
296 chunk_generation,
297 offset_bytes,
298 size_bytes,
299 )?;
300 let removed = self.remove_extent(chunk_ordinal, offset_bytes)?;
301 debug_assert_eq!(removed.length_bytes, length_bytes);
302 debug_assert_eq!(removed.chunk_generation, chunk_generation);
303 if size_bytes < length_bytes {
304 if let Err(error) = self.insert_extent(
305 chunk_ordinal,
306 chunk_generation,
307 offset_bytes + size_bytes,
308 length_bytes - size_bytes,
309 ) {
310 let restore =
311 self.insert_extent(chunk_ordinal, chunk_generation, offset_bytes, length_bytes);
312 return Err(match restore {
313 Ok(()) => error,
314 Err(rollback) => invalid_resource(format!(
315 "contiguous allocator failed and could not restore its selected extent: {error}; rollback: {rollback}"
316 )),
317 });
318 }
319 }
320 Ok(Some(segment))
321 }
322
323 pub(super) fn allocate_paged(
324 &mut self,
325 pool_id: &DynamicBackingPoolId,
326 size_bytes: u64,
327 block_bytes: u64,
328 ) -> Result<Option<Vec<BackingSegment>>, VNextError> {
329 if size_bytes == 0 || size_bytes % block_bytes != 0 {
330 return Err(invalid_resource(
331 "paged backing reservation is not block aligned",
332 ));
333 }
334 if self.free_bytes < size_bytes {
335 return Ok(None);
336 }
337 let mut remaining = size_bytes;
338 let mut segments = Vec::new();
339 while remaining != 0 {
340 self.search_probes = self.search_probes.saturating_add(1);
341 let Some((&(chunk_ordinal, offset_bytes), &extent)) = self.by_offset.first_key_value()
342 else {
343 self.rollback_segments(&segments)?;
344 return Ok(None);
345 };
346 if extent.length_bytes % block_bytes != 0 {
347 let error = invalid_resource("paged free extent lost fixed-block alignment");
348 return Err(self.with_rollback_context(&segments, error));
349 }
350 let take = extent.length_bytes.min(remaining);
351 let segment = match BackingSegment::from_chunk(
352 pool_id,
353 chunk_ordinal,
354 extent.chunk_generation,
355 offset_bytes,
356 take,
357 ) {
358 Ok(segment) => segment,
359 Err(error) => return Err(self.with_rollback_context(&segments, error)),
360 };
361 if let Err(error) = self.remove_extent(chunk_ordinal, offset_bytes) {
362 return Err(self.with_rollback_context(&segments, error));
363 }
364 if take < extent.length_bytes {
365 if let Err(error) = self.insert_extent(
366 chunk_ordinal,
367 extent.chunk_generation,
368 offset_bytes + take,
369 extent.length_bytes - take,
370 ) {
371 let restore = self.insert_extent(
372 chunk_ordinal,
373 extent.chunk_generation,
374 offset_bytes,
375 extent.length_bytes,
376 );
377 let error = match restore {
378 Ok(()) => error,
379 Err(rollback) => invalid_resource(format!(
380 "paged allocator failed and could not restore its selected extent: {error}; rollback: {rollback}"
381 )),
382 };
383 return Err(self.with_rollback_context(&segments, error));
384 }
385 }
386 segments.push(segment);
387 remaining -= take;
388 }
389 Ok(Some(segments))
390 }
391
392 pub(super) fn release(&mut self, segment: &BackingSegment) -> Result<(), VNextError> {
393 let chunk_ordinal = segment.chunk_ordinal();
394 let chunk_generation = segment.chunk_generation();
395 let mut offset_bytes = segment.offset_bytes();
396 let mut length_bytes = segment.length_bytes();
397 if let Some((&(ordinal, previous_offset), &previous)) = self
398 .by_offset
399 .range(..(chunk_ordinal, offset_bytes))
400 .next_back()
401 {
402 if ordinal == chunk_ordinal
403 && previous.chunk_generation == chunk_generation
404 && previous_offset + previous.length_bytes == offset_bytes
405 {
406 self.remove_extent(ordinal, previous_offset)?;
407 offset_bytes = previous_offset;
408 length_bytes = length_bytes
409 .checked_add(previous.length_bytes)
410 .ok_or_else(|| invalid_resource("coalesced free extent overflows u64"))?;
411 }
412 }
413 if let Some((&(ordinal, next_offset), &next)) =
414 self.by_offset.range((chunk_ordinal, offset_bytes)..).next()
415 {
416 if ordinal == chunk_ordinal
417 && next.chunk_generation == chunk_generation
418 && offset_bytes + length_bytes == next_offset
419 {
420 self.remove_extent(ordinal, next_offset)?;
421 length_bytes = length_bytes
422 .checked_add(next.length_bytes)
423 .ok_or_else(|| invalid_resource("coalesced free extent overflows u64"))?;
424 }
425 }
426 self.insert_extent(chunk_ordinal, chunk_generation, offset_bytes, length_bytes)
427 }
428
429 pub(super) fn largest_contiguous_bytes(&self) -> u64 {
430 self.by_size
431 .last()
432 .map_or(0, |(length_bytes, _, _, _)| *length_bytes)
433 }
434}