1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
use std::cell::UnsafeCell;
use std::io::ErrorKind;
use std::ops::Range;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use fs_err as fs;
use parking_lot::Mutex;
use roaring::RoaringBitmap;
use crate::common::generic_consts::AccessPattern;
use crate::common::universal_io::simple_disk_cache::BLOCK_SIZE;
use crate::common::universal_io::{
MmapFile, MmapFs, OpenOptions, Populate, UioResult, UniversalIoError, UniversalRead,
UniversalReadFs, UniversalWrite, mmap as mmap_file,
};
#[derive(Debug)]
pub(crate) struct LocalState {
/// UnsafeCell so that we can write to it under non-mut reference.
/// Such as when the pipeline reads from remote first.
pub mmap: UnsafeCell<MmapFile>,
/// Bitmask to know which blocks have been fetched so far.
pub fetched: Mutex<RoaringBitmap>,
/// Fast-path flag: when true, the mmap is fully populated and the
/// `fetched` bitmap can be skipped on the read hot path.
pub fully_populated: AtomicBool,
}
unsafe impl Sync for LocalState {}
impl LocalState {
pub(super) fn new(
local_path: impl AsRef<Path>,
len: u64,
options: OpenOptions,
) -> UioResult<Self> {
if let Some(parent) = local_path.as_ref().parent() {
fs::create_dir_all(parent)?;
}
let OpenOptions {
writeable: _, // always needs to be writeable
need_sequential,
populate: _, // this is handled externally to LocalState
advice,
} = options;
let file = fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(true)
.open(local_path.as_ref())?;
file.set_len(len)?;
let mmap = MmapFs.open(
local_path.as_ref(),
OpenOptions {
writeable: true,
need_sequential,
populate: Populate::No,
advice,
},
(),
)?;
Ok(LocalState {
mmap: UnsafeCell::new(mmap),
fetched: Mutex::new(RoaringBitmap::new()),
fully_populated: AtomicBool::new(false),
})
}
pub(super) fn resize(&mut self, local_path: impl AsRef<Path>, new_len: u64) -> UioResult<()> {
let mmap = self.mmap.get_mut();
let current_len = mmap.len::<u8>()?;
if current_len == new_len {
return Ok(());
}
if current_len > new_len {
return Err(UniversalIoError::Io(std::io::Error::new(
ErrorKind::Unsupported,
"Shrinking the file is not supported",
)));
}
let file = fs::OpenOptions::new()
.read(true)
.write(true)
.create(false)
.open(local_path.as_ref())?;
file.set_len(new_len)?;
// We just sized the file ourselves — no need for a full `reopen`
// that stats it again.
mmap.grow_mapping(new_len)?;
*self.fully_populated.get_mut() = false;
// The previous tail block may have been only partially populated (its
// fetch was clamped to the old EOF). `set_len` zero-filled the bytes
// between the old EOF and the next block boundary, so the block is no
// longer accurate. Drop it from `fetched` to force a re-fetch on the
// next read.
if current_len % BLOCK_SIZE as u64 != 0 {
let partial_tail_block = (current_len / BLOCK_SIZE as u64) as u32;
self.fetched.lock().remove(partial_tail_block);
}
Ok(())
}
pub(super) fn mmap(&self) -> &MmapFile {
// SAFETY: we have `&self` reference
unsafe { self.mmap.get().as_ref_unchecked() }
}
/// Whether `blocks_range` is already cached locally.
///
/// Cheap when the file is fully populated (one relaxed atomic load);
/// otherwise locks `fetched` and checks the bitmap.
pub(super) fn contains(&self, blocks_range: Range<u32>) -> bool {
if self.fully_populated.load(Ordering::Acquire) {
return true;
}
self.fetched.lock().contains_range(blocks_range)
}
/// # Safety
/// `byte_range` must have been populated first, caller must ensure `self.fetched` references the
/// blocks for the byte range.
pub(super) unsafe fn read_mmap_bytes<P: AccessPattern>(
&self,
range: Range<u64>,
) -> UioResult<&[u8]> {
let mmap_bytes = self.mmap().as_bytes::<P>();
mmap_file::read_bytes(mmap_bytes, range)
}
/// # Safety
/// `DiskCache` is only used for append-only remotes with an immutable
/// prefix, so every fetch of `blocks_range` yields the same bytes the
/// mirror already holds for it. Since `blocks_range` can include
/// already-fetched data, it is possible that some sections get
/// overwritten; however, it is the same data, so it is fine.
///
/// Assumes the bytes slice covers the entirety of `blocks_range`.
pub(super) unsafe fn write_mmap_bytes(&self, bytes: &[u8], blocks_range: Range<u32>) {
// SAFETY:
// 1. The remote's existing bytes are immutable, so worst case, same
// data is overwritten.
// 2. The `fetched` bitmap should track which blocks are already present.
let mmap = unsafe { self.mmap.get().as_mut_unchecked() };
if self.fully_populated.load(Ordering::Acquire) {
return;
}
let mut fetched = self.fetched.lock();
if fetched.contains_range(blocks_range.clone()) {
return;
}
let byte_offset = (blocks_range.start as usize * BLOCK_SIZE) as u64;
let max_len = mmap
.len::<u8>()
.expect("MmapFile::len is infallible")
.saturating_sub(byte_offset);
assert_eq!(
bytes.len() as u64,
max_len.min((blocks_range.len() * BLOCK_SIZE) as u64)
);
mmap.write(byte_offset, bytes)
.expect("MmapFile::write is infallible");
fetched.insert_range(blocks_range);
// If every block is now populated, turn `fully_populated` on.
let total_blocks = mmap
.len::<u8>()
.expect("MmapFile::len is infallible")
.div_ceil(BLOCK_SIZE as u64);
if fetched.len() == total_blocks {
self.fully_populated.store(true, Ordering::Release);
}
}
}