Skip to main content

agentmux_srv/backend/storage/filestore/
offset_ops.rs

1// Copyright 2025, Command Line Inc.
2// Copyright 2025-2026, AgentMux Corp.
3// SPDX-License-Identifier: Apache-2.0
4
5//! FileStore offset-based read/write operations.
6
7
8use std::collections::HashMap;
9
10use rusqlite::{params, OptionalExtension};
11
12use super::core::{FileStore, PART_DATA_SIZE};
13use super::counter::DROP_EPOCH_SQL;
14use super::FileOpts;
15use crate::backend::storage::error::StoreError;
16
17impl FileStore {
18    /// Write data at a specific offset.
19    /// The offset must be <= current file size.
20    #[allow(dead_code)]
21    #[track_caller]
22    pub fn write_at(
23        &self,
24        zone_id: &str,
25        name: &str,
26        offset: i64,
27        data: &[u8],
28    ) -> Result<(), StoreError> {
29        if data.is_empty() {
30            return Ok(());
31        }
32
33        let key = (zone_id.to_string(), name.to_string());
34        let now = agentmux_common::time::now_ms();
35
36        let pds = PART_DATA_SIZE as i64;
37
38        // Size and opts read from the database in the same transaction as the
39        // writes (see `write_txn`): a cached size can predate another
40        // process's append, and writing with it would shrink the file.
41        let written = self.write_txn(|tx| {
42            let (size, opts_json): (i64, String) = tx
43                .query_row(
44                    "SELECT size, opts FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
45                    params![zone_id, name],
46                    |row| Ok((row.get(0)?, row.get(1)?)),
47                )
48                .optional()?
49                .ok_or(StoreError::NotFound)?;
50            let opts: FileOpts = serde_json::from_str(&opts_json).unwrap_or_default();
51            if offset > size {
52                return Err(StoreError::Other(format!(
53                    "offset {} exceeds file size {}",
54                    offset, size
55                )));
56            }
57
58            let new_size = std::cmp::max(size, offset + data.len() as i64);
59
60            // Handle circular file data truncation
61            let (actual_offset, actual_data) = if opts.circular && opts.maxsize > 0 {
62                let start_cir_offset = new_size - opts.maxsize;
63                if start_cir_offset > 0 {
64                    let end = offset + data.len() as i64;
65                    if end <= start_cir_offset {
66                        // Entire write is before the circular window — no-op
67                        return Ok(None);
68                    }
69                    if offset < start_cir_offset {
70                        let skip = (start_cir_offset - offset) as usize;
71                        (start_cir_offset, &data[skip..])
72                    } else {
73                        (offset, data)
74                    }
75                } else {
76                    (offset, data)
77                }
78            } else {
79                (offset, data)
80            };
81
82            // Compute affected parts
83            let start_part = (actual_offset / pds) as i32;
84            let end_part = ((actual_offset + actual_data.len() as i64 - 1) / pds) as i32;
85            let mut data_pos = 0usize;
86
87            for part_idx in start_part..=end_part {
88                let part_start = part_idx as i64 * pds;
89                let offset_in_part = if part_idx == start_part {
90                    (actual_offset - part_start) as usize
91                } else {
92                    0
93                };
94
95                // Load existing part if needed
96                let mut part_data: Vec<u8> = tx
97                    .query_row(
98                        "SELECT data FROM db_file_data WHERE zoneid = ?1 AND name = ?2 AND partidx = ?3",
99                        params![zone_id, name, part_idx],
100                        |row| row.get(0),
101                    )
102                    .optional()?
103                    .unwrap_or_default();
104                // Ensure part is large enough
105                if part_data.len() < offset_in_part {
106                    part_data.resize(offset_in_part, 0);
107                }
108
109                // Copy data into part
110                let remaining = actual_data.len() - data_pos;
111                let space = PART_DATA_SIZE - offset_in_part;
112                let to_copy = remaining.min(space);
113
114                if offset_in_part < part_data.len() {
115                    // Overwrite existing bytes
116                    let overwrite_end = (offset_in_part + to_copy).min(part_data.len());
117                    let overwrite_len = overwrite_end - offset_in_part;
118                    part_data[offset_in_part..offset_in_part + overwrite_len]
119                        .copy_from_slice(&actual_data[data_pos..data_pos + overwrite_len]);
120                    if to_copy > overwrite_len {
121                        part_data.extend_from_slice(
122                            &actual_data[data_pos + overwrite_len..data_pos + to_copy],
123                        );
124                    }
125                } else {
126                    part_data.extend_from_slice(&actual_data[data_pos..data_pos + to_copy]);
127                }
128
129                tx.execute(
130                    "REPLACE INTO db_file_data (zoneid, name, partidx, data) VALUES (?1, ?2, ?3, ?4)",
131                    params![zone_id, name, part_idx, part_data],
132                )?;
133                data_pos += to_copy;
134            }
135
136            // An in-place overwrite can change any line: the counter can't
137            // follow it, so the epoch is dropped (counter.rs).
138            tx.execute(
139                &format!(
140                    "UPDATE db_wave_file SET size = ?1, modts = ?2, {DROP_EPOCH_SQL}
141                     WHERE zoneid = ?3 AND name = ?4"
142                ),
143                params![new_size, now, zone_id, name],
144            )?;
145            Ok(Some((new_size, start_part, end_part)))
146        })?;
147        let Some((new_size, start_part, end_part)) = written else {
148            return Ok(());
149        };
150
151        // Update cache
152        let mut cache = self.cache.lock().unwrap();
153        if let Some(entry) = cache.get_mut(&key) {
154            if let Some(ref mut f) = entry.file {
155                f.size = new_size;
156                f.modts = now;
157            }
158            // Invalidate cached data entries for affected parts
159            for part_idx in start_part..=end_part {
160                entry.data_entries.remove(&part_idx);
161            }
162        }
163
164        Ok(())
165    }
166
167    /// Read data at a specific offset and size.
168    /// For circular files, adjusts offset if it falls before valid data range.
169    /// Returns (adjusted_offset, data).
170    #[track_caller]
171    pub fn read_at(
172        &self,
173        zone_id: &str,
174        name: &str,
175        offset: i64,
176        size: i64,
177    ) -> Result<(i64, Vec<u8>), StoreError> {
178        let file = self
179            .stat(zone_id, name)?
180            .ok_or(StoreError::NotFound)?;
181
182        if file.size == 0 {
183            return Ok((0, Vec::new()));
184        }
185
186        let data_len = file.data_length();
187        let data_start = file.data_start_idx();
188
189        // Adjust offset for circular files
190        let mut actual_offset = offset;
191        let mut actual_size = if size == 0 { data_len } else { size };
192
193        if file.opts.circular && file.opts.maxsize > 0 {
194            if actual_offset < data_start {
195                let skip = data_start - actual_offset;
196                actual_offset = data_start;
197                actual_size -= skip;
198            }
199            if actual_size <= 0 {
200                return Ok((data_start, Vec::new()));
201            }
202        }
203
204        // Clamp to available data
205        if actual_offset >= file.size {
206            return Ok((actual_offset, Vec::new()));
207        }
208        let available = file.size - actual_offset;
209        actual_size = actual_size.min(available);
210
211        if actual_size <= 0 {
212            return Ok((actual_offset, Vec::new()));
213        }
214
215        let pds = PART_DATA_SIZE as i64;
216        let start_part = (actual_offset / pds) as i32;
217        let end_part = ((actual_offset + actual_size - 1) / pds) as i32;
218
219        // Load parts from DB: one range query, not one query per 64 KB part —
220        // reading an 800 MB transcript part by part was 12,270 queries and
221        // took 12 s in a release build.
222        let conn = self.lock_conn();
223        let mut parts_map: HashMap<i32, Vec<u8>> = HashMap::new();
224        {
225            let mut stmt = conn.prepare_cached(
226                "SELECT partidx, data FROM db_file_data \
227                 WHERE zoneid = ?1 AND name = ?2 AND partidx BETWEEN ?3 AND ?4",
228            )?;
229            let rows = stmt.query_map(params![zone_id, name, start_part, end_part], |row| {
230                Ok((row.get::<_, i32>(0)?, row.get::<_, Vec<u8>>(1)?))
231            })?;
232            for row in rows {
233                let (part_idx, data) = row?;
234                parts_map.insert(part_idx, data);
235            }
236        }
237        drop(conn);
238
239        // Assemble result
240        let mut result = Vec::with_capacity(actual_size as usize);
241        for part_idx in start_part..=end_part {
242            if let Some(part_data) = parts_map.get(&part_idx) {
243                let part_start = part_idx as i64 * pds;
244                let skip = if part_start < actual_offset {
245                    (actual_offset - part_start) as usize
246                } else {
247                    0
248                };
249                let remaining = actual_size as usize - result.len();
250                let take = remaining.min(part_data.len().saturating_sub(skip));
251                if take > 0 {
252                    result.extend_from_slice(&part_data[skip..skip + take]);
253                }
254            }
255        }
256
257        Ok((actual_offset, result))
258    }
259}