Skip to main content

we_trust_sqlite/writer/
pager.rs

1use super::*;
2use yykv_types::DsError;
3
4impl NativeWriter {
5    /// 获取数据库页大小
6    pub(super) async fn get_page_size(&self) -> Result<u32> {
7        self.ensure_initialized().await?;
8        let mut file = File::open(&self.path)
9            .await
10            .map_err(|e| DsError::io(format!("Failed to open database file: {}", e)))?;
11        let mut buf = [0u8; 2];
12        file.seek(SeekFrom::Start(16))
13            .await
14            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
15        file.read_exact(&mut buf)
16            .await
17            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
18        let ps = u16::from_be_bytes(buf);
19        Ok(if ps == 1 { 65536u32 } else { ps as u32 })
20    }
21
22    /// 分配一个新页面 (优先从 FreeList 获取)
23    pub async fn allocate_page(&self) -> Result<u32> {
24        self.ensure_initialized().await?;
25        let mut file = OpenOptions::new()
26            .read(true)
27            .write(true)
28            .open(&self.path)
29            .await
30            .map_err(|e| DsError::io(format!("Failed to open database file: {}", e)))?;
31        let page_size = self.get_page_size().await? as u64;
32
33        // 读取 FreeList Trunk Page ID (bytes 32-35) 和总计数 (bytes 36-39)
34        file.seek(SeekFrom::Start(32))
35            .await
36            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
37        let mut freelist_info = [0u8; 8];
38        file.read_exact(&mut freelist_info)
39            .await
40            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
41
42        let trunk_id = u32::from_be_bytes([
43            freelist_info[0],
44            freelist_info[1],
45            freelist_info[2],
46            freelist_info[3],
47        ]);
48        let free_count = u32::from_be_bytes([
49            freelist_info[4],
50            freelist_info[5],
51            freelist_info[6],
52            freelist_info[7],
53        ]);
54
55        if trunk_id > 0 && free_count > 0 {
56            let offset = (trunk_id as u64 - 1) * page_size;
57            let mut trunk_header = [0u8; 8];
58            file.seek(SeekFrom::Start(offset))
59                .await
60                .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
61            file.read_exact(&mut trunk_header)
62                .await
63                .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
64
65            let next_trunk = u32::from_be_bytes([
66                trunk_header[0],
67                trunk_header[1],
68                trunk_header[2],
69                trunk_header[3],
70            ]);
71            let leaf_count = u32::from_be_bytes([
72                trunk_header[4],
73                trunk_header[5],
74                trunk_header[6],
75                trunk_header[7],
76            ]);
77
78            let allocated_page_id;
79            if leaf_count > 0 {
80                // 从 Trunk 中取出一个叶子节点
81                let leaf_idx_offset = 8 + (leaf_count as u64 - 1) * 4;
82                file.seek(SeekFrom::Start(offset + leaf_idx_offset))
83                    .await
84                    .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
85                let mut leaf_id_buf = [0u8; 4];
86                file.read_exact(&mut leaf_id_buf)
87                    .await
88                    .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
89                allocated_page_id = u32::from_be_bytes(leaf_id_buf);
90
91                // 更新 leaf_count
92                file.seek(SeekFrom::Start(offset + 4))
93                    .await
94                    .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
95                file.write_all(&(leaf_count - 1).to_be_bytes())
96                    .await
97                    .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
98            } else {
99                // 使用 Trunk Page 本身
100                allocated_page_id = trunk_id;
101                // 更新 Header 指向下一个 Trunk
102                file.seek(SeekFrom::Start(32))
103                    .await
104                    .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
105                file.write_all(&next_trunk.to_be_bytes())
106                    .await
107                    .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
108            }
109
110            // 更新总 FreeList 计数
111            file.seek(SeekFrom::Start(36))
112                .await
113                .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
114            file.write_all(&(free_count - 1).to_be_bytes())
115                .await
116                .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
117
118            return Ok(allocated_page_id);
119        }
120
121        // 没可用页面,追加到文件末尾
122        let file_len = file
123            .metadata()
124            .await
125            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?
126            .len();
127        let new_page_id = (file_len / page_size) as u32 + 1;
128
129        // 更新数据库大小 (bytes 28-31)
130        file.seek(SeekFrom::Start(28))
131            .await
132            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
133        file.write_all(&new_page_id.to_be_bytes())
134            .await
135            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
136
137        // 填充新页面并确保文件增长
138        file.seek(SeekFrom::Start(file_len))
139            .await
140            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
141        file.write_all(&vec![0u8; page_size as usize])
142            .await
143            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
144
145        Ok(new_page_id)
146    }
147
148    /// 将页面释放回 FreeList
149    pub async fn deallocate_page(&self, page_id: u32) -> Result<()> {
150        if page_id == 0 {
151            return Ok(());
152        }
153
154        let mut file = OpenOptions::new()
155            .read(true)
156            .write(true)
157            .open(&self.path)
158            .await
159            .map_err(|e| DsError::io(format!("Failed to open database file: {}", e)))?;
160        let page_size = self.get_page_size().await?;
161
162        // 读取 Header 中的 FreeList 信息 (trunk_id: 32-35, total_count: 36-39)
163        file.seek(SeekFrom::Start(32))
164            .await
165            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
166        let mut freelist_info = [0u8; 8];
167        file.read_exact(&mut freelist_info)
168            .await
169            .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
170        let trunk_id = u32::from_be_bytes([
171            freelist_info[0],
172            freelist_info[1],
173            freelist_info[2],
174            freelist_info[3],
175        ]);
176        let free_count = u32::from_be_bytes([
177            freelist_info[4],
178            freelist_info[5],
179            freelist_info[6],
180            freelist_info[7],
181        ]);
182
183        let mut added_as_trunk = false;
184        if trunk_id > 0 {
185            // 尝试放入第一个 Trunk Page
186            let trunk_offset = (trunk_id as u64 - 1) * page_size as u64;
187            file.seek(SeekFrom::Start(trunk_offset + 4))
188                .await
189                .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
190            let mut leaf_count_buf = [0u8; 4];
191            file.read_exact(&mut leaf_count_buf)
192                .await
193                .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
194            let leaf_count = u32::from_be_bytes(leaf_count_buf);
195
196            let max_leaves = (page_size - 8) / 4;
197            if leaf_count < max_leaves {
198                // 有空间,作为叶子节点插入
199                let leaf_pos = trunk_offset + 8 + (leaf_count as u64 * 4);
200                file.seek(SeekFrom::Start(leaf_pos))
201                    .await
202                    .map_err(|e| DsError::io_raw(e, Some(self.path.clone().into())))?;
203                file.write_all(&page_id.to_be_bytes())
204                    .await
205                    .map_err(|e| DsError::io(e.to_string()))?;
206
207                // 更新 leaf_count
208                file.seek(SeekFrom::Start(trunk_offset + 4))
209                    .await
210                    .map_err(|e| DsError::io(e.to_string()))?;
211                file.write_all(&(leaf_count + 1).to_be_bytes())
212                    .await
213                    .map_err(|e| DsError::io(e.to_string()))?;
214            } else {
215                // Trunk Page 已满,将释放的页面作为新的第一个 Trunk Page
216                added_as_trunk = true;
217            }
218        } else {
219            // 没有 Trunk Page,将释放的页面作为第一个 Trunk Page
220            added_as_trunk = true;
221        }
222
223        if added_as_trunk {
224            let offset = (page_id as u64 - 1) * page_size as u64;
225            let mut trunk_header = [0u8; 8];
226            // Next trunk ID = 旧的 trunk_id
227            trunk_header[0..4].copy_from_slice(&trunk_id.to_be_bytes());
228            // Leaf count = 0
229            trunk_header[4..8].copy_from_slice(&0u32.to_be_bytes());
230
231            file.seek(SeekFrom::Start(offset))
232                .await
233                .map_err(|e| DsError::io(e.to_string()))?;
234            file.write_all(&trunk_header)
235                .await
236                .map_err(|e| DsError::io(e.to_string()))?;
237
238            // 更新 Header 指向新的 Trunk Page
239            file.seek(SeekFrom::Start(32))
240                .await
241                .map_err(|e| DsError::io(e.to_string()))?;
242            file.write_all(&page_id.to_be_bytes())
243                .await
244                .map_err(|e| DsError::io(e.to_string()))?;
245        }
246
247        // 更新总 FreeList 计数
248        file.seek(SeekFrom::Start(36))
249            .await
250            .map_err(|e| DsError::io(e.to_string()))?;
251        file.write_all(&(free_count + 1).to_be_bytes())
252            .await
253            .map_err(|e| DsError::io(e.to_string()))?;
254
255        Ok(())
256    }
257
258    /// 初始化一个新的“伪 SQLite”文件头
259    pub async fn init_file(&self) -> Result<()> {
260        let mut file = OpenOptions::new()
261            .write(true)
262            .create(true)
263            .truncate(true)
264            .open(&self.path)
265            .await
266            .map_err(|e| DsError::io(format!("Failed to open database file: {}", e)))?;
267
268        let mut header = vec![0u8; 100];
269        header[0..16].copy_from_slice(crate::reader::SQLITE_MAGIC);
270
271        // Page Size: 4096 (0x1000)
272        header[16] = 0x10;
273        header[17] = 0x00;
274
275        // Write/Read version
276        header[18] = 1;
277        header[19] = 1;
278
279        file.write_all(&header)
280            .await
281            .map_err(|e| DsError::io(e.to_string()))?;
282
283        Ok(())
284    }
285
286    /// 写入一个完整页面内容
287    pub async fn write_page(&self, page_id: u32, data: &[u8]) -> Result<()> {
288        let page_size = self.get_page_size().await?;
289        if data.len() != page_size as usize {
290            return Err(DsError::storage(format!(
291                "Invalid page data size: expected {}, got {}",
292                page_size,
293                data.len()
294            )));
295        }
296
297        let mut file = OpenOptions::new()
298            .write(true)
299            .open(&self.path)
300            .await
301            .map_err(|e| DsError::io(format!("Failed to open database file: {}", e)))?;
302        let offset = (page_id as u64 - 1) * (page_size as u64);
303        file.seek(SeekFrom::Start(offset))
304            .await
305            .map_err(|e| DsError::io(e.to_string()))?;
306        file.write_all(data)
307            .await
308            .map_err(|e| DsError::io(e.to_string()))?;
309        Ok(())
310    }
311}