forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
4.9 KiB
1//! Blob storage for release assets. A thin wrapper over object_store's local filesystem backend,
2//! so an S3 backend can replace it later without touching the pages.
3//!
4//! Keys are `assets/<asset id>`. Writes stream through object_store's multipart upload, which
5//! assembles the object in a temporary file and renames it into place on `finish`. Every chunk
6//! is hashed as it passes, so the sha256 is known without a second read. Blobs are immutable:
7//! an asset is never rewritten, only added or deleted.
8
9use std::path::Path;
10use std::sync::Arc;
11
12use bytes::Bytes;
13use forge_core::{Error, Result};
14use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
15use object_store::local::LocalFileSystem;
16use object_store::path::Path as ObjectPath;
17use object_store::{ObjectStore, ObjectStoreExt, WriteMultipart};
18use sha2::{Digest, Sha256};
19
20/// Bytes buffered per upload part before it is handed to the store.
21const PART_BYTES: usize = 8 * 1024 * 1024;
22/// Parts uploaded at once while a stream is being written.
23const PARTS_IN_FLIGHT: usize = 4;
24
25/// The object key of a release asset.
26pub fn asset_key(asset_id: &str) -> String {
27 format!("assets/{asset_id}")
28}
29
30fn storage_err<E>(e: E) -> Error
31where
32 E: std::error::Error + Send + Sync + 'static,
33{
34 Error::Io(std::io::Error::other(e))
35}
36
37/// The blob store rooted at `<data_dir>/blobs`. Cheap to clone.
38#[derive(Clone)]
39pub struct BlobStore {
40 store: Arc<dyn ObjectStore>,
41}
42
43/// A blob opened for reading: its size and its bytes.
44pub struct BlobRead {
45 pub size: u64,
46 pub stream: BoxStream<'static, std::io::Result<Bytes>>,
47}
48
49impl BlobStore {
50 /// Opens the store at `root`, creating the directory if needed.
51 pub fn open(root: &Path) -> Result<Self> {
52 std::fs::create_dir_all(root)?;
53 let fs = LocalFileSystem::new_with_prefix(root).map_err(storage_err)?;
54 Ok(Self {
55 store: Arc::new(fs),
56 })
57 }
58
59 /// Starts a write of `key`. The bytes are not visible under `key` until `finish`.
60 pub async fn begin(&self, key: &str, max_bytes: u64) -> Result<BlobWriter> {
61 let upload = self
62 .store
63 .put_multipart(&ObjectPath::from(key))
64 .await
65 .map_err(storage_err)?;
66 Ok(BlobWriter {
67 writer: WriteMultipart::new_with_chunk_size(upload, PART_BYTES),
68 hasher: Sha256::new(),
69 size: 0,
70 max_bytes,
71 })
72 }
73
74 /// Opens `key` for reading, or `None` if there is no such blob.
75 pub async fn read(&self, key: &str) -> Result<Option<BlobRead>> {
76 match self.store.get(&ObjectPath::from(key)).await {
77 Ok(result) => {
78 let size = result.meta.size;
79 let stream = result.into_stream().map_err(std::io::Error::other).boxed();
80 Ok(Some(BlobRead { size, stream }))
81 }
82 Err(object_store::Error::NotFound { .. }) => Ok(None),
83 Err(e) => Err(storage_err(e)),
84 }
85 }
86
87 /// Moves `from` to `to`, replacing `to`. Used to give a finished upload its final key.
88 pub async fn rename(&self, from: &str, to: &str) -> Result<()> {
89 self.store
90 .rename(&ObjectPath::from(from), &ObjectPath::from(to))
91 .await
92 .map_err(storage_err)
93 }
94
95 /// Deletes `key`. A blob that is already gone is not an error.
96 pub async fn delete(&self, key: &str) -> Result<()> {
97 match self.store.delete(&ObjectPath::from(key)).await {
98 Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()),
99 Err(e) => Err(storage_err(e)),
100 }
101 }
102}
103
104/// An upload in progress. Dropping it without `finish` leaves no object behind only if `abort`
105/// is called, so every error path in the caller must call `abort`.
106pub struct BlobWriter {
107 writer: WriteMultipart,
108 hasher: Sha256,
109 size: u64,
110 max_bytes: u64,
111}
112
113impl BlobWriter {
114 /// Appends a chunk. Going past the size limit is `Invalid` and writes nothing more.
115 pub async fn push(&mut self, chunk: &[u8]) -> Result<()> {
116 let next = self.size.saturating_add(chunk.len() as u64);
117 if next > self.max_bytes {
118 return Err(Error::Invalid(
119 "the file is larger than the upload limit".into(),
120 ));
121 }
122 self.size = next;
123 self.hasher.update(chunk);
124 self.writer
125 .wait_for_capacity(PARTS_IN_FLIGHT)
126 .await
127 .map_err(storage_err)?;
128 self.writer.write(chunk);
129 Ok(())
130 }
131
132 /// Completes the object. Returns its size and lowercase hex sha256.
133 pub async fn finish(self) -> Result<(u64, String)> {
134 let digest = hex::encode(self.hasher.finalize());
135 self.writer.finish().await.map_err(storage_err)?;
136 Ok((self.size, digest))
137 }
138
139 /// Discards the partial object. Errors are ignored: the object was never visible.
140 pub async fn abort(self) {
141 let _ = self.writer.abort().await;
142 }
143}