Line data Source code
1 : //! This file contains generic utility functions over the interface types,
2 : //! which could be handy for any compaction implementation.
3 : use crate::interface::*;
4 :
5 : use futures::future::BoxFuture;
6 : use futures::{Stream, StreamExt};
7 : use itertools::Itertools;
8 : use pageserver_api::shard::ShardIdentity;
9 : use pin_project_lite::pin_project;
10 : use std::collections::BinaryHeap;
11 : use std::collections::VecDeque;
12 : use std::fmt::Display;
13 : use std::future::Future;
14 : use std::ops::{DerefMut, Range};
15 : use std::pin::Pin;
16 : use std::task::{ready, Poll};
17 : use utils::lsn::Lsn;
18 :
19 : pub const PAGE_SZ: u64 = 8192;
20 :
21 6 : pub fn keyspace_total_size<K>(
22 6 : keyspace: &CompactionKeySpace<K>,
23 6 : shard_identity: &ShardIdentity,
24 6 : ) -> u64
25 6 : where
26 6 : K: CompactionKey,
27 6 : {
28 6 : keyspace
29 6 : .iter()
30 6 : .map(|r| K::key_range_size(r, shard_identity) as u64)
31 6 : .sum()
32 6 : }
33 :
34 27098 : pub fn overlaps_with<T: Ord>(a: &Range<T>, b: &Range<T>) -> bool {
35 27098 : !(a.end <= b.start || b.end <= a.start)
36 27098 : }
37 :
38 4794 : pub fn union_to_keyspace<K: Ord>(a: &mut CompactionKeySpace<K>, b: CompactionKeySpace<K>) {
39 4794 : let x = std::mem::take(a);
40 4794 : let mut all_ranges_iter = [x.into_iter(), b.into_iter()]
41 4794 : .into_iter()
42 4794 : .kmerge_by(|a, b| a.start < b.start);
43 4794 : let mut ranges = Vec::new();
44 4794 : if let Some(first) = all_ranges_iter.next() {
45 4794 : let (mut start, mut end) = (first.start, first.end);
46 :
47 9584 : for r in all_ranges_iter {
48 4790 : assert!(r.start >= start);
49 4790 : if r.start > end {
50 0 : ranges.push(start..end);
51 0 : start = r.start;
52 0 : end = r.end;
53 4790 : } else if r.end > end {
54 0 : end = r.end;
55 4790 : }
56 : }
57 4794 : ranges.push(start..end);
58 0 : }
59 4794 : *a = ranges
60 4794 : }
61 :
62 6 : pub fn intersect_keyspace<K: Ord + Clone + Copy>(
63 6 : a: &CompactionKeySpace<K>,
64 6 : r: &Range<K>,
65 6 : ) -> CompactionKeySpace<K> {
66 6 : let mut ranges: Vec<Range<K>> = Vec::new();
67 :
68 6 : for x in a.iter() {
69 6 : if x.end <= r.start {
70 0 : continue;
71 6 : }
72 6 : if x.start >= r.end {
73 0 : break;
74 6 : }
75 6 : ranges.push(x.clone())
76 : }
77 :
78 : // trim the ends
79 6 : if let Some(first) = ranges.first_mut() {
80 6 : first.start = std::cmp::max(first.start, r.start);
81 6 : }
82 6 : if let Some(last) = ranges.last_mut() {
83 6 : last.end = std::cmp::min(last.end, r.end);
84 6 : }
85 6 : ranges
86 6 : }
87 :
88 : /// Create a stream that iterates through all DeltaEntrys among all input
89 : /// layers, in key-lsn order.
90 : ///
91 : /// This is public because the create_delta() implementation likely wants to use this too
92 : /// TODO: move to a more shared place
93 98 : pub fn merge_delta_keys<'a, E: CompactionJobExecutor>(
94 98 : layers: &'a [E::DeltaLayer],
95 98 : ctx: &'a E::RequestContext,
96 98 : ) -> MergeDeltaKeys<'a, E> {
97 98 : // Use a binary heap to merge the layers. Each input layer is initially
98 98 : // represented by a LazyLoadLayer::Unloaded element, which uses the start of
99 98 : // the layer's key range as the key. The first time a layer reaches the top
100 98 : // of the heap, all the keys of the layer are loaded into a sorted vector.
101 98 : //
102 98 : // This helps to keep the memory usage reasonable: we only need to hold in
103 98 : // memory the DeltaEntrys of the layers that overlap with the "current" key.
104 98 : let mut heap: BinaryHeap<LazyLoadLayer<'a, E>> = BinaryHeap::new();
105 3220 : for l in layers {
106 3122 : heap.push(LazyLoadLayer::Unloaded(l));
107 3122 : }
108 98 : MergeDeltaKeys {
109 98 : heap,
110 98 : ctx,
111 98 : load_future: None,
112 98 : }
113 98 : }
114 :
115 6 : pub async fn merge_delta_keys_buffered<'a, E: CompactionJobExecutor + 'a>(
116 6 : layers: &'a [E::DeltaLayer],
117 6 : ctx: &'a E::RequestContext,
118 6 : ) -> anyhow::Result<impl Stream<Item = <E::DeltaLayer as CompactionDeltaLayer<E>>::DeltaEntry<'a>>>
119 6 : {
120 6 : let mut keys = Vec::new();
121 100 : for l in layers {
122 : // Boxing and casting to LoadFuture is required to obtain the right Sync bound.
123 : // If we do l.load_keys(ctx).await? directly, there is a compilation error.
124 94 : let load_future: LoadFuture<'a, _> = Box::pin(l.load_keys(ctx));
125 94 : keys.extend(load_future.await?.into_iter());
126 : }
127 10632670 : keys.sort_by_key(|k| (k.key(), k.lsn()));
128 6 : let stream = futures::stream::iter(keys.into_iter());
129 6 : Ok(stream)
130 6 : }
131 :
132 : enum LazyLoadLayer<'a, E: CompactionJobExecutor> {
133 : Loaded(VecDeque<<E::DeltaLayer as CompactionDeltaLayer<E>>::DeltaEntry<'a>>),
134 : Unloaded(&'a E::DeltaLayer),
135 : }
136 : impl<'a, E: CompactionJobExecutor> LazyLoadLayer<'a, E> {
137 81841588 : fn min_key(&self) -> E::Key {
138 81841588 : match self {
139 67232654 : Self::Loaded(entries) => entries.front().unwrap().key(),
140 14608934 : Self::Unloaded(dl) => dl.key_range().start,
141 : }
142 81841588 : }
143 81841588 : fn min_lsn(&self) -> Lsn {
144 81841588 : match self {
145 67232654 : Self::Loaded(entries) => entries.front().unwrap().lsn(),
146 14608934 : Self::Unloaded(dl) => dl.lsn_range().start,
147 : }
148 81841588 : }
149 : }
150 : impl<'a, E: CompactionJobExecutor> PartialOrd for LazyLoadLayer<'a, E> {
151 40920794 : fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
152 40920794 : Some(self.cmp(other))
153 40920794 : }
154 : }
155 : impl<'a, E: CompactionJobExecutor> Ord for LazyLoadLayer<'a, E> {
156 40920794 : fn cmp(&self, other: &Self) -> std::cmp::Ordering {
157 40920794 : // reverse order so that we get a min-heap
158 40920794 : (other.min_key(), other.min_lsn()).cmp(&(self.min_key(), self.min_lsn()))
159 40920794 : }
160 : }
161 : impl<'a, E: CompactionJobExecutor> PartialEq for LazyLoadLayer<'a, E> {
162 0 : fn eq(&self, other: &Self) -> bool {
163 0 : self.cmp(other) == std::cmp::Ordering::Equal
164 0 : }
165 : }
166 : impl<'a, E: CompactionJobExecutor> Eq for LazyLoadLayer<'a, E> {}
167 :
168 : type LoadFuture<'a, E> = BoxFuture<'a, anyhow::Result<Vec<E>>>;
169 :
170 : // Stream returned by `merge_delta_keys`
171 : pin_project! {
172 : #[allow(clippy::type_complexity)]
173 : pub struct MergeDeltaKeys<'a, E: CompactionJobExecutor> {
174 : heap: BinaryHeap<LazyLoadLayer<'a, E>>,
175 :
176 : #[pin]
177 : load_future: Option<LoadFuture<'a, <E::DeltaLayer as CompactionDeltaLayer<E>>::DeltaEntry<'a>>>,
178 :
179 : ctx: &'a E::RequestContext,
180 : }
181 : }
182 :
183 : impl<'a, E> Stream for MergeDeltaKeys<'a, E>
184 : where
185 : E: CompactionJobExecutor + 'a,
186 : {
187 : type Item = anyhow::Result<<E::DeltaLayer as CompactionDeltaLayer<E>>::DeltaEntry<'a>>;
188 :
189 11045220 : fn poll_next(
190 11045220 : self: Pin<&mut Self>,
191 11045220 : cx: &mut std::task::Context<'_>,
192 11045220 : ) -> Poll<std::option::Option<<Self as futures::Stream>::Item>> {
193 11045220 : let mut this = self.project();
194 : loop {
195 11048342 : if let Some(mut load_future) = this.load_future.as_mut().as_pin_mut() {
196 : // We are waiting for loading the keys to finish
197 3122 : match ready!(load_future.as_mut().poll(cx)) {
198 3122 : Ok(entries) => {
199 3122 : this.load_future.set(None);
200 3122 : *this.heap.peek_mut().unwrap() =
201 3122 : LazyLoadLayer::Loaded(VecDeque::from(entries));
202 3122 : }
203 0 : Err(e) => {
204 0 : return Poll::Ready(Some(Err(e)));
205 : }
206 : }
207 11045220 : }
208 :
209 : // If the topmost layer in the heap hasn't been loaded yet, start
210 : // loading it. Otherwise return the next entry from it and update
211 : // the layer's position in the heap (this decreaseKey operation is
212 : // performed implicitly when `top` is dropped).
213 11048342 : if let Some(mut top) = this.heap.peek_mut() {
214 11048244 : match top.deref_mut() {
215 3122 : LazyLoadLayer::Unloaded(ref mut l) => {
216 3122 : let fut = l.load_keys(this.ctx);
217 3122 : this.load_future.set(Some(Box::pin(fut)));
218 3122 : continue;
219 : }
220 11045122 : LazyLoadLayer::Loaded(ref mut entries) => {
221 11045122 : let result = entries.pop_front().unwrap();
222 11045122 : if entries.is_empty() {
223 3122 : std::collections::binary_heap::PeekMut::pop(top);
224 11042000 : }
225 11045122 : return Poll::Ready(Some(Ok(result)));
226 : }
227 : }
228 : } else {
229 98 : return Poll::Ready(None);
230 : }
231 : }
232 11045220 : }
233 : }
234 :
235 : // Accumulate values at key boundaries
236 : pub struct KeySize<K> {
237 : pub key: K,
238 : pub num_values: u64,
239 : pub size: u64,
240 : /// The lsns to partition at (if empty then no per-lsn partitioning)
241 : pub partition_lsns: Vec<(Lsn, u64)>,
242 : }
243 :
244 6 : pub fn accum_key_values<'a, I, K, D, E>(
245 6 : input: I,
246 6 : target_size: u64,
247 6 : ) -> impl Stream<Item = Result<KeySize<K>, E>>
248 6 : where
249 6 : K: Eq + PartialOrd + Display + Copy,
250 6 : I: Stream<Item = Result<D, E>>,
251 6 : D: CompactionDeltaEntry<'a, K>,
252 6 : {
253 : async_stream::try_stream! {
254 : // Initialize the state from the first value
255 : let mut input = std::pin::pin!(input);
256 :
257 : if let Some(first) = input.next().await {
258 : let first = first?;
259 : let mut part_size = first.size();
260 : let mut accum: KeySize<K> = KeySize {
261 : key: first.key(),
262 : num_values: 1,
263 : size: part_size,
264 : partition_lsns: Vec::new(),
265 : };
266 : let mut last_key = accum.key;
267 : while let Some(this) = input.next().await {
268 : let this = this?;
269 : if this.key() == accum.key {
270 : let add_size = this.size();
271 : if part_size + add_size > target_size {
272 : accum.partition_lsns.push((this.lsn(), part_size));
273 : part_size = 0;
274 : }
275 : part_size += add_size;
276 : accum.size += add_size;
277 : accum.num_values += 1;
278 : } else {
279 : assert!(last_key <= accum.key, "last_key={last_key} <= accum.key={}", accum.key);
280 : last_key = accum.key;
281 : yield accum;
282 : part_size = this.size();
283 : accum = KeySize {
284 : key: this.key(),
285 : num_values: 1,
286 : size: part_size,
287 : partition_lsns: Vec::new(),
288 : };
289 : }
290 : }
291 : assert!(last_key <= accum.key, "last_key={last_key} <= accum.key={}", accum.key);
292 : yield accum;
293 : }
294 : }
295 6 : }
|
