From c9ba4672a2c3ec43d6fa2ecf781d2ad0985fac7f Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 10:05:06 -0500 Subject: [PATCH 01/27] Add postgres dependencies --- Cargo.lock | 288 +++++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 5 + 2 files changed, 293 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 5711e02..f856629 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -670,6 +670,28 @@ dependencies = [ "parking_lot_core 0.9.8", ] +[[package]] +name = "deadpool" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "421fe0f90f2ab22016f32a9881be5134fdd71c65298917084b0c7477cbc3856e" +dependencies = [ + "async-trait", + "deadpool-runtime", + "num_cpus", + "retain_mut", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaa37046cc0f6c3cc6090fbdbf73ef0b8ef4cfcc37f6befc0020f63e8cf121e1" +dependencies = [ + "tokio", +] + [[package]] name = "deranged" version = "0.3.8" @@ -692,6 +714,54 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "diesel" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98235fdc2f355d330a8244184ab6b4b33c28679c0b4158f63138e51d6cf7e88" +dependencies = [ + "bitflags 2.4.0", + "byteorder", + "diesel_derives", + "itoa", +] + +[[package]] +name = "diesel-async" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acada1517534c92d3f382217b485db8a8638f111b0e3f2a2a8e26165050f77be" +dependencies = [ + "async-trait", + "deadpool", + "diesel", + "futures-util", + "scoped-futures", + "tokio", + "tokio-postgres", +] + +[[package]] +name = "diesel_derives" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e054665eaf6d97d1e7125512bb2d35d07c73ac86cc6920174cb42d1ab697a554" +dependencies = [ + "diesel_table_macro_syntax", + "proc-macro2", + "quote", + "syn 2.0.29", +] + +[[package]] +name = "diesel_table_macro_syntax" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc5557efc453706fed5e4fa85006fe9817c224c3f480a34c7e5959fd700921c5" +dependencies = [ + "syn 2.0.29", +] + [[package]] name = "digest" version = "0.10.7" @@ -740,6 +810,12 @@ dependencies = [ "once_cell", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "flate2" version = "1.0.27" @@ -1684,6 +1760,24 @@ dependencies = [ "sha2", ] +[[package]] +name = "phf" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_shared" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" +dependencies = [ + "siphasher", +] + [[package]] name = "pict-rs" version = "0.5.0-alpha.17" @@ -1700,6 +1794,9 @@ dependencies = [ "config", "console-subscriber", "dashmap", + "deadpool", + "diesel", + "diesel-async", "flume", "futures-core", "hex", @@ -1713,6 +1810,7 @@ dependencies = [ "opentelemetry-otlp", "pin-project-lite", "quick-xml 0.30.0", + "refinery", "reqwest", "reqwest-middleware", "reqwest-tracing", @@ -1728,6 +1826,7 @@ dependencies = [ "thiserror", "time", "tokio", + "tokio-postgres", "tokio-uring", "tokio-util", "toml 0.7.6", @@ -1780,6 +1879,53 @@ version = "1.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31114a898e107c51bb1609ffaf55a0e011cf6a4d7f1170d0015a165082c0338b" +[[package]] +name = "postgres" +version = "0.19.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7915b33ed60abc46040cbcaa25ffa1c7ec240668e0477c4f3070786f5916d451" +dependencies = [ + "bytes", + "fallible-iterator", + "futures-util", + "log", + "tokio", + "tokio-postgres", +] + +[[package]] +name = "postgres-protocol" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49b6c5ef183cd3ab4ba005f1ca64c21e8bd97ce4699cfea9e8d9a2c4958ca520" +dependencies = [ + "base64 0.21.3", + "byteorder", + "bytes", + "fallible-iterator", + "hmac", + "md-5", + "memchr", + "rand", + "sha2", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d2234cdee9408b523530a9b6d2d6b373d1db34f6a8e51dc03ded1828d7fb67c" +dependencies = [ + "bytes", + "fallible-iterator", + "postgres-protocol", + "serde", + "serde_json", + "time", + "uuid", +] + [[package]] name = "ppv-lite86" version = "0.2.17" @@ -1929,6 +2075,52 @@ dependencies = [ "bitflags 1.3.2", ] +[[package]] +name = "refinery" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdb0436d0dd7bd8d4fce1e828751fa79742b08e35f27cfea7546f8a322b5ef24" +dependencies = [ + "refinery-core", + "refinery-macros", +] + +[[package]] +name = "refinery-core" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19206547cd047e8f4dfa6b20c30d3ecaf24be05841b6aa0aa926a47a3d0662bb" +dependencies = [ + "async-trait", + "cfg-if", + "lazy_static", + "log", + "postgres", + "regex", + "serde", + "siphasher", + "thiserror", + "time", + "tokio", + "tokio-postgres", + "toml 0.7.6", + "url", + "walkdir", +] + +[[package]] +name = "refinery-macros" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d94d4b9241859ba19eaa5c04c86e782eb3aa0aae2c5868e0cfa90c856e58a174" +dependencies = [ + "proc-macro2", + "quote", + "refinery-core", + "regex", + "syn 2.0.29", +] + [[package]] name = "regex" version = "1.9.4" @@ -2046,6 +2238,12 @@ dependencies = [ "tracing", ] +[[package]] +name = "retain_mut" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4389f1d5789befaf6029ebd9f7dac4af7f7e3d61b69d4f30e2ac02b57e7712b0" + [[package]] name = "ring" version = "0.16.20" @@ -2168,6 +2366,25 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "scoped-futures" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1473e24c637950c9bd38763220bea91ec3e095a89f672bbd7a10d03e77ba467" +dependencies = [ + "cfg-if", + "pin-utils", +] + [[package]] name = "scoped-tls" version = "1.0.1" @@ -2316,6 +2533,12 @@ dependencies = [ "libc", ] +[[package]] +name = "siphasher" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" + [[package]] name = "sketches-ddsketch" version = "0.2.1" @@ -2395,6 +2618,16 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f11d35dae9818c4313649da4a97c8329e29357a7fe584526c1d78f5b63ef836" +[[package]] +name = "stringprep" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db3737bde7edce97102e0e2b15365bf7a20bfdb5f60f4f9e8d7004258a51a8da" +dependencies = [ + "unicode-bidi", + "unicode-normalization", +] + [[package]] name = "strsim" version = "0.10.0" @@ -2558,6 +2791,32 @@ dependencies = [ "syn 2.0.29", ] +[[package]] +name = "tokio-postgres" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d340244b32d920260ae7448cb72b6e238bddc3d4f7603394e7dd46ed8e48f5b8" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot 0.12.1", + "percent-encoding", + "phf", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand", + "socket2 0.5.3", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.24.1" @@ -2936,6 +3195,16 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" +[[package]] +name = "walkdir" +version = "2.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36df944cda56c7d8d8b7496af378e6b16de9284591917d307c9b4d313c44e698" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -3046,6 +3315,16 @@ version = "0.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "14247bb57be4f377dfb94c72830b8ce8fc6beac03cf4bf7b9732eadd414123fc" +[[package]] +name = "whoami" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22fc3756b8a9133049b26c7f61ab35416c130e8c09b660f5b3958b446f52cc50" +dependencies = [ + "wasm-bindgen", + "web-sys", +] + [[package]] name = "winapi" version = "0.3.9" @@ -3062,6 +3341,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +dependencies = [ + "winapi", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" diff --git a/Cargo.toml b/Cargo.toml index ea1e0d9..9146a07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,9 @@ color-eyre = "0.6" config = "0.13.0" console-subscriber = "0.1" dashmap = "5.1.0" +deadpool = { version = "0.9.5", features = ["rt_tokio_1"] } +diesel = "2.1.1" +diesel-async = { version = "0.4.1", features = ["postgres", "deadpool"] } flume = "0.11.0" futures-core = "0.3" hex = "0.4.3" @@ -39,6 +42,7 @@ opentelemetry = { version = "0.20", features = ["rt-tokio"] } opentelemetry-otlp = "0.13" pin-project-lite = "0.2.7" quick-xml = { version = "0.30.0", features = ["serialize"] } +refinery = { version = "0.8.10", features = ["tokio-postgres", "postgres"] } reqwest = { version = "0.11.18", default-features = false, features = ["json", "rustls-tls", "stream"] } reqwest-middleware = "0.2.2" reqwest-tracing = { version = "0.4.5" } @@ -54,6 +58,7 @@ storage-path-generator = "0.1.0" thiserror = "1.0" time = { version = "0.3.0", features = ["serde", "serde-well-known"] } tokio = { version = "1", features = ["full", "tracing"] } +tokio-postgres = { version = "0.7.10", features = ["with-uuid-1", "with-time-0_3", "with-serde_json-1"] } tokio-uring = { version = "0.4", optional = true, features = ["bytes"] } tokio-util = { version = "0.7", default-features = false, features = [ "codec", From 8c532c97e691215b87a654e8245a1e30208ec568 Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 11:52:55 -0500 Subject: [PATCH 02/27] Initial postgres work --- Cargo.lock | 7 ++ Cargo.toml | 1 + flake.nix | 1 + src/config.rs | 4 +- src/config/commandline.rs | 91 ++++++++++++++++++- src/config/defaults.rs | 12 +++ src/config/file.rs | 7 ++ src/lib.rs | 34 +++---- src/repo.rs | 10 +- src/repo/postgres.rs | 81 +++++++++++++++++ .../migrations/V0001__create_hashes.rs | 31 +++++++ src/store/file_store.rs | 40 +++----- src/store/object_store.rs | 38 +++----- 13 files changed, 282 insertions(+), 75 deletions(-) create mode 100644 src/repo/postgres.rs create mode 100644 src/repo/postgres/migrations/V0001__create_hashes.rs diff --git a/Cargo.lock b/Cargo.lock index f856629..c7e9463 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -388,6 +388,12 @@ dependencies = [ "rustc-demangle", ] +[[package]] +name = "barrel" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad9e605929a6964efbec5ac0884bd0fe93f12a3b1eb271f52c251316640c68d9" + [[package]] name = "base64" version = "0.13.1" @@ -1788,6 +1794,7 @@ dependencies = [ "actix-web", "anyhow", "async-trait", + "barrel", "base64 0.21.3", "clap", "color-eyre", diff --git a/Cargo.toml b/Cargo.toml index 9146a07..7f6e6db 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ actix-server = "2.0.0" actix-web = { version = "4.0.0", default-features = false } anyhow = "1.0" async-trait = "0.1.51" +barrel = { version = "0.7.0", features = ["pg"] } base64 = "0.21.0" clap = { version = "4.0.2", features = ["derive"] } color-eyre = "0.6" diff --git a/flake.nix b/flake.nix index 8538a2d..c0c7e2a 100644 --- a/flake.nix +++ b/flake.nix @@ -32,6 +32,7 @@ cargo cargo-outdated clippy + diesel-cli exiftool ffmpeg_6-full garage diff --git a/src/config.rs b/src/config.rs index 23ed003..25cf0e7 100644 --- a/src/config.rs +++ b/src/config.rs @@ -12,8 +12,8 @@ use defaults::Defaults; pub(crate) use commandline::Operation; pub(crate) use file::{ - Animation, ConfigFile as Configuration, Image, Media, ObjectStorage, OpenTelemetry, Repo, Sled, - Store, Tracing, Video, + Animation, ConfigFile as Configuration, Image, Media, ObjectStorage, OpenTelemetry, Postgres, + Repo, Sled, Store, Tracing, Video, }; pub(crate) use primitives::{Filesystem, LogFormat}; diff --git a/src/config/commandline.rs b/src/config/commandline.rs index d08a52c..c7ffad5 100644 --- a/src/config/commandline.rs +++ b/src/config/commandline.rs @@ -369,8 +369,64 @@ impl Args { from: from.into(), to: to.into(), }, - config_file, save_to, + config_file, + }, + MigrateRepoTo::Postgres(MigratePostgresInner { to }) => Output { + config_format: ConfigFormat { + server, + client, + old_repo, + tracing, + metrics, + media, + repo: None, + store: None, + }, + operation: Operation::MigrateRepo { + from: from.into(), + to: to.into(), + }, + save_to, + config_file, + }, + }, + MigrateRepoFrom::Postgres(MigratePostgresRepo { from, to }) => match to { + MigrateRepoTo::Sled(MigrateSledInner { to }) => Output { + config_format: ConfigFormat { + server, + client, + old_repo, + tracing, + metrics, + media, + repo: None, + store: None, + }, + operation: Operation::MigrateRepo { + from: from.into(), + to: to.into(), + }, + save_to, + config_file, + }, + MigrateRepoTo::Postgres(MigratePostgresInner { to }) => Output { + config_format: ConfigFormat { + server, + client, + old_repo, + tracing, + metrics, + media, + repo: None, + store: None, + }, + operation: Operation::MigrateRepo { + from: from.into(), + to: to.into(), + }, + save_to, + config_file, }, }, } @@ -1058,6 +1114,7 @@ enum MigrateStoreFrom { #[derive(Debug, Subcommand)] enum MigrateRepoFrom { Sled(MigrateSledRepo), + Postgres(MigratePostgresRepo), } /// Configure the destination storage for pict-rs storage migration @@ -1075,8 +1132,10 @@ enum MigrateStoreTo { /// Configure the destination repo for pict-rs repo migration #[derive(Debug, Subcommand)] enum MigrateRepoTo { - /// Migrate to the provided sled storage + /// Migrate to the provided sled repo Sled(MigrateSledInner), + /// Migrate to the provided postgres repo + Postgres(MigratePostgresInner), } /// Migrate pict-rs' storage from the provided filesystem storage @@ -1099,6 +1158,16 @@ struct MigrateSledRepo { to: MigrateRepoTo, } +/// Migrate pict-rs' repo from the provided postgres repo +#[derive(Debug, Parser)] +struct MigratePostgresRepo { + #[command(flatten)] + from: Postgres, + + #[command(subcommand)] + to: MigrateRepoTo, +} + /// Migrate pict-rs' storage to the provided filesystem storage #[derive(Debug, Parser)] struct MigrateFilesystemInner { @@ -1116,6 +1185,13 @@ struct MigrateSledInner { to: Sled, } +/// Migrate pict-rs' repo to the provided postgres repo +#[derive(Debug, Parser)] +struct MigratePostgresInner { + #[command(flatten)] + to: Postgres, +} + /// Migrate pict-rs' storage from the provided object storage #[derive(Debug, Parser)] struct MigrateObjectStorage { @@ -1163,6 +1239,8 @@ struct RunObjectStorage { enum Repo { /// Run pict-rs with the provided sled-backed data repository Sled(Sled), + /// Run pict-rs with the provided postgres-backed data repository + Postgres(Postgres), } /// Configuration for filesystem media storage @@ -1254,6 +1332,15 @@ pub(super) struct Sled { pub(super) export_path: Option, } +/// Configuration for the postgres-backed data repository +#[derive(Debug, Parser, serde::Serialize)] +#[serde(rename_all = "snake_case")] +pub(super) struct Postgres { + /// The URL of the postgres database + #[arg(short, long)] + pub(super) url: Url, +} + #[derive(Debug, Parser, serde::Serialize)] #[serde(rename_all = "snake_case")] struct OldSled { diff --git a/src/config/defaults.rs b/src/config/defaults.rs index fc4a0d2..935d5c9 100644 --- a/src/config/defaults.rs +++ b/src/config/defaults.rs @@ -363,8 +363,20 @@ impl From for crate::config::file::Sled { } } +impl From for crate::config::file::Postgres { + fn from(value: crate::config::commandline::Postgres) -> Self { + crate::config::file::Postgres { url: value.url } + } +} + impl From for crate::config::file::Repo { fn from(value: crate::config::commandline::Sled) -> Self { crate::config::file::Repo::Sled(value.into()) } } + +impl From for crate::config::file::Repo { + fn from(value: crate::config::commandline::Postgres) -> Self { + crate::config::file::Repo::Postgres(value.into()) + } +} diff --git a/src/config/file.rs b/src/config/file.rs index 9e78d19..213c715 100644 --- a/src/config/file.rs +++ b/src/config/file.rs @@ -88,6 +88,7 @@ pub(crate) struct ObjectStorage { #[serde(tag = "type")] pub(crate) enum Repo { Sled(Sled), + Postgres(Postgres), } #[derive(Clone, Debug, serde::Deserialize, serde::Serialize)] @@ -421,3 +422,9 @@ pub(crate) struct Sled { pub(crate) export_path: PathBuf, } + +#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) struct Postgres { + pub(crate) url: Url, +} diff --git a/src/lib.rs b/src/lib.rs index b7faff8..21952fc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1810,7 +1810,7 @@ async fn launch_object_store( - repo: Repo, + repo: ArcRepo, client: ClientWithMiddleware, from: S1, to: config::primitives::Store, @@ -1824,11 +1824,7 @@ where config::primitives::Store::Filesystem(config::Filesystem { path }) => { let to = FileStore::build(path.clone(), repo.clone()).await?; - match repo { - Repo::Sled(repo) => { - migrate_store(Arc::new(repo), from, to, skip_missing_files, timeout).await? - } - } + migrate_store(repo, from, to, skip_missing_files, timeout).await? } config::primitives::Store::ObjectStorage(config::primitives::ObjectStorage { endpoint, @@ -1862,11 +1858,7 @@ where .await? .build(client); - match repo { - Repo::Sled(repo) => { - migrate_store(Arc::new(repo), from, to, skip_missing_files, timeout).await? - } - } + migrate_store(repo, from, to, skip_missing_files, timeout).await? } } @@ -1970,7 +1962,7 @@ impl PictRsConfiguration { from, to, } => { - let repo = Repo::open(config.repo.clone())?; + let repo = Repo::open(config.repo.clone()).await?.to_arc(); match from { config::primitives::Store::Filesystem(config::Filesystem { path }) => { @@ -2034,15 +2026,15 @@ impl PictRsConfiguration { return Ok(()); } Operation::MigrateRepo { from, to } => { - let from = Repo::open(from)?.to_arc(); - let to = Repo::open(to)?.to_arc(); + let from = Repo::open(from).await?.to_arc(); + let to = Repo::open(to).await?.to_arc(); repo::migrate_repo(from, to).await?; return Ok(()); } } - let repo = Repo::open(config.repo.clone())?; + let repo = Repo::open(config.repo.clone()).await?; if config.server.read_only { tracing::warn!("Launching in READ ONLY mode"); @@ -2050,10 +2042,10 @@ impl PictRsConfiguration { match config.store.clone() { config::Store::Filesystem(config::Filesystem { path }) => { - let store = FileStore::build(path, repo.clone()).await?; - let arc_repo = repo.to_arc(); + let store = FileStore::build(path, arc_repo.clone()).await?; + if arc_repo.get("migrate-0.4").await?.is_none() { if let Some(old_repo) = repo_04::open(&config.old_repo)? { repo::migrate_04(old_repo, arc_repo.clone(), store.clone(), config.clone()) @@ -2075,6 +2067,7 @@ impl PictRsConfiguration { ) .await?; } + Repo::Postgres(_) => todo!(), } } config::Store::ObjectStorage(config::ObjectStorage { @@ -2089,6 +2082,8 @@ impl PictRsConfiguration { client_timeout, public_endpoint, }) => { + let arc_repo = repo.to_arc(); + let store = ObjectStore::build( endpoint, bucket_name, @@ -2104,13 +2099,11 @@ impl PictRsConfiguration { signature_duration, client_timeout, public_endpoint, - repo.clone(), + arc_repo.clone(), ) .await? .build(client.clone()); - let arc_repo = repo.to_arc(); - if arc_repo.get("migrate-0.4").await?.is_none() { if let Some(old_repo) = repo_04::open(&config.old_repo)? { repo::migrate_04(old_repo, arc_repo.clone(), store.clone(), config.clone()) @@ -2128,6 +2121,7 @@ impl PictRsConfiguration { }) .await?; } + Repo::Postgres(_) => todo!(), } } } diff --git a/src/repo.rs b/src/repo.rs index 2518cb3..51553f4 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -12,6 +12,7 @@ use uuid::Uuid; mod hash; mod migrate; +pub(crate) mod postgres; pub(crate) mod sled; pub(crate) use hash::Hash; @@ -22,6 +23,7 @@ pub(crate) type ArcRepo = Arc; #[derive(Clone, Debug)] pub(crate) enum Repo { Sled(self::sled::SledRepo), + Postgres(self::postgres::PostgresRepo), } #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -791,7 +793,7 @@ where impl Repo { #[tracing::instrument] - pub(crate) fn open(config: config::Repo) -> color_eyre::Result { + pub(crate) async fn open(config: config::Repo) -> color_eyre::Result { match config { config::Repo::Sled(config::Sled { path, @@ -802,12 +804,18 @@ impl Repo { Ok(Self::Sled(repo)) } + config::Repo::Postgres(config::Postgres { url }) => { + let repo = self::postgres::PostgresRepo::connect(url).await?; + + Ok(Self::Postgres(repo)) + } } } pub(crate) fn to_arc(&self) -> ArcRepo { match self { Self::Sled(sled_repo) => Arc::new(sled_repo.clone()), + Self::Postgres(_) => todo!(), } } } diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs new file mode 100644 index 0000000..7c3b031 --- /dev/null +++ b/src/repo/postgres.rs @@ -0,0 +1,81 @@ +mod embedded { + use refinery::embed_migrations; + + embed_migrations!("./src/repo/postgres/migrations"); +} + +use diesel_async::{ + pooled_connection::{ + deadpool::{BuildError, Pool}, + AsyncDieselConnectionManager, + }, + AsyncPgConnection, +}; +use url::Url; + +use super::{BaseRepo, HashRepo}; + +#[derive(Clone)] +pub(crate) struct PostgresRepo { + pool: Pool, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum ConnectPostgresError { + #[error("Failed to connect to postgres for migrations")] + ConnectForMigration(#[source] tokio_postgres::Error), + + #[error("Failed to run migrations")] + Migration(#[source] refinery::Error), + + #[error("Failed to build postgres connection pool")] + BuildPool(#[source] BuildError), +} + +#[derive(Debug, thiserror::Error)] +enum PostgresError {} + +impl PostgresRepo { + pub(crate) async fn connect(postgres_url: Url) -> Result { + let (mut client, conn) = + tokio_postgres::connect(postgres_url.as_str(), tokio_postgres::tls::NoTls) + .await + .map_err(ConnectPostgresError::ConnectForMigration)?; + + let handle = actix_rt::spawn(conn); + + embedded::migrations::runner() + .run_async(&mut client) + .await + .map_err(ConnectPostgresError::Migration)?; + + handle.abort(); + let _ = handle.await; + + let config = AsyncDieselConnectionManager::::new(postgres_url); + let pool = Pool::builder(config) + .build() + .map_err(ConnectPostgresError::BuildPool)?; + + Ok(PostgresRepo { pool }) + } +} + +impl BaseRepo for PostgresRepo {} + +/* +#[async_trait::async_trait] +impl HashRepo for PostgresRepo { + async fn size(&self) -> Result { + let conn = self.pool.get().await.map_err(PostgresError::from)?; + } +} +*/ + +impl std::fmt::Debug for PostgresRepo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PostgresRepo") + .field("pool", &"pool") + .finish() + } +} diff --git a/src/repo/postgres/migrations/V0001__create_hashes.rs b/src/repo/postgres/migrations/V0001__create_hashes.rs new file mode 100644 index 0000000..30455c3 --- /dev/null +++ b/src/repo/postgres/migrations/V0001__create_hashes.rs @@ -0,0 +1,31 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("hashes", |t| { + t.add_column( + "hash", + types::binary() + .primary(true) + .unique(true) + .nullable(false) + .size(128), + ); + t.add_column("identifier", types::text().unique(true).nullable(false)); + t.add_column( + "motion_identifier", + types::text().unique(true).nullable(true), + ); + t.add_column( + "created_at", + types::datetime() + .nullable(false) + .default(AutogenFunction::CurrentTimestamp), + ); + }); + + m.make::().to_string() +} diff --git a/src/store/file_store.rs b/src/store/file_store.rs index b2a1cb7..b1738cb 100644 --- a/src/store/file_store.rs +++ b/src/store/file_store.rs @@ -1,9 +1,4 @@ -use crate::{ - error_code::ErrorCode, - file::File, - repo::{Repo, SettingsRepo}, - store::Store, -}; +use crate::{error_code::ErrorCode, file::File, repo::ArcRepo, store::Store}; use actix_web::web::Bytes; use futures_core::Stream; use std::{ @@ -58,7 +53,7 @@ impl FileError { pub(crate) struct FileStore { path_gen: Generator, root_dir: PathBuf, - repo: Repo, + repo: ArcRepo, } #[async_trait::async_trait(?Send)] @@ -189,7 +184,7 @@ impl Store for FileStore { impl FileStore { #[tracing::instrument(skip(repo))] - pub(crate) async fn build(root_dir: PathBuf, repo: Repo) -> color_eyre::Result { + pub(crate) async fn build(root_dir: PathBuf, repo: ArcRepo) -> color_eyre::Result { let path_gen = init_generator(&repo).await?; tokio::fs::create_dir_all(&root_dir).await?; @@ -204,13 +199,9 @@ impl FileStore { async fn next_directory(&self) -> Result { let path = self.path_gen.next(); - match self.repo { - Repo::Sled(ref sled_repo) => { - sled_repo - .set(GENERATOR_KEY, path.to_be_bytes().into()) - .await?; - } - } + self.repo + .set(GENERATOR_KEY, path.to_be_bytes().into()) + .await?; let mut target_path = self.root_dir.clone(); for dir in path.to_strings() { @@ -308,18 +299,13 @@ pub(crate) async fn safe_create_parent>(path: P) -> Result<(), Fi Ok(()) } -async fn init_generator(repo: &Repo) -> Result { - match repo { - Repo::Sled(sled_repo) => { - if let Some(ivec) = sled_repo.get(GENERATOR_KEY).await? { - Ok(Generator::from_existing( - storage_path_generator::Path::from_be_bytes(ivec.to_vec()) - .map_err(FileError::from)?, - )) - } else { - Ok(Generator::new()) - } - } +async fn init_generator(repo: &ArcRepo) -> Result { + if let Some(ivec) = repo.get(GENERATOR_KEY).await? { + Ok(Generator::from_existing( + storage_path_generator::Path::from_be_bytes(ivec.to_vec()).map_err(FileError::from)?, + )) + } else { + Ok(Generator::new()) } } diff --git a/src/store/object_store.rs b/src/store/object_store.rs index a37f465..301c181 100644 --- a/src/store/object_store.rs +++ b/src/store/object_store.rs @@ -1,7 +1,7 @@ use crate::{ bytes_stream::BytesStream, error_code::ErrorCode, - repo::{Repo, SettingsRepo}, + repo::ArcRepo, store::Store, stream::{IntoStreamer, StreamMap}, }; @@ -107,7 +107,7 @@ impl From for ObjectError { #[derive(Clone)] pub(crate) struct ObjectStore { path_gen: Generator, - repo: Repo, + repo: ArcRepo, bucket: Bucket, credentials: Credentials, client: ClientWithMiddleware, @@ -119,7 +119,7 @@ pub(crate) struct ObjectStore { #[derive(Clone)] pub(crate) struct ObjectStoreConfig { path_gen: Generator, - repo: Repo, + repo: ArcRepo, bucket: Bucket, credentials: Credentials, signature_expiration: u64, @@ -493,7 +493,7 @@ impl ObjectStore { signature_expiration: u64, client_timeout: u64, public_endpoint: Option, - repo: Repo, + repo: ArcRepo, ) -> Result { let path_gen = init_generator(&repo).await?; @@ -714,13 +714,9 @@ impl ObjectStore { async fn next_directory(&self) -> Result { let path = self.path_gen.next(); - match self.repo { - Repo::Sled(ref sled_repo) => { - sled_repo - .set(GENERATOR_KEY, path.to_be_bytes().into()) - .await?; - } - } + self.repo + .set(GENERATOR_KEY, path.to_be_bytes().into()) + .await?; Ok(path) } @@ -733,18 +729,14 @@ impl ObjectStore { } } -async fn init_generator(repo: &Repo) -> Result { - match repo { - Repo::Sled(sled_repo) => { - if let Some(ivec) = sled_repo.get(GENERATOR_KEY).await? { - Ok(Generator::from_existing( - storage_path_generator::Path::from_be_bytes(ivec.to_vec()) - .map_err(ObjectError::from)?, - )) - } else { - Ok(Generator::new()) - } - } +async fn init_generator(repo: &ArcRepo) -> Result { + if let Some(ivec) = repo.get(GENERATOR_KEY).await? { + Ok(Generator::from_existing( + storage_path_generator::Path::from_be_bytes(ivec.to_vec()) + .map_err(ObjectError::from)?, + )) + } else { + Ok(Generator::new()) } } From d4757230876ea6e81d83fd4575e2f9d51ddf7c6c Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 13:35:30 -0500 Subject: [PATCH 03/27] Finish writing migrations & generate diesel schema --- src/repo/postgres.rs | 1 + .../migrations/V0000__enable_pgcrypto.rs | 11 ++ .../migrations/V0001__create_hashes.rs | 2 + .../migrations/V0002__create_variants.rs | 30 +++++ .../migrations/V0003__create_aliases.rs | 25 ++++ .../migrations/V0004__create_settings.rs | 21 ++++ .../migrations/V0005__create_details.rs | 17 +++ .../migrations/V0006__create_queue.rs | 53 ++++++++ .../V0007__create_store_migrations.rs | 16 +++ .../migrations/V0008__create_proxies.rs | 25 ++++ .../migrations/V0009__create_uploads.rs | 38 ++++++ src/repo/postgres/schema.rs | 114 ++++++++++++++++++ 12 files changed, 353 insertions(+) create mode 100644 src/repo/postgres/migrations/V0000__enable_pgcrypto.rs create mode 100644 src/repo/postgres/migrations/V0002__create_variants.rs create mode 100644 src/repo/postgres/migrations/V0003__create_aliases.rs create mode 100644 src/repo/postgres/migrations/V0004__create_settings.rs create mode 100644 src/repo/postgres/migrations/V0005__create_details.rs create mode 100644 src/repo/postgres/migrations/V0006__create_queue.rs create mode 100644 src/repo/postgres/migrations/V0007__create_store_migrations.rs create mode 100644 src/repo/postgres/migrations/V0008__create_proxies.rs create mode 100644 src/repo/postgres/migrations/V0009__create_uploads.rs create mode 100644 src/repo/postgres/schema.rs diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 7c3b031..f4e71da 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -3,6 +3,7 @@ mod embedded { embed_migrations!("./src/repo/postgres/migrations"); } +mod schema; use diesel_async::{ pooled_connection::{ diff --git a/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs b/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs new file mode 100644 index 0000000..44eb85c --- /dev/null +++ b/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs @@ -0,0 +1,11 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.inject_custom("CREATE EXTENSION pgcrypto;"); + + m.make::().to_string() +} diff --git a/src/repo/postgres/migrations/V0001__create_hashes.rs b/src/repo/postgres/migrations/V0001__create_hashes.rs index 30455c3..9829585 100644 --- a/src/repo/postgres/migrations/V0001__create_hashes.rs +++ b/src/repo/postgres/migrations/V0001__create_hashes.rs @@ -25,6 +25,8 @@ pub(crate) fn migration() -> String { .nullable(false) .default(AutogenFunction::CurrentTimestamp), ); + + t.add_index("ordered_hash_index", types::index(["created_at", "hash"])); }); m.make::().to_string() diff --git a/src/repo/postgres/migrations/V0002__create_variants.rs b/src/repo/postgres/migrations/V0002__create_variants.rs new file mode 100644 index 0000000..8377ddc --- /dev/null +++ b/src/repo/postgres/migrations/V0002__create_variants.rs @@ -0,0 +1,30 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("variants", |t| { + t.inject_custom(r#""id" UUID PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL UNIQUE"#); + t.add_column("hash", types::binary().nullable(false)); + t.add_column("variant", types::text().nullable(false)); + t.add_column("identifier", types::text().nullable(false)); + t.add_column( + "accessed", + types::datetime() + .nullable(false) + .default(AutogenFunction::CurrentTimestamp), + ); + + t.add_foreign_key(&["hash"], "hashes", &["hash"]); + t.add_index( + "hash_variant_index", + types::index(["hash", "variant"]).unique(true), + ); + }); + + let s = m.make::().to_string(); + println!("{s}"); + s +} diff --git a/src/repo/postgres/migrations/V0003__create_aliases.rs b/src/repo/postgres/migrations/V0003__create_aliases.rs new file mode 100644 index 0000000..a8aa37f --- /dev/null +++ b/src/repo/postgres/migrations/V0003__create_aliases.rs @@ -0,0 +1,25 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("aliases", |t| { + t.add_column( + "alias", + types::text() + .size(60) + .primary(true) + .unique(true) + .nullable(false), + ); + t.add_column("hash", types::binary().nullable(false)); + t.add_column("token", types::text().size(60).nullable(false)); + + t.add_foreign_key(&["hash"], "hashes", &["hash"]); + t.add_index("aliases_hash_index", types::index(["hash"])); + }); + + m.make::().to_string() +} diff --git a/src/repo/postgres/migrations/V0004__create_settings.rs b/src/repo/postgres/migrations/V0004__create_settings.rs new file mode 100644 index 0000000..37011ac --- /dev/null +++ b/src/repo/postgres/migrations/V0004__create_settings.rs @@ -0,0 +1,21 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("settings", |t| { + t.add_column( + "key", + types::text() + .size(80) + .primary(true) + .unique(true) + .nullable(false), + ); + t.add_column("value", types::text().size(80).nullable(false)); + }); + + m.make::().to_string() +} diff --git a/src/repo/postgres/migrations/V0005__create_details.rs b/src/repo/postgres/migrations/V0005__create_details.rs new file mode 100644 index 0000000..e324caa --- /dev/null +++ b/src/repo/postgres/migrations/V0005__create_details.rs @@ -0,0 +1,17 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("details", |t| { + t.add_column( + "identifier", + types::text().primary(true).unique(true).nullable(false), + ); + t.add_column("details", types::custom("jsonb").nullable(false)); + }); + + m.make::().to_string() +} diff --git a/src/repo/postgres/migrations/V0006__create_queue.rs b/src/repo/postgres/migrations/V0006__create_queue.rs new file mode 100644 index 0000000..293ff97 --- /dev/null +++ b/src/repo/postgres/migrations/V0006__create_queue.rs @@ -0,0 +1,53 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.inject_custom("CREATE TYPE job_status AS ENUM ('new', 'running');"); + + m.create_table("queue", |t| { + t.inject_custom(r#""id" UUID PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL UNIQUE"#); + t.add_column("queue", types::text().size(50).nullable(false)); + t.add_column("job", types::custom("jsonb").nullable(false)); + t.add_column("status", types::custom("job_status").nullable(false)); + t.add_column( + "queue_time", + types::datetime() + .nullable(false) + .default(AutogenFunction::CurrentTimestamp), + ); + t.add_column("heartbeat", types::datetime()); + + t.add_index("queue_status_index", types::index(["queue", "status"])); + t.add_index("heartbeat_index", types::index(["heartbeat"])); + }); + + m.inject_custom( + r#" +CREATE OR REPLACE FUNCTION queue_status_notify() + RETURNS trigger AS +$$ +BEGIN + PERFORM pg_notify('queue_status_channel', NEW.id::text); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + "# + .trim(), + ); + + m.inject_custom( + r#" +CREATE TRIGGER queue_status + AFTER INSERT OR UPDATE OF status + ON queue + FOR EACH ROW +EXECUTE PROCEDURE queue_status_notify(); + "# + .trim(), + ); + + m.make::().to_string() +} diff --git a/src/repo/postgres/migrations/V0007__create_store_migrations.rs b/src/repo/postgres/migrations/V0007__create_store_migrations.rs new file mode 100644 index 0000000..fe957b0 --- /dev/null +++ b/src/repo/postgres/migrations/V0007__create_store_migrations.rs @@ -0,0 +1,16 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("store_migrations", |t| { + t.add_column( + "identifier", + types::text().primary(true).nullable(false).unique(true), + ); + }); + + m.make::().to_string() +} diff --git a/src/repo/postgres/migrations/V0008__create_proxies.rs b/src/repo/postgres/migrations/V0008__create_proxies.rs new file mode 100644 index 0000000..bbb6c5d --- /dev/null +++ b/src/repo/postgres/migrations/V0008__create_proxies.rs @@ -0,0 +1,25 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("proxies", |t| { + t.add_column( + "url", + types::text().primary(true).unique(true).nullable(false), + ); + t.add_column("alias", types::text().nullable(false)); + t.add_column( + "accessed", + types::datetime() + .nullable(false) + .default(AutogenFunction::CurrentTimestamp), + ); + + t.add_foreign_key(&["alias"], "aliases", &["alias"]); + }); + + m.make::().to_string() +} diff --git a/src/repo/postgres/migrations/V0009__create_uploads.rs b/src/repo/postgres/migrations/V0009__create_uploads.rs new file mode 100644 index 0000000..9214ec4 --- /dev/null +++ b/src/repo/postgres/migrations/V0009__create_uploads.rs @@ -0,0 +1,38 @@ +use barrel::backend::Pg; +use barrel::functions::AutogenFunction; +use barrel::{types, Migration}; + +pub(crate) fn migration() -> String { + let mut m = Migration::new(); + + m.create_table("uploads", |t| { + t.inject_custom(r#""id" UUID PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL UNIQUE"#); + t.add_column("result", types::custom("jsonb")); + }); + + m.inject_custom( + r#" +CREATE OR REPLACE FUNCTION upload_completion_notify() + RETURNS trigger AS +$$ +BEGIN + PERFORM pg_notify('upload_completion_channel', NEW.id::text); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + "# + .trim(), + ); + + m.inject_custom( + r#" +CREATE TRIGGER upload_result + AFTER INSERT OR UPDATE OF result + ON uploads + FOR EACH ROW +EXECUTE PROCEDURE upload_completion_notify(); + "#, + ); + + m.make::().to_string() +} diff --git a/src/repo/postgres/schema.rs b/src/repo/postgres/schema.rs new file mode 100644 index 0000000..3ab3e46 --- /dev/null +++ b/src/repo/postgres/schema.rs @@ -0,0 +1,114 @@ +// @generated automatically by Diesel CLI. + +pub mod sql_types { + #[derive(diesel::sql_types::SqlType)] + #[diesel(postgres_type(name = "job_status"))] + pub struct JobStatus; +} + +diesel::table! { + aliases (alias) { + alias -> Text, + hash -> Bytea, + token -> Text, + } +} + +diesel::table! { + details (identifier) { + identifier -> Text, + #[sql_name = "details"] + details_json -> Jsonb, + } +} + +diesel::table! { + hashes (hash) { + hash -> Bytea, + identifier -> Text, + motion_identifier -> Nullable, + created_at -> Timestamp, + } +} + +diesel::table! { + proxies (url) { + url -> Text, + alias -> Text, + accessed -> Timestamp, + } +} + +diesel::table! { + use diesel::sql_types::*; + use super::sql_types::JobStatus; + + queue (id) { + id -> Uuid, + #[sql_name = "queue"] + queue_name -> Text, + job -> Jsonb, + status -> JobStatus, + queue_time -> Timestamp, + heartbeat -> Timestamp, + } +} + +diesel::table! { + refinery_schema_history (version) { + version -> Int4, + #[max_length = 255] + name -> Nullable, + #[max_length = 255] + applied_on -> Nullable, + #[max_length = 255] + checksum -> Nullable, + } +} + +diesel::table! { + settings (key) { + key -> Text, + value -> Text, + } +} + +diesel::table! { + store_migrations (identifier) { + identifier -> Text, + } +} + +diesel::table! { + uploads (id) { + id -> Uuid, + result -> Jsonb, + } +} + +diesel::table! { + variants (id) { + id -> Uuid, + hash -> Bytea, + variant -> Text, + identifier -> Text, + accessed -> Timestamp, + } +} + +diesel::joinable!(aliases -> hashes (hash)); +diesel::joinable!(proxies -> aliases (alias)); +diesel::joinable!(variants -> hashes (hash)); + +diesel::allow_tables_to_appear_in_same_query!( + aliases, + details, + hashes, + proxies, + queue, + refinery_schema_history, + settings, + store_migrations, + uploads, + variants, +); From 5167fe331e9e8fd1f153533bccb9747421dfcbe5 Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 13:36:03 -0500 Subject: [PATCH 04/27] Update docs, dev.toml --- dev.toml | 6 ++---- docs/postgres-planning.md | 4 ++-- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/dev.toml b/dev.toml index 209af82..ae32642 100644 --- a/dev.toml +++ b/dev.toml @@ -59,10 +59,8 @@ crf_2160 = 15 crf_max = 12 [repo] -type = 'sled' -path = 'data/sled-repo-local' -cache_capacity = 67108864 -export_path = "data/exports-local" +type = 'postgres' +url = 'postgres://postgres:1234@localhost:5432/postgres' [store] type = 'filesystem' diff --git a/docs/postgres-planning.md b/docs/postgres-planning.md index 158ce69..49b4d6f 100644 --- a/docs/postgres-planning.md +++ b/docs/postgres-planning.md @@ -88,13 +88,13 @@ methods: ```sql CREATE TABLE aliases ( - alias VARCHAR(30) PRIMARY KEY, + alias VARCHAR(50) PRIMARY KEY, hash BYTEA NOT NULL REFERENCES hashes(hash) ON DELETE CASCADE, delete_token VARCHAR(30) NOT NULL ); -CREATE INDEX alias_hashes_index ON aliases (hash); +CREATE INDEX aliases_hash_index ON aliases (hash); ``` From 8eb4cda256f21b8a96d8b837362f4678784e5f1c Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 13:52:15 -0500 Subject: [PATCH 05/27] Implement first query --- src/repo.rs | 3 ++ src/repo/postgres.rs | 37 ++++++++++++------- src/repo/postgres/embedded.rs | 3 ++ .../migrations/V0002__create_variants.rs | 4 +- 4 files changed, 31 insertions(+), 16 deletions(-) create mode 100644 src/repo/postgres/embedded.rs diff --git a/src/repo.rs b/src/repo.rs index 51553f4..4c5d923 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -72,6 +72,9 @@ pub(crate) enum RepoError { #[error("Error in sled")] SledError(#[from] crate::repo::sled::SledError), + #[error("Error in postgres")] + PostgresError(#[from] crate::repo::postgres::PostgresError), + #[error("Upload was already claimed")] AlreadyClaimed, diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index f4e71da..c66303c 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -1,20 +1,17 @@ -mod embedded { - use refinery::embed_migrations; - - embed_migrations!("./src/repo/postgres/migrations"); -} +mod embedded; mod schema; +use diesel::prelude::*; use diesel_async::{ pooled_connection::{ - deadpool::{BuildError, Pool}, + deadpool::{BuildError, Pool, PoolError}, AsyncDieselConnectionManager, }, - AsyncPgConnection, + AsyncPgConnection, RunQueryDsl, }; use url::Url; -use super::{BaseRepo, HashRepo}; +use super::{BaseRepo, HashRepo, RepoError}; #[derive(Clone)] pub(crate) struct PostgresRepo { @@ -34,7 +31,13 @@ pub(crate) enum ConnectPostgresError { } #[derive(Debug, thiserror::Error)] -enum PostgresError {} +pub(crate) enum PostgresError { + #[error("Error in db pool")] + Pool(#[source] PoolError), + + #[error("Error in database")] + Diesel(#[source] diesel::result::Error), +} impl PostgresRepo { pub(crate) async fn connect(postgres_url: Url) -> Result { @@ -64,14 +67,22 @@ impl PostgresRepo { impl BaseRepo for PostgresRepo {} -/* -#[async_trait::async_trait] +#[async_trait::async_trait(?Send)] impl HashRepo for PostgresRepo { async fn size(&self) -> Result { - let conn = self.pool.get().await.map_err(PostgresError::from)?; + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let count = hashes + .count() + .get_result::(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(count.try_into().expect("non-negative count")) } } -*/ impl std::fmt::Debug for PostgresRepo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { diff --git a/src/repo/postgres/embedded.rs b/src/repo/postgres/embedded.rs new file mode 100644 index 0000000..482a885 --- /dev/null +++ b/src/repo/postgres/embedded.rs @@ -0,0 +1,3 @@ +use refinery::embed_migrations; + +embed_migrations!("./src/repo/postgres/migrations"); diff --git a/src/repo/postgres/migrations/V0002__create_variants.rs b/src/repo/postgres/migrations/V0002__create_variants.rs index 8377ddc..c49c62c 100644 --- a/src/repo/postgres/migrations/V0002__create_variants.rs +++ b/src/repo/postgres/migrations/V0002__create_variants.rs @@ -24,7 +24,5 @@ pub(crate) fn migration() -> String { ); }); - let s = m.make::().to_string(); - println!("{s}"); - s + m.make::().to_string() } From 8921f57a21ff202a21f270f218ca21a1dbbfd610 Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 18:30:45 -0500 Subject: [PATCH 06/27] Remove Identifier trait, Replace with Arc --- Cargo.lock | 3 + Cargo.toml | 2 +- src/backgrounded.rs | 29 ++- src/details.rs | 4 +- src/ffmpeg.rs | 4 +- src/generate.rs | 6 +- src/ingest.rs | 27 +-- src/lib.rs | 74 +++---- src/magick.rs | 4 +- src/migrate_store.rs | 29 ++- src/queue.rs | 50 ++--- src/queue/cleanup.rs | 30 +-- src/queue/process.rs | 20 +- src/range.rs | 4 +- src/repo.rs | 185 ++++++++++------ src/repo/hash.rs | 47 +++- src/repo/migrate.rs | 18 +- src/repo/postgres.rs | 201 +++++++++++++++++- .../migrations/V0001__create_hashes.rs | 2 +- .../migrations/V0002__create_variants.rs | 2 +- .../migrations/V0003__create_aliases.rs | 2 +- .../migrations/V0005__create_details.rs | 2 +- .../migrations/V0006__create_queue.rs | 4 +- src/repo/postgres/schema.rs | 40 ++-- src/repo/sled.rs | 163 +++++++------- src/repo_04.rs | 20 +- src/repo_04/sled.rs | 54 ++--- src/store.rs | 200 ++++++++--------- src/store/file_store.rs | 47 ++-- src/store/file_store/file_id.rs | 57 ----- src/store/object_store.rs | 62 +++--- src/store/object_store/object_id.rs | 37 ---- 32 files changed, 799 insertions(+), 630 deletions(-) delete mode 100644 src/store/file_store/file_id.rs delete mode 100644 src/store/object_store/object_id.rs diff --git a/Cargo.lock b/Cargo.lock index c7e9463..1a31c37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -730,6 +730,9 @@ dependencies = [ "byteorder", "diesel_derives", "itoa", + "serde_json", + "time", + "uuid", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 7f6e6db..449d844 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,7 +28,7 @@ config = "0.13.0" console-subscriber = "0.1" dashmap = "5.1.0" deadpool = { version = "0.9.5", features = ["rt_tokio_1"] } -diesel = "2.1.1" +diesel = { version = "2.1.1", features = ["postgres_backend", "serde_json", "time", "uuid"] } diesel-async = { version = "0.4.1", features = ["postgres", "deadpool"] } flume = "0.11.0" futures-core = "0.3" diff --git a/src/backgrounded.rs b/src/backgrounded.rs index 0210dac..94839a1 100644 --- a/src/backgrounded.rs +++ b/src/backgrounded.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{ error::Error, repo::{ArcRepo, UploadId}, @@ -9,19 +11,13 @@ use futures_core::Stream; use mime::APPLICATION_OCTET_STREAM; use tracing::{Instrument, Span}; -pub(crate) struct Backgrounded -where - S: Store, -{ +pub(crate) struct Backgrounded { repo: ArcRepo, - identifier: Option, + identifier: Option>, upload_id: Option, } -impl Backgrounded -where - S: Store, -{ +impl Backgrounded { pub(crate) fn disarm(mut self) { let _ = self.identifier.take(); let _ = self.upload_id.take(); @@ -31,12 +27,13 @@ where self.upload_id } - pub(crate) fn identifier(&self) -> Option<&S::Identifier> { + pub(crate) fn identifier(&self) -> Option<&Arc> { self.identifier.as_ref() } - pub(crate) async fn proxy

(repo: ArcRepo, store: S, stream: P) -> Result + pub(crate) async fn proxy(repo: ArcRepo, store: S, stream: P) -> Result where + S: Store, P: Stream> + Unpin + 'static, { let mut this = Self { @@ -50,8 +47,9 @@ where Ok(this) } - async fn do_proxy

(&mut self, store: S, stream: P) -> Result<(), Error> + async fn do_proxy(&mut self, store: S, stream: P) -> Result<(), Error> where + S: Store, P: Stream> + Unpin + 'static, { self.upload_id = Some(self.repo.create_upload().await?); @@ -68,10 +66,7 @@ where } } -impl Drop for Backgrounded -where - S: Store, -{ +impl Drop for Backgrounded { fn drop(&mut self) { let any_items = self.identifier.is_some() || self.upload_id.is_some(); @@ -90,7 +85,7 @@ where tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { actix_rt::spawn( async move { - let _ = crate::queue::cleanup_identifier(&repo, identifier).await; + let _ = crate::queue::cleanup_identifier(&repo, &identifier).await; } .instrument(cleanup_span), ) diff --git a/src/details.rs b/src/details.rs index edf921b..d632e32 100644 --- a/src/details.rs +++ b/src/details.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{ bytes_stream::BytesStream, discover::Discovery, @@ -103,7 +105,7 @@ impl Details { pub(crate) async fn from_store( store: &S, - identifier: &S::Identifier, + identifier: &Arc, timeout: u64, ) -> Result { let mut buf = BytesStream::new(); diff --git a/src/ffmpeg.rs b/src/ffmpeg.rs index 3602b5f..bab1582 100644 --- a/src/ffmpeg.rs +++ b/src/ffmpeg.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{ error_code::ErrorCode, formats::InternalVideoFormat, @@ -132,7 +134,7 @@ impl ThumbnailFormat { #[tracing::instrument(skip(store))] pub(crate) async fn thumbnail( store: S, - from: S::Identifier, + from: Arc, input_format: InternalVideoFormat, format: ThumbnailFormat, timeout: u64, diff --git a/src/generate.rs b/src/generate.rs index 1976012..6808590 100644 --- a/src/generate.rs +++ b/src/generate.rs @@ -5,7 +5,7 @@ use crate::{ ffmpeg::ThumbnailFormat, formats::{InputProcessableFormat, InternalVideoFormat}, repo::{Alias, ArcRepo, Hash, VariantAlreadyExists}, - store::{Identifier, Store}, + store::Store, }; use actix_web::web::Bytes; use std::{path::PathBuf, time::Instant}; @@ -91,7 +91,7 @@ async fn process( let permit = crate::PROCESS_SEMAPHORE.acquire().await; let identifier = if let Some(identifier) = repo.still_identifier_from_alias(&alias).await? { - S::Identifier::from_arc(identifier)? + identifier } else { let Some(identifier) = repo.identifier(hash.clone()).await? else { return Err(UploadError::MissingIdentifier.into()); @@ -101,7 +101,7 @@ async fn process( let reader = crate::ffmpeg::thumbnail( store.clone(), - S::Identifier::from_arc(identifier)?, + identifier, input_format.unwrap_or(InternalVideoFormat::Mp4), thumbnail_format, media.process_timeout, diff --git a/src/ingest.rs b/src/ingest.rs index 938e277..e0cd34b 100644 --- a/src/ingest.rs +++ b/src/ingest.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{ bytes_stream::BytesStream, either::Either, @@ -15,15 +17,12 @@ mod hasher; use hasher::Hasher; #[derive(Debug)] -pub(crate) struct Session -where - S: Store, -{ +pub(crate) struct Session { repo: ArcRepo, delete_token: DeleteToken, hash: Option, alias: Option, - identifier: Option, + identifier: Option>, } #[tracing::instrument(skip(stream))] @@ -49,7 +48,7 @@ pub(crate) async fn ingest( stream: impl Stream> + Unpin + 'static, declared_alias: Option, media: &crate::config::Media, -) -> Result, Error> +) -> Result where S: Store, { @@ -131,11 +130,11 @@ where #[tracing::instrument(level = "trace", skip_all)] async fn save_upload( - session: &mut Session, + session: &mut Session, repo: &ArcRepo, store: &S, hash: Hash, - identifier: &S::Identifier, + identifier: &Arc, ) -> Result<(), Error> where S: Store, @@ -153,10 +152,7 @@ where Ok(()) } -impl Session -where - S: Store, -{ +impl Session { pub(crate) fn disarm(mut self) -> DeleteToken { let _ = self.hash.take(); let _ = self.alias.take(); @@ -206,10 +202,7 @@ where } } -impl Drop for Session -where - S: Store, -{ +impl Drop for Session { fn drop(&mut self) { let any_items = self.hash.is_some() || self.alias.is_some() || self.identifier.is_some(); @@ -258,7 +251,7 @@ where tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { actix_rt::spawn( async move { - let _ = crate::queue::cleanup_identifier(&repo, identifier).await; + let _ = crate::queue::cleanup_identifier(&repo, &identifier).await; } .instrument(cleanup_span), ) diff --git a/src/lib.rs b/src/lib.rs index 21952fc..e6476c7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,6 +45,7 @@ use reqwest_middleware::{ClientBuilder, ClientWithMiddleware}; use reqwest_tracing::TracingMiddleware; use rusty_s3::UrlStyle; use std::{ + marker::PhantomData, path::Path, path::PathBuf, sync::Arc, @@ -69,7 +70,7 @@ use self::{ queue::queue_generate, repo::{sled::SledRepo, Alias, DeleteToken, Hash, Repo, UploadId, UploadResult}, serde_str::Serde, - store::{file_store::FileStore, object_store::ObjectStore, Identifier, Store}, + store::{file_store::FileStore, object_store::ObjectStore, Store}, stream::{empty, once, StreamLimit, StreamMap, StreamTimeout}, }; @@ -93,7 +94,7 @@ async fn ensure_details( config: &Configuration, alias: &Alias, ) -> Result { - let Some(identifier) = repo.identifier_from_alias(alias).await?.map(S::Identifier::from_arc).transpose()? else { + let Some(identifier) = repo.identifier_from_alias(alias).await? else { return Err(UploadError::MissingAlias.into()); }; @@ -117,10 +118,10 @@ async fn ensure_details( } } -struct Upload(Value>); +struct Upload(Value, PhantomData); impl FormData for Upload { - type Item = Session; + type Item = Session; type Error = Error; fn form(req: &HttpRequest) -> Form { @@ -172,14 +173,14 @@ impl FormData for Upload { } fn extract(value: Value) -> Result { - Ok(Upload(value)) + Ok(Upload(value, PhantomData)) } } -struct Import(Value>); +struct Import(Value, PhantomData); impl FormData for Import { - type Item = Session; + type Item = Session; type Error = Error; fn form(req: &actix_web::HttpRequest) -> Form { @@ -241,14 +242,14 @@ impl FormData for Import { where Self: Sized, { - Ok(Import(value)) + Ok(Import(value, PhantomData)) } } /// Handle responding to successful uploads #[tracing::instrument(name = "Uploaded files", skip(value, repo, store, config))] async fn upload( - Multipart(Upload(value)): Multipart>, + Multipart(Upload(value, _)): Multipart>, repo: web::Data, store: web::Data, config: web::Data, @@ -259,7 +260,7 @@ async fn upload( /// Handle responding to successful uploads #[tracing::instrument(name = "Imported files", skip(value, repo, store, config))] async fn import( - Multipart(Import(value)): Multipart>, + Multipart(Import(value, _)): Multipart>, repo: web::Data, store: web::Data, config: web::Data, @@ -270,7 +271,7 @@ async fn import( /// Handle responding to successful uploads #[tracing::instrument(name = "Uploaded files", skip(value, repo, store, config))] async fn handle_upload( - value: Value>, + value: Value, repo: web::Data, store: web::Data, config: web::Data, @@ -312,10 +313,10 @@ async fn handle_upload( }))) } -struct BackgroundedUpload(Value>); +struct BackgroundedUpload(Value, PhantomData); impl FormData for BackgroundedUpload { - type Item = Backgrounded; + type Item = Backgrounded; type Error = Error; fn form(req: &actix_web::HttpRequest) -> Form { @@ -371,13 +372,13 @@ impl FormData for BackgroundedUpload { where Self: Sized, { - Ok(BackgroundedUpload(value)) + Ok(BackgroundedUpload(value, PhantomData)) } } #[tracing::instrument(name = "Uploaded files", skip(value, repo))] async fn upload_backgrounded( - Multipart(BackgroundedUpload(value)): Multipart>, + Multipart(BackgroundedUpload(value, _)): Multipart>, repo: web::Data, ) -> Result { let images = value @@ -394,11 +395,7 @@ async fn upload_backgrounded( for image in &images { let upload_id = image.result.upload_id().expect("Upload ID exists"); - let identifier = image - .result - .identifier() - .expect("Identifier exists") - .to_bytes()?; + let identifier = image.result.identifier().expect("Identifier exists"); queue::queue_ingest(&repo, identifier, upload_id, None).await?; @@ -560,10 +557,7 @@ async fn do_download_backgrounded( let backgrounded = Backgrounded::proxy((**repo).clone(), (**store).clone(), stream).await?; let upload_id = backgrounded.upload_id().expect("Upload ID exists"); - let identifier = backgrounded - .identifier() - .expect("Identifier exists") - .to_bytes()?; + let identifier = backgrounded.identifier().expect("Identifier exists"); queue::queue_ingest(&repo, identifier, upload_id, None).await?; @@ -764,8 +758,6 @@ async fn process_details( let identifier = repo .variant_identifier(hash, thumbnail_string) .await? - .map(S::Identifier::from_arc) - .transpose()? .ok_or(UploadError::MissingAlias)?; let details = repo.details(&identifier).await?; @@ -856,11 +848,7 @@ async fn process( .await?; } - let identifier_opt = repo - .variant_identifier(hash.clone(), path_string) - .await? - .map(S::Identifier::from_arc) - .transpose()?; + let identifier_opt = repo.variant_identifier(hash.clone(), path_string).await?; if let Some(identifier) = identifier_opt { let details = repo.details(&identifier).await?; @@ -980,11 +968,7 @@ async fn process_head( .await?; } - let identifier_opt = repo - .variant_identifier(hash.clone(), path_string) - .await? - .map(S::Identifier::from_arc) - .transpose()?; + let identifier_opt = repo.variant_identifier(hash.clone(), path_string).await?; if let Some(identifier) = identifier_opt { let details = repo.details(&identifier).await?; @@ -1047,11 +1031,7 @@ async fn process_backgrounded( return Ok(HttpResponse::BadRequest().finish()); }; - let identifier_opt = repo - .variant_identifier(hash.clone(), path_string) - .await? - .map(S::Identifier::from_arc) - .transpose()?; + let identifier_opt = repo.variant_identifier(hash.clone(), path_string).await?; if identifier_opt.is_some() { return Ok(HttpResponse::Accepted().finish()); @@ -1185,7 +1165,7 @@ async fn do_serve( (hash, alias, true) }; - let Some(identifier) = repo.identifier(hash.clone()).await?.map(Identifier::from_arc).transpose()? else { + let Some(identifier) = repo.identifier(hash.clone()).await? else { tracing::warn!( "Original File identifier for hash {hash:?} is missing, queue cleanup task", ); @@ -1250,7 +1230,7 @@ async fn do_serve_head( store: web::Data, config: web::Data, ) -> Result { - let Some(identifier) = repo.identifier_from_alias(&alias).await?.map(S::Identifier::from_arc).transpose()? else { + let Some(identifier) = repo.identifier_from_alias(&alias).await? else { // Invalid alias return Ok(HttpResponse::NotFound().finish()); }; @@ -1268,7 +1248,7 @@ async fn do_serve_head( async fn ranged_file_head_resp( store: &S, - identifier: S::Identifier, + identifier: Arc, range: Option>, details: Details, ) -> Result { @@ -1303,7 +1283,7 @@ async fn ranged_file_head_resp( async fn ranged_file_resp( store: &S, - identifier: S::Identifier, + identifier: Arc, range: Option>, details: Details, not_found: bool, @@ -1555,7 +1535,7 @@ async fn identifier( } }; - let Some(identifier) = repo.identifier_from_alias(&alias).await?.map(S::Identifier::from_arc).transpose()? else { + let Some(identifier) = repo.identifier_from_alias(&alias).await? else { // Invalid alias return Ok(HttpResponse::NotFound().json(serde_json::json!({ "msg": "No identifiers associated with provided alias" @@ -1564,7 +1544,7 @@ async fn identifier( Ok(HttpResponse::Ok().json(&serde_json::json!({ "msg": "ok", - "identifier": identifier.string_repr(), + "identifier": identifier.as_ref(), }))) } diff --git a/src/magick.rs b/src/magick.rs index b9efaa6..29c9636 100644 --- a/src/magick.rs +++ b/src/magick.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{ error_code::ErrorCode, formats::ProcessableFormat, @@ -140,7 +142,7 @@ where pub(crate) async fn process_image_store_read( store: &S, - identifier: &S::Identifier, + identifier: &Arc, args: Vec, input_format: ProcessableFormat, format: ProcessableFormat, diff --git a/src/migrate_store.rs b/src/migrate_store.rs index 5be16b9..78cb0b0 100644 --- a/src/migrate_store.rs +++ b/src/migrate_store.rs @@ -1,6 +1,9 @@ use std::{ rc::Rc, - sync::atomic::{AtomicU64, Ordering}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, time::{Duration, Instant}, }; @@ -8,7 +11,7 @@ use crate::{ details::Details, error::{Error, UploadError}, repo::{ArcRepo, Hash}, - store::{Identifier, Store}, + store::Store, stream::IntoStreamer, }; @@ -103,7 +106,7 @@ where } // Hashes are read in a consistent order - let mut stream = repo.hashes().await.into_streamer(); + let mut stream = repo.hashes().into_streamer(); let state = Rc::new(MigrateState { repo: repo.clone(), @@ -169,7 +172,7 @@ where let current_index = index.fetch_add(1, Ordering::Relaxed); let original_identifier = match repo.identifier(hash.clone()).await { - Ok(Some(identifier)) => S1::Identifier::from_arc(identifier)?, + Ok(Some(identifier)) => identifier, Ok(None) => { tracing::warn!( "Original File identifier for hash {hash:?} is missing, queue cleanup task", @@ -214,8 +217,6 @@ where } if let Some(identifier) = repo.motion_identifier(hash.clone()).await? { - let identifier = S1::Identifier::from_arc(identifier)?; - if !repo.is_migrated(&identifier).await? { match migrate_file(repo, from, to, &identifier, *skip_missing_files, *timeout).await { Ok(new_identifier) => { @@ -245,8 +246,6 @@ where } for (variant, identifier) in repo.variants(hash.clone()).await? { - let identifier = S1::Identifier::from_arc(identifier)?; - if !repo.is_migrated(&identifier).await? { match migrate_file(repo, from, to, &identifier, *skip_missing_files, *timeout).await { Ok(new_identifier) => { @@ -339,10 +338,10 @@ async fn migrate_file( repo: &ArcRepo, from: &S1, to: &S2, - identifier: &S1::Identifier, + identifier: &Arc, skip_missing_files: bool, timeout: u64, -) -> Result +) -> Result, MigrateError> where S1: Store, S2: Store, @@ -382,9 +381,9 @@ async fn do_migrate_file( repo: &ArcRepo, from: &S1, to: &S2, - identifier: &S1::Identifier, + identifier: &Arc, timeout: u64, -) -> Result +) -> Result, MigrateError> where S1: Store, S2: Store, @@ -421,11 +420,7 @@ where Ok(new_identifier) } -async fn migrate_details(repo: &ArcRepo, from: &I1, to: &I2) -> Result<(), Error> -where - I1: Identifier, - I2: Identifier, -{ +async fn migrate_details(repo: &ArcRepo, from: &Arc, to: &Arc) -> Result<(), Error> { if let Some(details) = repo.details(from).await? { repo.relate_details(to, &details).await?; repo.cleanup_details(from).await?; diff --git a/src/queue.rs b/src/queue.rs index 8a3fc20..a2c9deb 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -5,7 +5,7 @@ use crate::{ formats::InputProcessableFormat, repo::{Alias, DeleteToken, FullRepo, Hash, JobId, UploadId}, serde_str::Serde, - store::{Identifier, Store}, + store::Store, }; use base64::{prelude::BASE64_STANDARD, Engine}; use std::{ @@ -55,7 +55,7 @@ enum Cleanup { hash: Hash, }, Identifier { - identifier: Base64Bytes, + identifier: String, }, Alias { alias: Serde, @@ -74,7 +74,7 @@ enum Cleanup { #[derive(Debug, serde::Deserialize, serde::Serialize)] enum Process { Ingest { - identifier: Base64Bytes, + identifier: String, upload_id: Serde, declared_alias: Option>, }, @@ -91,7 +91,7 @@ pub(crate) async fn cleanup_alias( alias: Alias, token: DeleteToken, ) -> Result<(), Error> { - let job = serde_json::to_vec(&Cleanup::Alias { + let job = serde_json::to_string(&Cleanup::Alias { alias: Serde::new(alias), token: Serde::new(token), }) @@ -101,17 +101,17 @@ pub(crate) async fn cleanup_alias( } pub(crate) async fn cleanup_hash(repo: &Arc, hash: Hash) -> Result<(), Error> { - let job = serde_json::to_vec(&Cleanup::Hash { hash }).map_err(UploadError::PushJob)?; + let job = serde_json::to_string(&Cleanup::Hash { hash }).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job.into()).await?; Ok(()) } -pub(crate) async fn cleanup_identifier( +pub(crate) async fn cleanup_identifier( repo: &Arc, - identifier: I, + identifier: &Arc, ) -> Result<(), Error> { - let job = serde_json::to_vec(&Cleanup::Identifier { - identifier: Base64Bytes(identifier.to_bytes()?), + let job = serde_json::to_string(&Cleanup::Identifier { + identifier: identifier.to_string(), }) .map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job.into()).await?; @@ -124,37 +124,37 @@ async fn cleanup_variants( variant: Option, ) -> Result<(), Error> { let job = - serde_json::to_vec(&Cleanup::Variant { hash, variant }).map_err(UploadError::PushJob)?; + serde_json::to_string(&Cleanup::Variant { hash, variant }).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job.into()).await?; Ok(()) } pub(crate) async fn cleanup_outdated_proxies(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_vec(&Cleanup::OutdatedProxies).map_err(UploadError::PushJob)?; + let job = serde_json::to_string(&Cleanup::OutdatedProxies).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job.into()).await?; Ok(()) } pub(crate) async fn cleanup_outdated_variants(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_vec(&Cleanup::OutdatedVariants).map_err(UploadError::PushJob)?; + let job = serde_json::to_string(&Cleanup::OutdatedVariants).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job.into()).await?; Ok(()) } pub(crate) async fn cleanup_all_variants(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_vec(&Cleanup::AllVariants).map_err(UploadError::PushJob)?; + let job = serde_json::to_string(&Cleanup::AllVariants).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job.into()).await?; Ok(()) } pub(crate) async fn queue_ingest( repo: &Arc, - identifier: Vec, + identifier: &Arc, upload_id: UploadId, declared_alias: Option, ) -> Result<(), Error> { - let job = serde_json::to_vec(&Process::Ingest { - identifier: Base64Bytes(identifier), + let job = serde_json::to_string(&Process::Ingest { + identifier: identifier.to_string(), declared_alias: declared_alias.map(Serde::new), upload_id: Serde::new(upload_id), }) @@ -170,7 +170,7 @@ pub(crate) async fn queue_generate( process_path: PathBuf, process_args: Vec, ) -> Result<(), Error> { - let job = serde_json::to_vec(&Process::Generate { + let job = serde_json::to_string(&Process::Generate { target_format, source: Serde::new(source), process_path, @@ -220,7 +220,7 @@ async fn process_jobs( &'a Arc, &'a S, &'a Configuration, - &'a [u8], + &'a str, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { @@ -284,13 +284,13 @@ where &'a Arc, &'a S, &'a Configuration, - &'a [u8], + &'a str, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { let fut = async { - let (job_id, bytes) = repo.pop(queue, worker_id).await?; + let (job_id, string) = repo.pop(queue, worker_id).await?; let span = tracing::info_span!("Running Job"); @@ -303,7 +303,7 @@ where queue, worker_id, job_id, - (callback)(repo, store, config, bytes.as_ref()), + (callback)(repo, store, config, string.as_ref()), ) }) .instrument(span) @@ -337,7 +337,7 @@ async fn process_image_jobs( &'a S, &'a ProcessMap, &'a Configuration, - &'a [u8], + &'a str, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { @@ -373,13 +373,13 @@ where &'a S, &'a ProcessMap, &'a Configuration, - &'a [u8], + &'a str, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { let fut = async { - let (job_id, bytes) = repo.pop(queue, worker_id).await?; + let (job_id, string) = repo.pop(queue, worker_id).await?; let span = tracing::info_span!("Running Job"); @@ -392,7 +392,7 @@ where queue, worker_id, job_id, - (callback)(repo, store, process_map, config, bytes.as_ref()), + (callback)(repo, store, process_map, config, string.as_ref()), ) }) .instrument(span) diff --git a/src/queue/cleanup.rs b/src/queue/cleanup.rs index 34b2424..a6b76fd 100644 --- a/src/queue/cleanup.rs +++ b/src/queue/cleanup.rs @@ -1,10 +1,12 @@ +use std::sync::Arc; + use crate::{ config::Configuration, error::{Error, UploadError}, - queue::{Base64Bytes, Cleanup, LocalBoxFuture}, + queue::{Cleanup, LocalBoxFuture}, repo::{Alias, ArcRepo, DeleteToken, Hash}, serde_str::Serde, - store::{Identifier, Store}, + store::Store, stream::IntoStreamer, }; @@ -12,18 +14,18 @@ pub(super) fn perform<'a, S>( repo: &'a ArcRepo, store: &'a S, configuration: &'a Configuration, - job: &'a [u8], + job: &'a str, ) -> LocalBoxFuture<'a, Result<(), Error>> where S: Store, { Box::pin(async move { - match serde_json::from_slice(job) { + match serde_json::from_str(job) { Ok(job) => match job { Cleanup::Hash { hash: in_hash } => hash(repo, in_hash).await?, Cleanup::Identifier { - identifier: Base64Bytes(in_identifier), - } => identifier(repo, store, in_identifier).await?, + identifier: in_identifier, + } => identifier(repo, store, Arc::from(in_identifier)).await?, Cleanup::Alias { alias: stored_alias, token, @@ -50,20 +52,18 @@ where } #[tracing::instrument(skip_all)] -async fn identifier(repo: &ArcRepo, store: &S, identifier: Vec) -> Result<(), Error> +async fn identifier(repo: &ArcRepo, store: &S, identifier: Arc) -> Result<(), Error> where S: Store, { - let identifier = S::Identifier::from_bytes(identifier)?; - let mut errors = Vec::new(); if let Err(e) = store.remove(&identifier).await { - errors.push(e); + errors.push(UploadError::from(e)); } if let Err(e) = repo.cleanup_details(&identifier).await { - errors.push(e); + errors.push(UploadError::from(e)); } for error in errors { @@ -100,7 +100,7 @@ async fn hash(repo: &ArcRepo, hash: Hash) -> Result<(), Error> { idents.extend(repo.motion_identifier(hash.clone()).await?); for identifier in idents { - let _ = super::cleanup_identifier(repo, identifier).await; + let _ = super::cleanup_identifier(repo, &identifier).await; } repo.cleanup_hash(hash).await?; @@ -136,7 +136,7 @@ async fn alias(repo: &ArcRepo, alias: Alias, token: DeleteToken) -> Result<(), E #[tracing::instrument(skip_all)] async fn all_variants(repo: &ArcRepo) -> Result<(), Error> { - let mut hash_stream = repo.hashes().await.into_streamer(); + let mut hash_stream = repo.hashes().into_streamer(); while let Some(res) = hash_stream.next().await { let hash = res?; @@ -193,7 +193,7 @@ async fn hash_variant( .variant_identifier(hash.clone(), target_variant.clone()) .await? { - super::cleanup_identifier(repo, identifier).await?; + super::cleanup_identifier(repo, &identifier).await?; } repo.remove_variant(hash.clone(), target_variant.clone()) @@ -203,7 +203,7 @@ async fn hash_variant( for (variant, identifier) in repo.variants(hash.clone()).await? { repo.remove_variant(hash.clone(), variant.clone()).await?; repo.remove_variant_access(hash.clone(), variant).await?; - super::cleanup_identifier(repo, identifier).await?; + super::cleanup_identifier(repo, &identifier).await?; } } diff --git a/src/queue/process.rs b/src/queue/process.rs index 39a1fb9..a5f8e43 100644 --- a/src/queue/process.rs +++ b/src/queue/process.rs @@ -4,36 +4,36 @@ use crate::{ error::{Error, UploadError}, formats::InputProcessableFormat, ingest::Session, - queue::{Base64Bytes, LocalBoxFuture, Process}, + queue::{LocalBoxFuture, Process}, repo::{Alias, ArcRepo, UploadId, UploadResult}, serde_str::Serde, - store::{Identifier, Store}, + store::Store, stream::StreamMap, }; -use std::path::PathBuf; +use std::{path::PathBuf, sync::Arc}; pub(super) fn perform<'a, S>( repo: &'a ArcRepo, store: &'a S, process_map: &'a ProcessMap, config: &'a Configuration, - job: &'a [u8], + job: &'a str, ) -> LocalBoxFuture<'a, Result<(), Error>> where S: Store + 'static, { Box::pin(async move { - match serde_json::from_slice(job) { + match serde_json::from_str(job) { Ok(job) => match job { Process::Ingest { - identifier: Base64Bytes(identifier), + identifier, upload_id, declared_alias, } => { process_ingest( repo, store, - identifier, + Arc::from(identifier), Serde::into_inner(upload_id), declared_alias.map(Serde::into_inner), &config.media, @@ -72,7 +72,7 @@ where async fn process_ingest( repo: &ArcRepo, store: &S, - unprocessed_identifier: Vec, + unprocessed_identifier: Arc, upload_id: UploadId, declared_alias: Option, media: &crate::config::Media, @@ -81,8 +81,6 @@ where S: Store + 'static, { let fut = async { - let unprocessed_identifier = S::Identifier::from_bytes(unprocessed_identifier)?; - let ident = unprocessed_identifier.clone(); let store2 = store.clone(); let repo = repo.clone(); @@ -97,7 +95,7 @@ where let session = crate::ingest::ingest(&repo, &store2, stream, declared_alias, &media).await?; - Ok(session) as Result, Error> + Ok(session) as Result }) .await; diff --git a/src/range.rs b/src/range.rs index 976f384..cf76319 100644 --- a/src/range.rs +++ b/src/range.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{ error::{Error, UploadError}, store::Store, @@ -26,7 +28,7 @@ pub(crate) fn chop_bytes( pub(crate) async fn chop_store( byte_range: &ByteRangeSpec, store: &S, - identifier: &S::Identifier, + identifier: &Arc, length: u64, ) -> Result>, Error> { if let Some((start, end)) = byte_range.to_satisfiable_range(length) { diff --git a/src/repo.rs b/src/repo.rs index 4c5d923..22318d4 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -2,7 +2,6 @@ use crate::{ config, details::Details, error_code::{ErrorCode, OwnedErrorCode}, - store::{Identifier, StoreError}, stream::LocalBoxStream, }; use base64::Engine; @@ -86,6 +85,7 @@ impl RepoError { pub(crate) const fn error_code(&self) -> ErrorCode { match self { Self::SledError(e) => e.error_code(), + Self::PostgresError(e) => e.error_code(), Self::AlreadyClaimed => ErrorCode::ALREADY_CLAIMED, Self::Canceled => ErrorCode::PANIC, } @@ -111,7 +111,7 @@ pub(crate) trait FullRepo: async fn health_check(&self) -> Result<(), RepoError>; #[tracing::instrument(skip(self))] - async fn identifier_from_alias(&self, alias: &Alias) -> Result>, RepoError> { + async fn identifier_from_alias(&self, alias: &Alias) -> Result>, RepoError> { let Some(hash) = self.hash(alias).await? else { return Ok(None); }; @@ -132,7 +132,7 @@ pub(crate) trait FullRepo: async fn still_identifier_from_alias( &self, alias: &Alias, - ) -> Result>, StoreError> { + ) -> Result>, RepoError> { let Some(hash) = self.hash(alias).await? else { return Ok(None); }; @@ -372,13 +372,13 @@ impl JobId { #[async_trait::async_trait(?Send)] pub(crate) trait QueueRepo: BaseRepo { - async fn push(&self, queue: &'static str, job: Arc<[u8]>) -> Result; + async fn push(&self, queue: &'static str, job: Arc) -> Result; async fn pop( &self, queue: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc<[u8]>), RepoError>; + ) -> Result<(JobId, Arc), RepoError>; async fn heartbeat( &self, @@ -400,7 +400,7 @@ impl QueueRepo for Arc where T: QueueRepo, { - async fn push(&self, queue: &'static str, job: Arc<[u8]>) -> Result { + async fn push(&self, queue: &'static str, job: Arc) -> Result { T::push(self, queue, job).await } @@ -408,7 +408,7 @@ where &self, queue: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc<[u8]>), RepoError> { + ) -> Result<(JobId, Arc), RepoError> { T::pop(self, queue, worker_id).await } @@ -460,12 +460,12 @@ where pub(crate) trait DetailsRepo: BaseRepo { async fn relate_details( &self, - identifier: &dyn Identifier, + identifier: &Arc, details: &Details, - ) -> Result<(), StoreError>; - async fn details(&self, identifier: &dyn Identifier) -> Result, StoreError>; + ) -> Result<(), RepoError>; + async fn details(&self, identifier: &Arc) -> Result, RepoError>; - async fn cleanup_details(&self, identifier: &dyn Identifier) -> Result<(), StoreError>; + async fn cleanup_details(&self, identifier: &Arc) -> Result<(), RepoError>; } #[async_trait::async_trait(?Send)] @@ -475,17 +475,17 @@ where { async fn relate_details( &self, - identifier: &dyn Identifier, + identifier: &Arc, details: &Details, - ) -> Result<(), StoreError> { + ) -> Result<(), RepoError> { T::relate_details(self, identifier, details).await } - async fn details(&self, identifier: &dyn Identifier) -> Result, StoreError> { + async fn details(&self, identifier: &Arc) -> Result, RepoError> { T::details(self, identifier).await } - async fn cleanup_details(&self, identifier: &dyn Identifier) -> Result<(), StoreError> { + async fn cleanup_details(&self, identifier: &Arc) -> Result<(), RepoError> { T::cleanup_details(self, identifier).await } } @@ -496,11 +496,11 @@ pub(crate) trait StoreMigrationRepo: BaseRepo { async fn mark_migrated( &self, - old_identifier: &dyn Identifier, - new_identifier: &dyn Identifier, - ) -> Result<(), StoreError>; + old_identifier: &Arc, + new_identifier: &Arc, + ) -> Result<(), RepoError>; - async fn is_migrated(&self, identifier: &dyn Identifier) -> Result; + async fn is_migrated(&self, identifier: &Arc) -> Result; async fn clear(&self) -> Result<(), RepoError>; } @@ -516,13 +516,13 @@ where async fn mark_migrated( &self, - old_identifier: &dyn Identifier, - new_identifier: &dyn Identifier, - ) -> Result<(), StoreError> { + old_identifier: &Arc, + new_identifier: &Arc, + ) -> Result<(), RepoError> { T::mark_migrated(self, old_identifier, new_identifier).await } - async fn is_migrated(&self, identifier: &dyn Identifier) -> Result { + async fn is_migrated(&self, identifier: &Arc) -> Result { T::is_migrated(self, identifier).await } @@ -569,12 +569,87 @@ impl HashPage { } } +pub(crate) struct HashStream { + repo: Option, + page_future: + Option>>>>, + page: Option, +} + +impl futures_core::Stream for HashStream { + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + + loop { + let Some(repo) = &this.repo else { + return std::task::Poll::Ready(None); + }; + + let slug = if let Some(page) = &mut this.page { + // popping last in page is fine - we reversed them + if let Some(hash) = page.hashes.pop() { + return std::task::Poll::Ready(Some(Ok(hash))); + } + + let slug = page.next(); + this.page.take(); + + if let Some(slug) = slug { + Some(slug) + } else { + this.repo.take(); + return std::task::Poll::Ready(None); + } + } else { + None + }; + + if let Some(page_future) = &mut this.page_future { + let res = std::task::ready!(page_future.as_mut().poll(cx)); + + this.page_future.take(); + + match res { + Ok(mut page) => { + // reverse because we use `pop` to fetch next + page.hashes.reverse(); + + this.page = Some(page); + } + Err(e) => { + this.repo.take(); + + return std::task::Poll::Ready(Some(Err(e))); + } + } + } else { + let repo = repo.clone(); + + this.page_future = Some(Box::pin(async move { repo.hash_page(slug, 100).await })); + } + } + } +} + +impl dyn FullRepo { + pub(crate) fn hashes(self: &Arc) -> HashStream { + HashStream { + repo: Some(self.clone()), + page_future: None, + page: None, + } + } +} + #[async_trait::async_trait(?Send)] pub(crate) trait HashRepo: BaseRepo { async fn size(&self) -> Result; - async fn hashes(&self) -> LocalBoxStream<'static, Result>; - async fn hash_page(&self, slug: Option, limit: usize) -> Result { let hash = slug.as_deref().and_then(hash_from_slug); @@ -604,8 +679,8 @@ pub(crate) trait HashRepo: BaseRepo { async fn create_hash( &self, hash: Hash, - identifier: &dyn Identifier, - ) -> Result, StoreError> { + identifier: &Arc, + ) -> Result, RepoError> { self.create_hash_with_timestamp(hash, identifier, time::OffsetDateTime::now_utc()) .await } @@ -613,38 +688,34 @@ pub(crate) trait HashRepo: BaseRepo { async fn create_hash_with_timestamp( &self, hash: Hash, - identifier: &dyn Identifier, + identifier: &Arc, timestamp: time::OffsetDateTime, - ) -> Result, StoreError>; + ) -> Result, RepoError>; - async fn update_identifier( - &self, - hash: Hash, - identifier: &dyn Identifier, - ) -> Result<(), StoreError>; + async fn update_identifier(&self, hash: Hash, identifier: &Arc) -> Result<(), RepoError>; - async fn identifier(&self, hash: Hash) -> Result>, RepoError>; + async fn identifier(&self, hash: Hash) -> Result>, RepoError>; async fn relate_variant_identifier( &self, hash: Hash, variant: String, - identifier: &dyn Identifier, - ) -> Result, StoreError>; + identifier: &Arc, + ) -> Result, RepoError>; async fn variant_identifier( &self, hash: Hash, variant: String, - ) -> Result>, RepoError>; - async fn variants(&self, hash: Hash) -> Result)>, RepoError>; + ) -> Result>, RepoError>; + async fn variants(&self, hash: Hash) -> Result)>, RepoError>; async fn remove_variant(&self, hash: Hash, variant: String) -> Result<(), RepoError>; async fn relate_motion_identifier( &self, hash: Hash, - identifier: &dyn Identifier, - ) -> Result<(), StoreError>; - async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError>; + identifier: &Arc, + ) -> Result<(), RepoError>; + async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError>; async fn cleanup_hash(&self, hash: Hash) -> Result<(), RepoError>; } @@ -658,10 +729,6 @@ where T::size(self).await } - async fn hashes(&self) -> LocalBoxStream<'static, Result> { - T::hashes(self).await - } - async fn bound(&self, hash: Hash) -> Result, RepoError> { T::bound(self, hash).await } @@ -685,21 +752,17 @@ where async fn create_hash_with_timestamp( &self, hash: Hash, - identifier: &dyn Identifier, + identifier: &Arc, timestamp: time::OffsetDateTime, - ) -> Result, StoreError> { + ) -> Result, RepoError> { T::create_hash_with_timestamp(self, hash, identifier, timestamp).await } - async fn update_identifier( - &self, - hash: Hash, - identifier: &dyn Identifier, - ) -> Result<(), StoreError> { + async fn update_identifier(&self, hash: Hash, identifier: &Arc) -> Result<(), RepoError> { T::update_identifier(self, hash, identifier).await } - async fn identifier(&self, hash: Hash) -> Result>, RepoError> { + async fn identifier(&self, hash: Hash) -> Result>, RepoError> { T::identifier(self, hash).await } @@ -707,8 +770,8 @@ where &self, hash: Hash, variant: String, - identifier: &dyn Identifier, - ) -> Result, StoreError> { + identifier: &Arc, + ) -> Result, RepoError> { T::relate_variant_identifier(self, hash, variant, identifier).await } @@ -716,11 +779,11 @@ where &self, hash: Hash, variant: String, - ) -> Result>, RepoError> { + ) -> Result>, RepoError> { T::variant_identifier(self, hash, variant).await } - async fn variants(&self, hash: Hash) -> Result)>, RepoError> { + async fn variants(&self, hash: Hash) -> Result)>, RepoError> { T::variants(self, hash).await } @@ -731,12 +794,12 @@ where async fn relate_motion_identifier( &self, hash: Hash, - identifier: &dyn Identifier, - ) -> Result<(), StoreError> { + identifier: &Arc, + ) -> Result<(), RepoError> { T::relate_motion_identifier(self, hash, identifier).await } - async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError> { + async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError> { T::motion_identifier(self, hash).await } diff --git a/src/repo/hash.rs b/src/repo/hash.rs index 18c6dab..b3c8036 100644 --- a/src/repo/hash.rs +++ b/src/repo/hash.rs @@ -1,13 +1,42 @@ +use diesel::{backend::Backend, sql_types::VarChar, AsExpression, FromSqlRow}; + use crate::formats::InternalFormat; use std::sync::Arc; -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, AsExpression, FromSqlRow)] +#[diesel(sql_type = VarChar)] pub(crate) struct Hash { hash: Arc<[u8; 32]>, size: u64, format: InternalFormat, } +impl diesel::serialize::ToSql for Hash { + fn to_sql<'b>( + &'b self, + out: &mut diesel::serialize::Output<'b, '_, diesel::pg::Pg>, + ) -> diesel::serialize::Result { + let s = self.to_base64(); + + >::to_sql( + &s, + &mut out.reborrow(), + ) + } +} + +impl diesel::deserialize::FromSql for Hash +where + B: Backend, + String: diesel::deserialize::FromSql, +{ + fn from_sql(bytes: ::RawValue<'_>) -> diesel::deserialize::Result { + let s = String::from_sql(bytes)?; + + Self::from_base64(s).ok_or_else(|| format!("Invalid base64 hash").into()) + } +} + impl Hash { pub(crate) fn new(hash: [u8; 32], size: u64, format: InternalFormat) -> Self { Self { @@ -30,6 +59,22 @@ impl Hash { hex::encode(self.to_bytes()) } + pub(crate) fn to_base64(&self) -> String { + use base64::Engine; + + base64::engine::general_purpose::STANDARD.encode(self.to_bytes()) + } + + pub(crate) fn from_base64(input: String) -> Option { + use base64::Engine; + + let bytes = base64::engine::general_purpose::STANDARD + .decode(input) + .ok()?; + + Self::from_bytes(&bytes) + } + pub(super) fn to_bytes(&self) -> Vec { let format_byte = self.format.to_byte(); diff --git a/src/repo/migrate.rs b/src/repo/migrate.rs index d27dd24..60dbe68 100644 --- a/src/repo/migrate.rs +++ b/src/repo/migrate.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use tokio::task::JoinSet; use crate::{ @@ -33,7 +35,7 @@ pub(crate) async fn migrate_repo(old_repo: ArcRepo, new_repo: ArcRepo) -> Result tracing::warn!("Checks complete, migrating repo"); tracing::warn!("{total_size} hashes will be migrated"); - let mut hash_stream = old_repo.hashes().await.into_streamer(); + let mut hash_stream = old_repo.hashes().into_streamer(); let mut index = 0; while let Some(res) = hash_stream.next().await { @@ -266,7 +268,7 @@ async fn do_migrate_hash_04( config: &Configuration, old_hash: sled::IVec, ) -> Result<(), Error> { - let Some(identifier) = old_repo.identifier::(old_hash.clone()).await? else { + let Some(identifier) = old_repo.identifier(old_hash.clone()).await? else { tracing::warn!("Skipping hash {}, no identifier", hex::encode(&old_hash)); return Ok(()); }; @@ -276,10 +278,8 @@ async fn do_migrate_hash_04( let hash_details = set_details(old_repo, new_repo, store, config, &identifier).await?; let aliases = old_repo.aliases_for_hash(old_hash.clone()).await?; - let variants = old_repo.variants::(old_hash.clone()).await?; - let motion_identifier = old_repo - .motion_identifier::(old_hash.clone()) - .await?; + let variants = old_repo.variants(old_hash.clone()).await?; + let motion_identifier = old_repo.motion_identifier(old_hash.clone()).await?; let hash = old_hash[..].try_into().expect("Invalid hash size"); @@ -326,7 +326,7 @@ async fn set_details( new_repo: &ArcRepo, store: &S, config: &Configuration, - identifier: &S::Identifier, + identifier: &Arc, ) -> Result { if let Some(details) = new_repo.details(identifier).await? { Ok(details) @@ -342,9 +342,9 @@ async fn fetch_or_generate_details( old_repo: &OldSledRepo, store: &S, config: &Configuration, - identifier: &S::Identifier, + identifier: &Arc, ) -> Result { - let details_opt = old_repo.details(identifier).await?; + let details_opt = old_repo.details(identifier.clone()).await?; if let Some(details) = details_opt { Ok(details) diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index c66303c..9ffd3f4 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -1,6 +1,8 @@ mod embedded; mod schema; +use std::sync::Arc; + use diesel::prelude::*; use diesel_async::{ pooled_connection::{ @@ -11,7 +13,12 @@ use diesel_async::{ }; use url::Url; -use super::{BaseRepo, HashRepo, RepoError}; +use crate::error_code::ErrorCode; + +use super::{ + BaseRepo, Hash, HashAlreadyExists, HashPage, HashRepo, OrderedHash, RepoError, + VariantAlreadyExists, +}; #[derive(Clone)] pub(crate) struct PostgresRepo { @@ -39,6 +46,12 @@ pub(crate) enum PostgresError { Diesel(#[source] diesel::result::Error), } +impl PostgresError { + pub(super) const fn error_code(&self) -> ErrorCode { + todo!() + } +} + impl PostgresRepo { pub(crate) async fn connect(postgres_url: Url) -> Result { let (mut client, conn) = @@ -65,6 +78,11 @@ impl PostgresRepo { } } +fn to_primitive(timestamp: time::OffsetDateTime) -> time::PrimitiveDateTime { + let timestamp = timestamp.to_offset(time::UtcOffset::UTC); + time::PrimitiveDateTime::new(timestamp.date(), timestamp.time()) +} + impl BaseRepo for PostgresRepo {} #[async_trait::async_trait(?Send)] @@ -82,6 +100,187 @@ impl HashRepo for PostgresRepo { Ok(count.try_into().expect("non-negative count")) } + + async fn bound(&self, input_hash: Hash) -> Result, RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let timestamp = hashes + .select(created_at) + .filter(hash.eq(&input_hash)) + .first(&mut conn) + .await + .map(time::PrimitiveDateTime::assume_utc) + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(timestamp.map(|timestamp| OrderedHash { + timestamp, + hash: input_hash, + })) + } + + async fn hash_page_by_date( + &self, + date: time::OffsetDateTime, + limit: usize, + ) -> Result { + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let timestamp = to_primitive(date); + + let ordered_hash = hashes + .select((created_at, hash)) + .filter(created_at.lt(timestamp)) + .order(created_at.desc()) + .first::<(time::PrimitiveDateTime, Hash)>(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .map(|tup| OrderedHash { + timestamp: tup.0.assume_utc(), + hash: tup.1, + }); + + self.hashes_ordered(ordered_hash, limit).await + } + + async fn hashes_ordered( + &self, + bound: Option, + limit: usize, + ) -> Result { + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let (mut page, prev) = if let Some(OrderedHash { + timestamp, + hash: bound_hash, + }) = bound + { + let timestamp = to_primitive(timestamp); + + let page = hashes + .select(hash) + .filter(created_at.lt(timestamp)) + .or_filter(created_at.eq(timestamp).and(hash.le(&bound_hash))) + .order(created_at.desc()) + .then_order_by(hash.desc()) + .limit(limit as i64 + 1) + .load::(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + let prev = hashes + .select(hash) + .filter(created_at.gt(timestamp)) + .or_filter(created_at.eq(timestamp).and(hash.gt(&bound_hash))) + .order(created_at) + .then_order_by(hash) + .offset(limit.saturating_sub(1) as i64) + .first::(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)?; + + (page, prev) + } else { + let page = hashes + .select(hash) + .order(created_at.desc()) + .then_order_by(hash.desc()) + .limit(limit as i64 + 1) + .load::(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + (page, None) + }; + + let next = if page.len() > limit { page.pop() } else { None }; + + Ok(HashPage { + limit, + prev, + next, + hashes: page, + }) + } + + async fn create_hash_with_timestamp( + &self, + input_hash: Hash, + input_identifier: &Arc, + timestamp: time::OffsetDateTime, + ) -> Result, RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let timestamp = to_primitive(timestamp); + + /* + insert_into(hashes).values(( + hash.eq(&input_hash), + identifier.eq(&input_identifier) + )) + */ + + todo!() + } + + async fn update_identifier(&self, hash: Hash, identifier: &Arc) -> Result<(), RepoError> { + todo!() + } + + async fn identifier(&self, hash: Hash) -> Result>, RepoError> { + todo!() + } + + async fn relate_variant_identifier( + &self, + hash: Hash, + variant: String, + identifier: &Arc, + ) -> Result, RepoError> { + todo!() + } + + async fn variant_identifier( + &self, + hash: Hash, + variant: String, + ) -> Result>, RepoError> { + todo!() + } + + async fn variants(&self, hash: Hash) -> Result)>, RepoError> { + todo!() + } + + async fn remove_variant(&self, hash: Hash, variant: String) -> Result<(), RepoError> { + todo!() + } + + async fn relate_motion_identifier( + &self, + hash: Hash, + identifier: &Arc, + ) -> Result<(), RepoError> { + todo!() + } + + async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError> { + todo!() + } + + async fn cleanup_hash(&self, hash: Hash) -> Result<(), RepoError> { + todo!() + } } impl std::fmt::Debug for PostgresRepo { diff --git a/src/repo/postgres/migrations/V0001__create_hashes.rs b/src/repo/postgres/migrations/V0001__create_hashes.rs index 9829585..1c97f08 100644 --- a/src/repo/postgres/migrations/V0001__create_hashes.rs +++ b/src/repo/postgres/migrations/V0001__create_hashes.rs @@ -8,7 +8,7 @@ pub(crate) fn migration() -> String { m.create_table("hashes", |t| { t.add_column( "hash", - types::binary() + types::text() .primary(true) .unique(true) .nullable(false) diff --git a/src/repo/postgres/migrations/V0002__create_variants.rs b/src/repo/postgres/migrations/V0002__create_variants.rs index c49c62c..4e62f1b 100644 --- a/src/repo/postgres/migrations/V0002__create_variants.rs +++ b/src/repo/postgres/migrations/V0002__create_variants.rs @@ -7,7 +7,7 @@ pub(crate) fn migration() -> String { m.create_table("variants", |t| { t.inject_custom(r#""id" UUID PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL UNIQUE"#); - t.add_column("hash", types::binary().nullable(false)); + t.add_column("hash", types::text().nullable(false)); t.add_column("variant", types::text().nullable(false)); t.add_column("identifier", types::text().nullable(false)); t.add_column( diff --git a/src/repo/postgres/migrations/V0003__create_aliases.rs b/src/repo/postgres/migrations/V0003__create_aliases.rs index a8aa37f..007f123 100644 --- a/src/repo/postgres/migrations/V0003__create_aliases.rs +++ b/src/repo/postgres/migrations/V0003__create_aliases.rs @@ -14,7 +14,7 @@ pub(crate) fn migration() -> String { .unique(true) .nullable(false), ); - t.add_column("hash", types::binary().nullable(false)); + t.add_column("hash", types::text().nullable(false)); t.add_column("token", types::text().size(60).nullable(false)); t.add_foreign_key(&["hash"], "hashes", &["hash"]); diff --git a/src/repo/postgres/migrations/V0005__create_details.rs b/src/repo/postgres/migrations/V0005__create_details.rs index e324caa..440aa02 100644 --- a/src/repo/postgres/migrations/V0005__create_details.rs +++ b/src/repo/postgres/migrations/V0005__create_details.rs @@ -10,7 +10,7 @@ pub(crate) fn migration() -> String { "identifier", types::text().primary(true).unique(true).nullable(false), ); - t.add_column("details", types::custom("jsonb").nullable(false)); + t.add_column("json", types::custom("jsonb").nullable(false)); }); m.make::().to_string() diff --git a/src/repo/postgres/migrations/V0006__create_queue.rs b/src/repo/postgres/migrations/V0006__create_queue.rs index 293ff97..d5161bf 100644 --- a/src/repo/postgres/migrations/V0006__create_queue.rs +++ b/src/repo/postgres/migrations/V0006__create_queue.rs @@ -7,7 +7,7 @@ pub(crate) fn migration() -> String { m.inject_custom("CREATE TYPE job_status AS ENUM ('new', 'running');"); - m.create_table("queue", |t| { + m.create_table("job_queue", |t| { t.inject_custom(r#""id" UUID PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL UNIQUE"#); t.add_column("queue", types::text().size(50).nullable(false)); t.add_column("job", types::custom("jsonb").nullable(false)); @@ -42,7 +42,7 @@ $$ LANGUAGE plpgsql; r#" CREATE TRIGGER queue_status AFTER INSERT OR UPDATE OF status - ON queue + ON job_queue FOR EACH ROW EXECUTE PROCEDURE queue_status_notify(); "# diff --git a/src/repo/postgres/schema.rs b/src/repo/postgres/schema.rs index 3ab3e46..1b6aba4 100644 --- a/src/repo/postgres/schema.rs +++ b/src/repo/postgres/schema.rs @@ -9,7 +9,7 @@ pub mod sql_types { diesel::table! { aliases (alias) { alias -> Text, - hash -> Bytea, + hash -> Text, token -> Text, } } @@ -17,20 +17,33 @@ diesel::table! { diesel::table! { details (identifier) { identifier -> Text, - #[sql_name = "details"] - details_json -> Jsonb, + json -> Jsonb, } } diesel::table! { hashes (hash) { - hash -> Bytea, + hash -> Text, identifier -> Text, motion_identifier -> Nullable, created_at -> Timestamp, } } +diesel::table! { + use diesel::sql_types::*; + use super::sql_types::JobStatus; + + job_queue (id) { + id -> Uuid, + queue -> Text, + job -> Jsonb, + status -> JobStatus, + queue_time -> Timestamp, + heartbeat -> Timestamp, + } +} + diesel::table! { proxies (url) { url -> Text, @@ -39,21 +52,6 @@ diesel::table! { } } -diesel::table! { - use diesel::sql_types::*; - use super::sql_types::JobStatus; - - queue (id) { - id -> Uuid, - #[sql_name = "queue"] - queue_name -> Text, - job -> Jsonb, - status -> JobStatus, - queue_time -> Timestamp, - heartbeat -> Timestamp, - } -} - diesel::table! { refinery_schema_history (version) { version -> Int4, @@ -89,7 +87,7 @@ diesel::table! { diesel::table! { variants (id) { id -> Uuid, - hash -> Bytea, + hash -> Text, variant -> Text, identifier -> Text, accessed -> Timestamp, @@ -104,8 +102,8 @@ diesel::allow_tables_to_appear_in_same_query!( aliases, details, hashes, + job_queue, proxies, - queue, refinery_schema_history, settings, store_migrations, diff --git a/src/repo/sled.rs b/src/repo/sled.rs index 57527c2..33cf611 100644 --- a/src/repo/sled.rs +++ b/src/repo/sled.rs @@ -2,7 +2,6 @@ use crate::{ details::HumanDate, error_code::{ErrorCode, OwnedErrorCode}, serde_str::Serde, - store::StoreError, stream::{from_iterator, LocalBoxStream}, }; use sled::{transaction::TransactionError, Db, IVec, Transactional, Tree}; @@ -21,9 +20,9 @@ use uuid::Uuid; use super::{ hash::Hash, Alias, AliasAccessRepo, AliasAlreadyExists, AliasRepo, BaseRepo, DeleteToken, - Details, DetailsRepo, FullRepo, HashAlreadyExists, HashPage, HashRepo, Identifier, JobId, - OrderedHash, ProxyRepo, QueueRepo, RepoError, SettingsRepo, StoreMigrationRepo, UploadId, - UploadRepo, UploadResult, VariantAccessRepo, VariantAlreadyExists, + Details, DetailsRepo, FullRepo, HashAlreadyExists, HashPage, HashRepo, JobId, OrderedHash, + ProxyRepo, QueueRepo, RepoError, SettingsRepo, StoreMigrationRepo, UploadId, UploadRepo, + UploadResult, VariantAccessRepo, VariantAlreadyExists, }; macro_rules! b { @@ -55,6 +54,9 @@ pub(crate) enum SledError { #[error("Error parsing variant key")] VariantKey(#[from] VariantKeyError), + #[error("Invalid string data in db")] + Utf8(#[source] std::str::Utf8Error), + #[error("Operation panicked")] Panic, @@ -65,7 +67,7 @@ pub(crate) enum SledError { impl SledError { pub(super) const fn error_code(&self) -> ErrorCode { match self { - Self::Sled(_) | Self::VariantKey(_) => ErrorCode::SLED_ERROR, + Self::Sled(_) | Self::VariantKey(_) | Self::Utf8(_) => ErrorCode::SLED_ERROR, Self::Details(_) => ErrorCode::EXTRACT_DETAILS, Self::UploadResult(_) => ErrorCode::EXTRACT_UPLOAD_RESULT, Self::Panic => ErrorCode::PANIC, @@ -648,10 +650,17 @@ fn job_key(queue: &'static str, job_id: JobId) -> Arc<[u8]> { Arc::from(key) } +fn try_into_arc_str(ivec: IVec) -> Result, SledError> { + std::str::from_utf8(&ivec[..]) + .map_err(SledError::Utf8) + .map(String::from) + .map(Arc::from) +} + #[async_trait::async_trait(?Send)] impl QueueRepo for SledRepo { - #[tracing::instrument(skip(self, job), fields(job = %String::from_utf8_lossy(&job)))] - async fn push(&self, queue_name: &'static str, job: Arc<[u8]>) -> Result { + #[tracing::instrument(skip(self))] + async fn push(&self, queue_name: &'static str, job: Arc) -> Result { let metrics_guard = PushMetricsGuard::guard(queue_name); let id = JobId::gen(); @@ -700,7 +709,7 @@ impl QueueRepo for SledRepo { &self, queue_name: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc<[u8]>), RepoError> { + ) -> Result<(JobId, Arc), RepoError> { let metrics_guard = PopMetricsGuard::guard(queue_name); let now = time::OffsetDateTime::now_utc(); @@ -753,11 +762,10 @@ impl QueueRepo for SledRepo { tracing::Span::current().record("job_id", &format!("{job_id:?}")); - let opt = queue - .get(&key)? - .map(|job_bytes| (job_id, Arc::from(job_bytes.to_vec()))); + let opt = queue.get(&key)?.map(try_into_arc_str).transpose()?; - return Ok(opt) as Result)>, SledError>; + return Ok(opt.map(|job| (job_id, job))) + as Result)>, SledError>; } Ok(None) @@ -949,43 +957,46 @@ fn variant_from_key(hash: &[u8], key: &[u8]) -> Option { #[async_trait::async_trait(?Send)] impl DetailsRepo for SledRepo { - #[tracing::instrument(level = "trace", skip(self, identifier), fields(identifier = identifier.string_repr()))] + #[tracing::instrument(level = "trace", skip(self))] async fn relate_details( &self, - identifier: &dyn Identifier, + identifier: &Arc, details: &Details, - ) -> Result<(), StoreError> { - let key = identifier.to_bytes()?; - let details = serde_json::to_vec(&details.inner) - .map_err(SledError::Details) - .map_err(RepoError::from)?; + ) -> Result<(), RepoError> { + let key = identifier.clone(); + let details = serde_json::to_vec(&details.inner).map_err(SledError::Details)?; b!( self.identifier_details, - identifier_details.insert(key, details) + identifier_details.insert(key.as_bytes(), details) ); Ok(()) } - #[tracing::instrument(level = "trace", skip(self, identifier), fields(identifier = identifier.string_repr()))] - async fn details(&self, identifier: &dyn Identifier) -> Result, StoreError> { - let key = identifier.to_bytes()?; + #[tracing::instrument(level = "trace", skip(self))] + async fn details(&self, identifier: &Arc) -> Result, RepoError> { + let key = identifier.clone(); - let opt = b!(self.identifier_details, identifier_details.get(key)); + let opt = b!( + self.identifier_details, + identifier_details.get(key.as_bytes()) + ); opt.map(|ivec| serde_json::from_slice(&ivec).map(|inner| Details { inner })) .transpose() .map_err(SledError::Details) .map_err(RepoError::from) - .map_err(StoreError::from) } - #[tracing::instrument(level = "trace", skip(self, identifier), fields(identifier = identifier.string_repr()))] - async fn cleanup_details(&self, identifier: &dyn Identifier) -> Result<(), StoreError> { - let key = identifier.to_bytes()?; + #[tracing::instrument(level = "trace", skip(self))] + async fn cleanup_details(&self, identifier: &Arc) -> Result<(), RepoError> { + let key = identifier.clone(); - b!(self.identifier_details, identifier_details.remove(key)); + b!( + self.identifier_details, + identifier_details.remove(key.as_bytes()) + ); Ok(()) } @@ -999,24 +1010,28 @@ impl StoreMigrationRepo for SledRepo { async fn mark_migrated( &self, - old_identifier: &dyn Identifier, - new_identifier: &dyn Identifier, - ) -> Result<(), StoreError> { - let key = new_identifier.to_bytes()?; - let value = old_identifier.to_bytes()?; + old_identifier: &Arc, + new_identifier: &Arc, + ) -> Result<(), RepoError> { + let key = new_identifier.clone(); + let value = old_identifier.clone(); b!( self.migration_identifiers, - migration_identifiers.insert(key, value) + migration_identifiers.insert(key.as_bytes(), value.as_bytes()) ); Ok(()) } - async fn is_migrated(&self, identifier: &dyn Identifier) -> Result { - let key = identifier.to_bytes()?; + async fn is_migrated(&self, identifier: &Arc) -> Result { + let key = identifier.clone(); - Ok(b!(self.migration_identifiers, migration_identifiers.get(key)).is_some()) + Ok(b!( + self.migration_identifiers, + migration_identifiers.get(key.as_bytes()) + ) + .is_some()) } async fn clear(&self) -> Result<(), RepoError> { @@ -1062,17 +1077,6 @@ impl HashRepo for SledRepo { )) } - async fn hashes(&self) -> LocalBoxStream<'static, Result> { - let iter = self.hashes.iter().keys().filter_map(|res| { - res.map_err(SledError::from) - .map_err(RepoError::from) - .map(Hash::from_ivec) - .transpose() - }); - - Box::pin(from_iterator(iter, 8)) - } - async fn bound(&self, hash: Hash) -> Result, RepoError> { let opt = b!(self.hashes, hashes.get(hash.to_ivec())); @@ -1197,10 +1201,10 @@ impl HashRepo for SledRepo { async fn create_hash_with_timestamp( &self, hash: Hash, - identifier: &dyn Identifier, + identifier: &Arc, timestamp: time::OffsetDateTime, - ) -> Result, StoreError> { - let identifier: sled::IVec = identifier.to_bytes()?.into(); + ) -> Result, RepoError> { + let identifier: sled::IVec = identifier.as_bytes().to_vec().into(); let hashes = self.hashes.clone(); let hashes_inverse = self.hashes_inverse.clone(); @@ -1234,63 +1238,56 @@ impl HashRepo for SledRepo { match res { Ok(res) => Ok(res), Err(TransactionError::Abort(e) | TransactionError::Storage(e)) => { - Err(StoreError::from(RepoError::from(SledError::from(e)))) + Err(RepoError::from(SledError::from(e))) } } } - async fn update_identifier( - &self, - hash: Hash, - identifier: &dyn Identifier, - ) -> Result<(), StoreError> { - let identifier = identifier.to_bytes()?; + async fn update_identifier(&self, hash: Hash, identifier: &Arc) -> Result<(), RepoError> { + let identifier = identifier.clone(); let hash = hash.to_ivec(); b!( self.hash_identifiers, - hash_identifiers.insert(hash, identifier) + hash_identifiers.insert(hash, identifier.as_bytes()) ); Ok(()) } #[tracing::instrument(level = "trace", skip(self))] - async fn identifier(&self, hash: Hash) -> Result>, RepoError> { + async fn identifier(&self, hash: Hash) -> Result>, RepoError> { let hash = hash.to_ivec(); - let Some(ivec) = b!(self.hash_identifiers, hash_identifiers.get(hash)) else { - return Ok(None); - }; + let opt = b!(self.hash_identifiers, hash_identifiers.get(hash)); - Ok(Some(Arc::from(ivec.to_vec()))) + Ok(opt.map(try_into_arc_str).transpose()?) } - #[tracing::instrument(level = "trace", skip(self, identifier), fields(identifier = identifier.string_repr()))] + #[tracing::instrument(level = "trace", skip(self))] async fn relate_variant_identifier( &self, hash: Hash, variant: String, - identifier: &dyn Identifier, - ) -> Result, StoreError> { + identifier: &Arc, + ) -> Result, RepoError> { let hash = hash.to_bytes(); let key = variant_key(&hash, &variant); - let value = identifier.to_bytes()?; + let value = identifier.clone(); let hash_variant_identifiers = self.hash_variant_identifiers.clone(); actix_rt::task::spawn_blocking(move || { hash_variant_identifiers - .compare_and_swap(key, Option::<&[u8]>::None, Some(value)) + .compare_and_swap(key, Option::<&[u8]>::None, Some(value.as_bytes())) .map(|res| res.map_err(|_| VariantAlreadyExists)) }) .await .map_err(|_| RepoError::Canceled)? .map_err(SledError::from) .map_err(RepoError::from) - .map_err(StoreError::from) } #[tracing::instrument(level = "trace", skip(self))] @@ -1298,7 +1295,7 @@ impl HashRepo for SledRepo { &self, hash: Hash, variant: String, - ) -> Result>, RepoError> { + ) -> Result>, RepoError> { let hash = hash.to_bytes(); let key = variant_key(&hash, &variant); @@ -1308,11 +1305,11 @@ impl HashRepo for SledRepo { hash_variant_identifiers.get(key) ); - Ok(opt.map(|ivec| Arc::from(ivec.to_vec()))) + Ok(opt.map(try_into_arc_str).transpose()?) } #[tracing::instrument(level = "debug", skip(self))] - async fn variants(&self, hash: Hash) -> Result)>, RepoError> { + async fn variants(&self, hash: Hash) -> Result)>, RepoError> { let hash = hash.to_ivec(); let vec = b!( @@ -1321,14 +1318,14 @@ impl HashRepo for SledRepo { .scan_prefix(hash.clone()) .filter_map(|res| res.ok()) .filter_map(|(key, ivec)| { - let identifier = Arc::from(ivec.to_vec()); + let identifier = try_into_arc_str(ivec).ok(); let variant = variant_from_key(&hash, &key); if variant.is_none() { tracing::warn!("Skipping a variant: {}", String::from_utf8_lossy(&key)); } - Some((variant?, identifier)) + Some((variant?, identifier?)) }) .collect::>()) as Result, SledError> ); @@ -1350,25 +1347,25 @@ impl HashRepo for SledRepo { Ok(()) } - #[tracing::instrument(level = "trace", skip(self, identifier), fields(identifier = identifier.string_repr()))] + #[tracing::instrument(level = "trace", skip(self))] async fn relate_motion_identifier( &self, hash: Hash, - identifier: &dyn Identifier, - ) -> Result<(), StoreError> { + identifier: &Arc, + ) -> Result<(), RepoError> { let hash = hash.to_ivec(); - let bytes = identifier.to_bytes()?; + let bytes = identifier.clone(); b!( self.hash_motion_identifiers, - hash_motion_identifiers.insert(hash, bytes) + hash_motion_identifiers.insert(hash, bytes.as_bytes()) ); Ok(()) } #[tracing::instrument(level = "trace", skip(self))] - async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError> { + async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError> { let hash = hash.to_ivec(); let opt = b!( @@ -1376,7 +1373,7 @@ impl HashRepo for SledRepo { hash_motion_identifiers.get(hash) ); - Ok(opt.map(|ivec| Arc::from(ivec.to_vec()))) + Ok(opt.map(try_into_arc_str).transpose()?) } #[tracing::instrument(skip(self))] diff --git a/src/repo_04.rs b/src/repo_04.rs index 174882a..0d63d55 100644 --- a/src/repo_04.rs +++ b/src/repo_04.rs @@ -2,10 +2,9 @@ use crate::{ config, details::Details, repo::{Alias, DeleteToken}, - store::{Identifier, StoreError}, }; use futures_core::Stream; -use std::fmt::Debug; +use std::{fmt::Debug, sync::Arc}; pub(crate) use self::sled::SledRepo; @@ -46,7 +45,7 @@ pub(crate) trait SettingsRepo: BaseRepo { #[async_trait::async_trait(?Send)] pub(crate) trait IdentifierRepo: BaseRepo { - async fn details(&self, identifier: &I) -> Result, StoreError>; + async fn details(&self, identifier: Arc) -> Result, RepoError>; } #[async_trait::async_trait(?Send)] @@ -57,20 +56,11 @@ pub(crate) trait HashRepo: BaseRepo { async fn hashes(&self) -> Self::Stream; - async fn identifier( - &self, - hash: Self::Bytes, - ) -> Result, StoreError>; + async fn identifier(&self, hash: Self::Bytes) -> Result>, RepoError>; - async fn variants( - &self, - hash: Self::Bytes, - ) -> Result, StoreError>; + async fn variants(&self, hash: Self::Bytes) -> Result)>, RepoError>; - async fn motion_identifier( - &self, - hash: Self::Bytes, - ) -> Result, StoreError>; + async fn motion_identifier(&self, hash: Self::Bytes) -> Result>, RepoError>; } #[async_trait::async_trait(?Send)] diff --git a/src/repo_04/sled.rs b/src/repo_04/sled.rs index 36595b7..09a684a 100644 --- a/src/repo_04/sled.rs +++ b/src/repo_04/sled.rs @@ -1,10 +1,9 @@ use crate::{ details::HumanDate, repo_04::{ - Alias, AliasRepo, BaseRepo, DeleteToken, Details, HashRepo, Identifier, IdentifierRepo, - RepoError, SettingsRepo, + Alias, AliasRepo, BaseRepo, DeleteToken, Details, HashRepo, IdentifierRepo, RepoError, + SettingsRepo, }, - store::StoreError, stream::{from_iterator, LocalBoxStream}, }; use sled::{Db, IVec, Tree}; @@ -56,6 +55,9 @@ pub(crate) enum SledError { #[error("Operation panicked")] Panic, + + #[error("Error reading string")] + Utf8(#[from] std::str::Utf8Error), } #[derive(Clone)] @@ -179,17 +181,17 @@ fn variant_from_key(hash: &[u8], key: &[u8]) -> Option { #[async_trait::async_trait(?Send)] impl IdentifierRepo for SledRepo { - #[tracing::instrument(level = "trace", skip(self, identifier), fields(identifier = identifier.string_repr()))] - async fn details(&self, identifier: &I) -> Result, StoreError> { - let key = identifier.to_bytes()?; - - let opt = b!(self.identifier_details, identifier_details.get(key)); + #[tracing::instrument(level = "trace", skip(self))] + async fn details(&self, key: Arc) -> Result, RepoError> { + let opt = b!( + self.identifier_details, + identifier_details.get(key.as_bytes()) + ); opt.map(|ivec| serde_json::from_slice::(&ivec)) .transpose() .map_err(SledError::from) .map_err(RepoError::from) - .map_err(StoreError::from) .map(|opt| opt.and_then(OldDetails::into_details)) } } @@ -219,29 +221,27 @@ impl HashRepo for SledRepo { } #[tracing::instrument(level = "trace", skip(self, hash), fields(hash = hex::encode(&hash)))] - async fn identifier( - &self, - hash: Self::Bytes, - ) -> Result, StoreError> { + async fn identifier(&self, hash: Self::Bytes) -> Result>, RepoError> { let Some(ivec) = b!(self.hash_identifiers, hash_identifiers.get(hash)) else { return Ok(None); }; - Ok(Some(I::from_bytes(ivec.to_vec())?)) + Ok(Some(Arc::from( + std::str::from_utf8(&ivec[..]) + .map_err(SledError::from)? + .to_string(), + ))) } #[tracing::instrument(level = "debug", skip(self, hash), fields(hash = hex::encode(&hash)))] - async fn variants( - &self, - hash: Self::Bytes, - ) -> Result, StoreError> { + async fn variants(&self, hash: Self::Bytes) -> Result)>, RepoError> { let vec = b!( self.hash_variant_identifiers, Ok(hash_variant_identifiers .scan_prefix(&hash) .filter_map(|res| res.ok()) .filter_map(|(key, ivec)| { - let identifier = I::from_bytes(ivec.to_vec()).ok(); + let identifier = String::from_utf8(ivec.to_vec()).ok(); if identifier.is_none() { tracing::warn!( "Skipping an identifier: {}", @@ -254,7 +254,7 @@ impl HashRepo for SledRepo { tracing::warn!("Skipping a variant: {}", String::from_utf8_lossy(&key)); } - Some((variant?, identifier?)) + Some((variant?, Arc::from(identifier?))) }) .collect::>()) as Result, SledError> ); @@ -263,16 +263,20 @@ impl HashRepo for SledRepo { } #[tracing::instrument(level = "trace", skip(self, hash), fields(hash = hex::encode(&hash)))] - async fn motion_identifier( - &self, - hash: Self::Bytes, - ) -> Result, StoreError> { + async fn motion_identifier(&self, hash: Self::Bytes) -> Result>, RepoError> { let opt = b!( self.hash_motion_identifiers, hash_motion_identifiers.get(hash) ); - opt.map(|ivec| I::from_bytes(ivec.to_vec())).transpose() + opt.map(|ivec| { + Ok(Arc::from( + std::str::from_utf8(&ivec[..]) + .map_err(SledError::from)? + .to_string(), + )) + }) + .transpose() } } diff --git a/src/store.rs b/src/store.rs index 66c9337..3293237 100644 --- a/src/store.rs +++ b/src/store.rs @@ -1,10 +1,9 @@ use actix_web::web::Bytes; -use base64::{prelude::BASE64_STANDARD, Engine}; use futures_core::Stream; use std::{fmt::Debug, sync::Arc}; use tokio::io::{AsyncRead, AsyncWrite}; -use crate::error_code::ErrorCode; +use crate::{error_code::ErrorCode, stream::LocalBoxStream}; pub(crate) mod file_store; pub(crate) mod object_store; @@ -70,32 +69,15 @@ impl From for StoreError { } } -pub(crate) trait Identifier: Send + Sync + Debug { - fn to_bytes(&self) -> Result, StoreError>; - - fn from_bytes(bytes: Vec) -> Result - where - Self: Sized; - - fn from_arc(arc: Arc<[u8]>) -> Result - where - Self: Sized; - - fn string_repr(&self) -> String; -} - #[async_trait::async_trait(?Send)] pub(crate) trait Store: Clone + Debug { - type Identifier: Identifier + Clone + 'static; - type Stream: Stream> + Unpin + 'static; - async fn health_check(&self) -> Result<(), StoreError>; async fn save_async_read( &self, reader: Reader, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static; @@ -103,7 +85,7 @@ pub(crate) trait Store: Clone + Debug { &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static; @@ -111,28 +93,28 @@ pub(crate) trait Store: Clone + Debug { &self, bytes: Bytes, content_type: mime::Mime, - ) -> Result; + ) -> Result, StoreError>; - fn public_url(&self, _: &Self::Identifier) -> Option; + fn public_url(&self, _: &Arc) -> Option; async fn to_stream( &self, - identifier: &Self::Identifier, + identifier: &Arc, from_start: Option, len: Option, - ) -> Result; + ) -> Result>, StoreError>; async fn read_into( &self, - identifier: &Self::Identifier, + identifier: &Arc, writer: &mut Writer, ) -> Result<(), std::io::Error> where Writer: AsyncWrite + Unpin; - async fn len(&self, identifier: &Self::Identifier) -> Result; + async fn len(&self, identifier: &Arc) -> Result; - async fn remove(&self, identifier: &Self::Identifier) -> Result<(), StoreError>; + async fn remove(&self, identifier: &Arc) -> Result<(), StoreError>; } #[async_trait::async_trait(?Send)] @@ -140,9 +122,6 @@ impl Store for actix_web::web::Data where T: Store, { - type Identifier = T::Identifier; - type Stream = T::Stream; - async fn health_check(&self) -> Result<(), StoreError> { T::health_check(self).await } @@ -151,7 +130,7 @@ where &self, reader: Reader, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static, { @@ -162,7 +141,7 @@ where &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static, { @@ -173,26 +152,26 @@ where &self, bytes: Bytes, content_type: mime::Mime, - ) -> Result { + ) -> Result, StoreError> { T::save_bytes(self, bytes, content_type).await } - fn public_url(&self, identifier: &Self::Identifier) -> Option { + fn public_url(&self, identifier: &Arc) -> Option { T::public_url(self, identifier) } async fn to_stream( &self, - identifier: &Self::Identifier, + identifier: &Arc, from_start: Option, len: Option, - ) -> Result { + ) -> Result>, StoreError> { T::to_stream(self, identifier, from_start, len).await } async fn read_into( &self, - identifier: &Self::Identifier, + identifier: &Arc, writer: &mut Writer, ) -> Result<(), std::io::Error> where @@ -201,11 +180,83 @@ where T::read_into(self, identifier, writer).await } - async fn len(&self, identifier: &Self::Identifier) -> Result { + async fn len(&self, identifier: &Arc) -> Result { T::len(self, identifier).await } - async fn remove(&self, identifier: &Self::Identifier) -> Result<(), StoreError> { + async fn remove(&self, identifier: &Arc) -> Result<(), StoreError> { + T::remove(self, identifier).await + } +} + +#[async_trait::async_trait(?Send)] +impl Store for Arc +where + T: Store, +{ + async fn health_check(&self) -> Result<(), StoreError> { + T::health_check(self).await + } + + async fn save_async_read( + &self, + reader: Reader, + content_type: mime::Mime, + ) -> Result, StoreError> + where + Reader: AsyncRead + Unpin + 'static, + { + T::save_async_read(self, reader, content_type).await + } + + async fn save_stream( + &self, + stream: S, + content_type: mime::Mime, + ) -> Result, StoreError> + where + S: Stream> + Unpin + 'static, + { + T::save_stream(self, stream, content_type).await + } + + async fn save_bytes( + &self, + bytes: Bytes, + content_type: mime::Mime, + ) -> Result, StoreError> { + T::save_bytes(self, bytes, content_type).await + } + + fn public_url(&self, identifier: &Arc) -> Option { + T::public_url(self, identifier) + } + + async fn to_stream( + &self, + identifier: &Arc, + from_start: Option, + len: Option, + ) -> Result>, StoreError> { + T::to_stream(self, identifier, from_start, len).await + } + + async fn read_into( + &self, + identifier: &Arc, + writer: &mut Writer, + ) -> Result<(), std::io::Error> + where + Writer: AsyncWrite + Unpin, + { + T::read_into(self, identifier, writer).await + } + + async fn len(&self, identifier: &Arc) -> Result { + T::len(self, identifier).await + } + + async fn remove(&self, identifier: &Arc) -> Result<(), StoreError> { T::remove(self, identifier).await } } @@ -215,9 +266,6 @@ impl<'a, T> Store for &'a T where T: Store, { - type Identifier = T::Identifier; - type Stream = T::Stream; - async fn health_check(&self) -> Result<(), StoreError> { T::health_check(self).await } @@ -226,7 +274,7 @@ where &self, reader: Reader, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static, { @@ -237,7 +285,7 @@ where &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static, { @@ -248,26 +296,26 @@ where &self, bytes: Bytes, content_type: mime::Mime, - ) -> Result { + ) -> Result, StoreError> { T::save_bytes(self, bytes, content_type).await } - fn public_url(&self, identifier: &Self::Identifier) -> Option { + fn public_url(&self, identifier: &Arc) -> Option { T::public_url(self, identifier) } async fn to_stream( &self, - identifier: &Self::Identifier, + identifier: &Arc, from_start: Option, len: Option, - ) -> Result { + ) -> Result>, StoreError> { T::to_stream(self, identifier, from_start, len).await } async fn read_into( &self, - identifier: &Self::Identifier, + identifier: &Arc, writer: &mut Writer, ) -> Result<(), std::io::Error> where @@ -276,59 +324,11 @@ where T::read_into(self, identifier, writer).await } - async fn len(&self, identifier: &Self::Identifier) -> Result { + async fn len(&self, identifier: &Arc) -> Result { T::len(self, identifier).await } - async fn remove(&self, identifier: &Self::Identifier) -> Result<(), StoreError> { + async fn remove(&self, identifier: &Arc) -> Result<(), StoreError> { T::remove(self, identifier).await } } - -impl Identifier for Vec { - fn from_bytes(bytes: Vec) -> Result - where - Self: Sized, - { - Ok(bytes) - } - - fn from_arc(arc: Arc<[u8]>) -> Result - where - Self: Sized, - { - Ok(Vec::from(&arc[..])) - } - - fn to_bytes(&self) -> Result, StoreError> { - Ok(self.clone()) - } - - fn string_repr(&self) -> String { - BASE64_STANDARD.encode(self.as_slice()) - } -} - -impl Identifier for Arc<[u8]> { - fn from_bytes(bytes: Vec) -> Result - where - Self: Sized, - { - Ok(Arc::from(bytes)) - } - - fn from_arc(arc: Arc<[u8]>) -> Result - where - Self: Sized, - { - Ok(arc) - } - - fn to_bytes(&self) -> Result, StoreError> { - Ok(Vec::from(&self[..])) - } - - fn string_repr(&self) -> String { - BASE64_STANDARD.encode(&self[..]) - } -} diff --git a/src/store/file_store.rs b/src/store/file_store.rs index b1738cb..dafa8f7 100644 --- a/src/store/file_store.rs +++ b/src/store/file_store.rs @@ -1,18 +1,17 @@ -use crate::{error_code::ErrorCode, file::File, repo::ArcRepo, store::Store}; +use crate::{ + error_code::ErrorCode, file::File, repo::ArcRepo, store::Store, stream::LocalBoxStream, +}; use actix_web::web::Bytes; use futures_core::Stream; use std::{ path::{Path, PathBuf}, - pin::Pin, + sync::Arc, }; use storage_path_generator::Generator; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_util::io::StreamReader; use tracing::Instrument; -mod file_id; -pub(crate) use file_id::FileId; - use super::StoreError; // - Settings Tree @@ -28,11 +27,8 @@ pub(crate) enum FileError { #[error("Failed to generate path")] PathGenerator(#[from] storage_path_generator::PathError), - #[error("Error formatting file store ID")] - IdError, - - #[error("Malformed file store ID")] - PrefixError, + #[error("Couldn't convert Path to String")] + StringError, #[error("Tried to save over existing file")] FileExists, @@ -44,7 +40,7 @@ impl FileError { Self::Io(_) => ErrorCode::FILE_IO_ERROR, Self::PathGenerator(_) => ErrorCode::PARSE_PATH_ERROR, Self::FileExists => ErrorCode::FILE_EXISTS, - Self::IdError | Self::PrefixError => ErrorCode::FORMAT_FILE_ID_ERROR, + Self::StringError => ErrorCode::FORMAT_FILE_ID_ERROR, } } } @@ -58,9 +54,6 @@ pub(crate) struct FileStore { #[async_trait::async_trait(?Send)] impl Store for FileStore { - type Identifier = FileId; - type Stream = Pin>>>; - async fn health_check(&self) -> Result<(), StoreError> { tokio::fs::metadata(&self.root_dir) .await @@ -74,7 +67,7 @@ impl Store for FileStore { &self, mut reader: Reader, _content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static, { @@ -92,7 +85,7 @@ impl Store for FileStore { &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static, { @@ -105,7 +98,7 @@ impl Store for FileStore { &self, bytes: Bytes, _content_type: mime::Mime, - ) -> Result { + ) -> Result, StoreError> { let path = self.next_file().await?; if let Err(e) = self.safe_save_bytes(&path, bytes).await { @@ -116,17 +109,17 @@ impl Store for FileStore { Ok(self.file_id_from_path(path)?) } - fn public_url(&self, _identifier: &Self::Identifier) -> Option { + fn public_url(&self, _identifier: &Arc) -> Option { None } #[tracing::instrument] async fn to_stream( &self, - identifier: &Self::Identifier, + identifier: &Arc, from_start: Option, len: Option, - ) -> Result { + ) -> Result>, StoreError> { let path = self.path_from_file_id(identifier); let file_span = tracing::trace_span!(parent: None, "File Stream"); @@ -147,7 +140,7 @@ impl Store for FileStore { #[tracing::instrument(skip(writer))] async fn read_into( &self, - identifier: &Self::Identifier, + identifier: &Arc, writer: &mut Writer, ) -> Result<(), std::io::Error> where @@ -161,7 +154,7 @@ impl Store for FileStore { } #[tracing::instrument] - async fn len(&self, identifier: &Self::Identifier) -> Result { + async fn len(&self, identifier: &Arc) -> Result { let path = self.path_from_file_id(identifier); let len = tokio::fs::metadata(path) @@ -173,7 +166,7 @@ impl Store for FileStore { } #[tracing::instrument] - async fn remove(&self, identifier: &Self::Identifier) -> Result<(), StoreError> { + async fn remove(&self, identifier: &Arc) -> Result<(), StoreError> { let path = self.path_from_file_id(identifier); self.safe_remove_file(path).await?; @@ -196,6 +189,14 @@ impl FileStore { }) } + fn file_id_from_path(&self, path: PathBuf) -> Result, FileError> { + path.to_str().ok_or(FileError::StringError).map(Into::into) + } + + fn path_from_file_id(&self, file_id: &Arc) -> PathBuf { + self.root_dir.join(file_id.as_ref()) + } + async fn next_directory(&self) -> Result { let path = self.path_gen.next(); diff --git a/src/store/file_store/file_id.rs b/src/store/file_store/file_id.rs deleted file mode 100644 index 279ddbc..0000000 --- a/src/store/file_store/file_id.rs +++ /dev/null @@ -1,57 +0,0 @@ -use crate::store::{ - file_store::{FileError, FileStore}, - Identifier, StoreError, -}; -use std::path::PathBuf; - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct FileId(PathBuf); - -impl Identifier for FileId { - fn to_bytes(&self) -> Result, StoreError> { - let vec = self - .0 - .to_str() - .ok_or(FileError::IdError)? - .as_bytes() - .to_vec(); - - Ok(vec) - } - - fn from_bytes(bytes: Vec) -> Result - where - Self: Sized, - { - let string = String::from_utf8(bytes).map_err(|_| FileError::IdError)?; - - let id = FileId(PathBuf::from(string)); - - Ok(id) - } - - fn from_arc(arc: std::sync::Arc<[u8]>) -> Result - where - Self: Sized, - { - Self::from_bytes(Vec::from(&arc[..])) - } - - fn string_repr(&self) -> String { - self.0.to_string_lossy().into_owned() - } -} - -impl FileStore { - pub(super) fn file_id_from_path(&self, path: PathBuf) -> Result { - let stripped = path - .strip_prefix(&self.root_dir) - .map_err(|_| FileError::PrefixError)?; - - Ok(FileId(stripped.to_path_buf())) - } - - pub(super) fn path_from_file_id(&self, file_id: &FileId) -> PathBuf { - self.root_dir.join(&file_id.0) - } -} diff --git a/src/store/object_store.rs b/src/store/object_store.rs index 301c181..4c1aab7 100644 --- a/src/store/object_store.rs +++ b/src/store/object_store.rs @@ -3,7 +3,7 @@ use crate::{ error_code::ErrorCode, repo::ArcRepo, store::Store, - stream::{IntoStreamer, StreamMap}, + stream::{IntoStreamer, LocalBoxStream, StreamMap}, }; use actix_rt::task::JoinError; use actix_web::{ @@ -19,16 +19,13 @@ use futures_core::Stream; use reqwest::{header::RANGE, Body, Response}; use reqwest_middleware::{ClientWithMiddleware, RequestBuilder}; use rusty_s3::{actions::S3Action, Bucket, BucketError, Credentials, UrlStyle}; -use std::{pin::Pin, string::FromUtf8Error, time::Duration}; +use std::{string::FromUtf8Error, sync::Arc, time::Duration}; use storage_path_generator::{Generator, Path}; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use tokio_util::io::ReaderStream; use tracing::Instrument; use url::Url; -mod object_id; -pub(crate) use object_id::ObjectId; - use super::StoreError; const CHUNK_SIZE: usize = 8_388_608; // 8 Mebibytes, min is 5 (5_242_880); @@ -189,9 +186,6 @@ async fn status_error(response: Response) -> StoreError { #[async_trait::async_trait(?Send)] impl Store for ObjectStore { - type Identifier = ObjectId; - type Stream = Pin>>>; - async fn health_check(&self) -> Result<(), StoreError> { let response = self .head_bucket_request() @@ -211,7 +205,7 @@ impl Store for ObjectStore { &self, reader: Reader, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static, { @@ -224,7 +218,7 @@ impl Store for ObjectStore { &self, mut stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static, { @@ -363,7 +357,7 @@ impl Store for ObjectStore { &self, bytes: Bytes, content_type: mime::Mime, - ) -> Result { + ) -> Result, StoreError> { let (req, object_id) = self.put_object_request(bytes.len(), content_type).await?; let response = req.body(bytes).send().await.map_err(ObjectError::from)?; @@ -375,9 +369,9 @@ impl Store for ObjectStore { Ok(object_id) } - fn public_url(&self, identifier: &Self::Identifier) -> Option { + fn public_url(&self, identifier: &Arc) -> Option { self.public_endpoint.clone().map(|mut endpoint| { - endpoint.set_path(identifier.as_str()); + endpoint.set_path(identifier.as_ref()); endpoint }) } @@ -385,10 +379,10 @@ impl Store for ObjectStore { #[tracing::instrument(skip(self))] async fn to_stream( &self, - identifier: &Self::Identifier, + identifier: &Arc, from_start: Option, len: Option, - ) -> Result { + ) -> Result>, StoreError> { let response = self .get_object_request(identifier, from_start, len) .send() @@ -409,7 +403,7 @@ impl Store for ObjectStore { #[tracing::instrument(skip(self, writer))] async fn read_into( &self, - identifier: &Self::Identifier, + identifier: &Arc, writer: &mut Writer, ) -> Result<(), std::io::Error> where @@ -440,7 +434,7 @@ impl Store for ObjectStore { } #[tracing::instrument(skip(self))] - async fn len(&self, identifier: &Self::Identifier) -> Result { + async fn len(&self, identifier: &Arc) -> Result { let response = self .head_object_request(identifier) .send() @@ -464,7 +458,7 @@ impl Store for ObjectStore { } #[tracing::instrument(skip(self))] - async fn remove(&self, identifier: &Self::Identifier) -> Result<(), StoreError> { + async fn remove(&self, identifier: &Arc) -> Result<(), StoreError> { let response = self .delete_object_request(identifier) .send() @@ -523,7 +517,7 @@ impl ObjectStore { &self, length: usize, content_type: mime::Mime, - ) -> Result<(RequestBuilder, ObjectId), StoreError> { + ) -> Result<(RequestBuilder, Arc), StoreError> { let path = self.next_file().await?; let mut action = self.bucket.put_object(Some(&self.credentials), &path); @@ -535,13 +529,13 @@ impl ObjectStore { .headers_mut() .insert("content-length", length.to_string()); - Ok((self.build_request(action), ObjectId::from_string(path))) + Ok((self.build_request(action), Arc::from(path))) } async fn create_multipart_request( &self, content_type: mime::Mime, - ) -> Result<(RequestBuilder, ObjectId), StoreError> { + ) -> Result<(RequestBuilder, Arc), StoreError> { let path = self.next_file().await?; let mut action = self @@ -552,13 +546,13 @@ impl ObjectStore { .headers_mut() .insert("content-type", content_type.as_ref()); - Ok((self.build_request(action), ObjectId::from_string(path))) + Ok((self.build_request(action), Arc::from(path))) } async fn create_upload_part_request( &self, buf: BytesStream, - object_id: &ObjectId, + object_id: &Arc, part_number: u16, upload_id: &str, ) -> Result { @@ -566,7 +560,7 @@ impl ObjectStore { let mut action = self.bucket.upload_part( Some(&self.credentials), - object_id.as_str(), + object_id.as_ref(), part_number, upload_id, ); @@ -601,13 +595,13 @@ impl ObjectStore { async fn send_complete_multipart_request<'a, I: Iterator>( &'a self, - object_id: &'a ObjectId, + object_id: &'a Arc, upload_id: &'a str, etags: I, ) -> Result { let mut action = self.bucket.complete_multipart_upload( Some(&self.credentials), - object_id.as_str(), + object_id.as_ref(), upload_id, etags, ); @@ -628,12 +622,12 @@ impl ObjectStore { fn create_abort_multipart_request( &self, - object_id: &ObjectId, + object_id: &Arc, upload_id: &str, ) -> RequestBuilder { let action = self.bucket.abort_multipart_upload( Some(&self.credentials), - object_id.as_str(), + object_id.as_ref(), upload_id, ); @@ -671,13 +665,13 @@ impl ObjectStore { fn get_object_request( &self, - identifier: &ObjectId, + identifier: &Arc, from_start: Option, len: Option, ) -> RequestBuilder { let action = self .bucket - .get_object(Some(&self.credentials), identifier.as_str()); + .get_object(Some(&self.credentials), identifier.as_ref()); let req = self.build_request(action); @@ -695,18 +689,18 @@ impl ObjectStore { ) } - fn head_object_request(&self, identifier: &ObjectId) -> RequestBuilder { + fn head_object_request(&self, identifier: &Arc) -> RequestBuilder { let action = self .bucket - .head_object(Some(&self.credentials), identifier.as_str()); + .head_object(Some(&self.credentials), identifier.as_ref()); self.build_request(action) } - fn delete_object_request(&self, identifier: &ObjectId) -> RequestBuilder { + fn delete_object_request(&self, identifier: &Arc) -> RequestBuilder { let action = self .bucket - .delete_object(Some(&self.credentials), identifier.as_str()); + .delete_object(Some(&self.credentials), identifier.as_ref()); self.build_request(action) } diff --git a/src/store/object_store/object_id.rs b/src/store/object_store/object_id.rs deleted file mode 100644 index f8a7dd9..0000000 --- a/src/store/object_store/object_id.rs +++ /dev/null @@ -1,37 +0,0 @@ -use crate::store::{object_store::ObjectError, Identifier, StoreError}; - -#[derive(Debug, Clone)] -pub(crate) struct ObjectId(String); - -impl Identifier for ObjectId { - fn to_bytes(&self) -> Result, StoreError> { - Ok(self.0.as_bytes().to_vec()) - } - - fn from_bytes(bytes: Vec) -> Result { - Ok(ObjectId( - String::from_utf8(bytes).map_err(ObjectError::from)?, - )) - } - - fn from_arc(arc: std::sync::Arc<[u8]>) -> Result - where - Self: Sized, - { - Self::from_bytes(Vec::from(&arc[..])) - } - - fn string_repr(&self) -> String { - self.0.clone() - } -} - -impl ObjectId { - pub(super) fn from_string(string: String) -> Self { - ObjectId(string) - } - - pub(super) fn as_str(&self) -> &str { - &self.0 - } -} From e580e7701eb75e6092bcac2cd1df908fc6ebedf0 Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 19:13:00 -0500 Subject: [PATCH 07/27] Finish implementing HashRepo --- docs/postgres-planning.md | 4 +- src/repo/postgres.rs | 202 ++++++++++++++++++++++++++++++++------ 2 files changed, 174 insertions(+), 32 deletions(-) diff --git a/docs/postgres-planning.md b/docs/postgres-planning.md index 49b4d6f..36801d0 100644 --- a/docs/postgres-planning.md +++ b/docs/postgres-planning.md @@ -165,8 +165,8 @@ CREATE TABLE queue ( ); -CREATE INDEX queue_status_index ON queue INCLUDE status; -CREATE INDEX heartbeat_index ON queue +CREATE INDEX queue_status_index ON queue INCLUDE queue, status; +CREATE INDEX heartbeat_index ON queue INCLUDE heartbeat; ``` claiming a job can be diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 9ffd3f4..190ac2b 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -9,7 +9,7 @@ use diesel_async::{ deadpool::{BuildError, Pool, PoolError}, AsyncDieselConnectionManager, }, - AsyncPgConnection, RunQueryDsl, + AsyncConnection, AsyncPgConnection, RunQueryDsl, }; use url::Url; @@ -223,63 +223,205 @@ impl HashRepo for PostgresRepo { let timestamp = to_primitive(timestamp); - /* - insert_into(hashes).values(( - hash.eq(&input_hash), - identifier.eq(&input_identifier) - )) - */ + let res = diesel::insert_into(hashes) + .values(( + hash.eq(&input_hash), + identifier.eq(input_identifier.as_ref()), + created_at.eq(×tamp), + )) + .execute(&mut conn) + .await; - todo!() + match res { + Ok(_) => Ok(Ok(())), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => Ok(Err(HashAlreadyExists)), + Err(e) => Err(PostgresError::Diesel(e).into()), + } } - async fn update_identifier(&self, hash: Hash, identifier: &Arc) -> Result<(), RepoError> { - todo!() + async fn update_identifier( + &self, + input_hash: Hash, + input_identifier: &Arc, + ) -> Result<(), RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::update(hashes) + .filter(hash.eq(&input_hash)) + .set(identifier.eq(input_identifier.as_ref())) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) } - async fn identifier(&self, hash: Hash) -> Result>, RepoError> { - todo!() + async fn identifier(&self, input_hash: Hash) -> Result>, RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = hashes + .select(identifier) + .filter(hash.eq(&input_hash)) + .get_result::(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt.map(Arc::from)) } async fn relate_variant_identifier( &self, - hash: Hash, - variant: String, - identifier: &Arc, + input_hash: Hash, + input_variant: String, + input_identifier: &Arc, ) -> Result, RepoError> { - todo!() + use schema::variants::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let res = diesel::insert_into(variants) + .values(( + hash.eq(&input_hash), + variant.eq(&input_variant), + identifier.eq(input_identifier.as_ref()), + )) + .execute(&mut conn) + .await; + + match res { + Ok(_) => Ok(Ok(())), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => Ok(Err(VariantAlreadyExists)), + Err(e) => Err(PostgresError::Diesel(e).into()), + } } async fn variant_identifier( &self, - hash: Hash, - variant: String, + input_hash: Hash, + input_variant: String, ) -> Result>, RepoError> { - todo!() + use schema::variants::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = variants + .select(identifier) + .filter(hash.eq(&input_hash)) + .filter(variant.eq(&input_variant)) + .get_result::(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .map(Arc::from); + + Ok(opt) } - async fn variants(&self, hash: Hash) -> Result)>, RepoError> { - todo!() + async fn variants(&self, input_hash: Hash) -> Result)>, RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let vec = variants + .select((variant, identifier)) + .filter(hash.eq(&input_hash)) + .get_results::<(String, String)>(&mut conn) + .await + .map_err(PostgresError::Diesel)? + .into_iter() + .map(|(s, i)| (s, Arc::from(i))) + .collect(); + + Ok(vec) } - async fn remove_variant(&self, hash: Hash, variant: String) -> Result<(), RepoError> { - todo!() + async fn remove_variant( + &self, + input_hash: Hash, + input_variant: String, + ) -> Result<(), RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(variants) + .filter(hash.eq(&input_hash)) + .filter(variant.eq(&input_variant)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) } async fn relate_motion_identifier( &self, - hash: Hash, - identifier: &Arc, + input_hash: Hash, + input_identifier: &Arc, ) -> Result<(), RepoError> { - todo!() + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::update(hashes) + .filter(hash.eq(&input_hash)) + .set(motion_identifier.eq(input_identifier.as_ref())) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) } - async fn motion_identifier(&self, hash: Hash) -> Result>, RepoError> { - todo!() + async fn motion_identifier(&self, input_hash: Hash) -> Result>, RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = hashes + .select(motion_identifier) + .filter(hash.eq(&input_hash)) + .get_result::>(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .flatten() + .map(Arc::from); + + Ok(opt) } - async fn cleanup_hash(&self, hash: Hash) -> Result<(), RepoError> { - todo!() + async fn cleanup_hash(&self, input_hash: Hash) -> Result<(), RepoError> { + let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + + conn.transaction(|conn| { + Box::pin(async move { + diesel::delete(schema::hashes::dsl::hashes) + .filter(schema::hashes::dsl::hash.eq(&input_hash)) + .execute(conn) + .await?; + + diesel::delete(schema::variants::dsl::variants) + .filter(schema::variants::dsl::hash.eq(&input_hash)) + .execute(conn) + .await + }) + }) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) } } From eac4cd54a4e86e41ec18699831752f23216c1d60 Mon Sep 17 00:00:00 2001 From: asonix Date: Sat, 2 Sep 2023 20:13:32 -0500 Subject: [PATCH 08/27] Initial work for pg notifications --- src/repo/postgres.rs | 65 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 61 insertions(+), 4 deletions(-) diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 190ac2b..a3720d5 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -7,10 +7,11 @@ use diesel::prelude::*; use diesel_async::{ pooled_connection::{ deadpool::{BuildError, Pool, PoolError}, - AsyncDieselConnectionManager, + AsyncDieselConnectionManager, ManagerConfig, }, AsyncConnection, AsyncPgConnection, RunQueryDsl, }; +use tokio_postgres::{AsyncMessage, Notification}; use url::Url; use crate::error_code::ErrorCode; @@ -23,6 +24,7 @@ use super::{ #[derive(Clone)] pub(crate) struct PostgresRepo { pool: Pool, + notifications: flume::Receiver, } #[derive(Debug, thiserror::Error)] @@ -69,15 +71,70 @@ impl PostgresRepo { handle.abort(); let _ = handle.await; - let config = AsyncDieselConnectionManager::::new(postgres_url); - let pool = Pool::builder(config) + let (tx, notifications) = flume::bounded(10); + + let mut config = ManagerConfig::default(); + config.custom_setup = build_handler(tx); + + let mgr = AsyncDieselConnectionManager::::new_with_config( + postgres_url, + config, + ); + let pool = Pool::builder(mgr) .build() .map_err(ConnectPostgresError::BuildPool)?; - Ok(PostgresRepo { pool }) + Ok(PostgresRepo { + pool, + notifications, + }) } } +type BoxFuture<'a, T> = std::pin::Pin + Send + 'a>>; +type ConfigFn = + Box BoxFuture<'_, ConnectionResult> + Send + Sync + 'static>; + +fn build_handler(sender: flume::Sender) -> ConfigFn { + Box::new( + move |config: &str| -> BoxFuture<'_, ConnectionResult> { + let sender = sender.clone(); + + Box::pin(async move { + let (client, mut conn) = + tokio_postgres::connect(config, tokio_postgres::tls::NoTls) + .await + .map_err(|e| ConnectionError::BadConnection(e.to_string()))?; + + // not very cash money (structured concurrency) of me + actix_rt::spawn(async move { + while let Some(res) = std::future::poll_fn(|cx| conn.poll_message(cx)).await { + match res { + Err(e) => { + tracing::error!("Database Connection {e:?}"); + return; + } + Ok(AsyncMessage::Notice(e)) => { + tracing::warn!("Database Notice {e:?}"); + } + Ok(AsyncMessage::Notification(notification)) => { + if sender.send_async(notification).await.is_err() { + tracing::warn!("Missed notification. Are we shutting down?"); + } + } + Ok(_) => { + tracing::warn!("Unhandled AsyncMessage!!! Please contact the developer of this application"); + } + } + } + }); + + AsyncPgConnection::try_from(client).await + }) + }, + ) +} + fn to_primitive(timestamp: time::OffsetDateTime) -> time::PrimitiveDateTime { let timestamp = timestamp.to_offset(time::UtcOffset::UTC); time::PrimitiveDateTime::new(timestamp.date(), timestamp.time()) From 443d327edfe171b618db46f723467e35e97f4c17 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 12:47:06 -0500 Subject: [PATCH 09/27] Implement a couple more repo traits --- Cargo.lock | 13 + Cargo.toml | 1 + docs/postgres-planning.md | 8 +- src/error_code.rs | 3 + src/queue.rs | 79 +-- src/queue/cleanup.rs | 4 +- src/queue/process.rs | 4 +- src/repo.rs | 155 +----- src/repo/alias.rs | 121 ++++ src/repo/delete_token.rs | 88 +++ src/repo/postgres.rs | 517 +++++++++++++++++- src/repo/postgres/job_status.rs | 6 + .../migrations/V0006__create_queue.rs | 5 +- .../V0007__create_store_migrations.rs | 3 +- src/repo/postgres/schema.rs | 8 +- src/repo/sled.rs | 25 +- 16 files changed, 800 insertions(+), 240 deletions(-) create mode 100644 src/repo/alias.rs create mode 100644 src/repo/delete_token.rs create mode 100644 src/repo/postgres/job_status.rs diff --git a/Cargo.lock b/Cargo.lock index 1a31c37..fae5c89 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -750,6 +750,18 @@ dependencies = [ "tokio-postgres", ] +[[package]] +name = "diesel-derive-enum" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81c5131a2895ef64741dad1d483f358c2a229a3a2d1b256778cdc5e146db64d4" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.29", +] + [[package]] name = "diesel_derives" version = "2.1.1" @@ -1807,6 +1819,7 @@ dependencies = [ "deadpool", "diesel", "diesel-async", + "diesel-derive-enum", "flume", "futures-core", "hex", diff --git a/Cargo.toml b/Cargo.toml index 449d844..abf9469 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ dashmap = "5.1.0" deadpool = { version = "0.9.5", features = ["rt_tokio_1"] } diesel = { version = "2.1.1", features = ["postgres_backend", "serde_json", "time", "uuid"] } diesel-async = { version = "0.4.1", features = ["postgres", "deadpool"] } +diesel-derive-enum = { version = "2.1.0", features = ["postgres"] } flume = "0.11.0" futures-core = "0.3" hex = "0.4.3" diff --git a/docs/postgres-planning.md b/docs/postgres-planning.md index 36801d0..d39d5cb 100644 --- a/docs/postgres-planning.md +++ b/docs/postgres-planning.md @@ -155,7 +155,7 @@ methods: CREATE TYPE job_status AS ENUM ('new', 'running'); -CREATE TABLE queue ( +CREATE TABLE job_queue ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), queue VARCHAR(30) NOT NULL, job JSONB NOT NULL, @@ -171,14 +171,14 @@ CREATE INDEX heartbeat_index ON queue INCLUDE heartbeat; claiming a job can be ```sql -UPDATE queue SET status = 'new', heartbeat = NULL +UPDATE job_queue SET status = 'new', heartbeat = NULL WHERE heartbeat IS NOT NULL AND heartbeat < NOW - INTERVAL '2 MINUTES'; -UPDATE queue SET status = 'running', heartbeat = CURRENT_TIMESTAMP +UPDATE job_queue SET status = 'running', heartbeat = CURRENT_TIMESTAMP WHERE id = ( SELECT id - FROM queue + FROM job_queue WHERE status = 'new' AND queue = '$QUEUE' ORDER BY queue_time ASC FOR UPDATE SKIP LOCKED diff --git a/src/error_code.rs b/src/error_code.rs index 67eaef2..365b261 100644 --- a/src/error_code.rs +++ b/src/error_code.rs @@ -62,6 +62,9 @@ impl ErrorCode { pub(crate) const EXTRACT_UPLOAD_RESULT: ErrorCode = ErrorCode { code: "extract-upload-result", }; + pub(crate) const EXTRACT_JOB: ErrorCode = ErrorCode { + code: "extract-job", + }; pub(crate) const CONFLICTED_RECORD: ErrorCode = ErrorCode { code: "conflicted-record", }; diff --git a/src/queue.rs b/src/queue.rs index a2c9deb..3f84fae 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -7,7 +7,6 @@ use crate::{ serde_str::Serde, store::Store, }; -use base64::{prelude::BASE64_STANDARD, Engine}; use std::{ future::Future, path::PathBuf, @@ -20,32 +19,6 @@ use tracing::Instrument; mod cleanup; mod process; -#[derive(Debug)] -struct Base64Bytes(Vec); - -impl serde::Serialize for Base64Bytes { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - let s = BASE64_STANDARD.encode(&self.0); - s.serialize(serializer) - } -} - -impl<'de> serde::Deserialize<'de> for Base64Bytes { - fn deserialize(deserializer: D) -> Result - where - D: serde::Deserializer<'de>, - { - let s: String = serde::Deserialize::deserialize(deserializer)?; - BASE64_STANDARD - .decode(s) - .map(Base64Bytes) - .map_err(|e| serde::de::Error::custom(e.to_string())) - } -} - const CLEANUP_QUEUE: &str = "cleanup"; const PROCESS_QUEUE: &str = "process"; @@ -91,18 +64,18 @@ pub(crate) async fn cleanup_alias( alias: Alias, token: DeleteToken, ) -> Result<(), Error> { - let job = serde_json::to_string(&Cleanup::Alias { + let job = serde_json::to_value(&Cleanup::Alias { alias: Serde::new(alias), token: Serde::new(token), }) .map_err(UploadError::PushJob)?; - repo.push(CLEANUP_QUEUE, job.into()).await?; + repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } pub(crate) async fn cleanup_hash(repo: &Arc, hash: Hash) -> Result<(), Error> { - let job = serde_json::to_string(&Cleanup::Hash { hash }).map_err(UploadError::PushJob)?; - repo.push(CLEANUP_QUEUE, job.into()).await?; + let job = serde_json::to_value(&Cleanup::Hash { hash }).map_err(UploadError::PushJob)?; + repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } @@ -110,11 +83,11 @@ pub(crate) async fn cleanup_identifier( repo: &Arc, identifier: &Arc, ) -> Result<(), Error> { - let job = serde_json::to_string(&Cleanup::Identifier { + let job = serde_json::to_value(&Cleanup::Identifier { identifier: identifier.to_string(), }) .map_err(UploadError::PushJob)?; - repo.push(CLEANUP_QUEUE, job.into()).await?; + repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } @@ -124,26 +97,26 @@ async fn cleanup_variants( variant: Option, ) -> Result<(), Error> { let job = - serde_json::to_string(&Cleanup::Variant { hash, variant }).map_err(UploadError::PushJob)?; - repo.push(CLEANUP_QUEUE, job.into()).await?; + serde_json::to_value(&Cleanup::Variant { hash, variant }).map_err(UploadError::PushJob)?; + repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } pub(crate) async fn cleanup_outdated_proxies(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_string(&Cleanup::OutdatedProxies).map_err(UploadError::PushJob)?; - repo.push(CLEANUP_QUEUE, job.into()).await?; + let job = serde_json::to_value(&Cleanup::OutdatedProxies).map_err(UploadError::PushJob)?; + repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } pub(crate) async fn cleanup_outdated_variants(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_string(&Cleanup::OutdatedVariants).map_err(UploadError::PushJob)?; - repo.push(CLEANUP_QUEUE, job.into()).await?; + let job = serde_json::to_value(&Cleanup::OutdatedVariants).map_err(UploadError::PushJob)?; + repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } pub(crate) async fn cleanup_all_variants(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_string(&Cleanup::AllVariants).map_err(UploadError::PushJob)?; - repo.push(CLEANUP_QUEUE, job.into()).await?; + let job = serde_json::to_value(&Cleanup::AllVariants).map_err(UploadError::PushJob)?; + repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } @@ -153,13 +126,13 @@ pub(crate) async fn queue_ingest( upload_id: UploadId, declared_alias: Option, ) -> Result<(), Error> { - let job = serde_json::to_string(&Process::Ingest { + let job = serde_json::to_value(&Process::Ingest { identifier: identifier.to_string(), declared_alias: declared_alias.map(Serde::new), upload_id: Serde::new(upload_id), }) .map_err(UploadError::PushJob)?; - repo.push(PROCESS_QUEUE, job.into()).await?; + repo.push(PROCESS_QUEUE, job).await?; Ok(()) } @@ -170,14 +143,14 @@ pub(crate) async fn queue_generate( process_path: PathBuf, process_args: Vec, ) -> Result<(), Error> { - let job = serde_json::to_string(&Process::Generate { + let job = serde_json::to_value(&Process::Generate { target_format, source: Serde::new(source), process_path, process_args, }) .map_err(UploadError::PushJob)?; - repo.push(PROCESS_QUEUE, job.into()).await?; + repo.push(PROCESS_QUEUE, job).await?; Ok(()) } @@ -220,7 +193,7 @@ async fn process_jobs( &'a Arc, &'a S, &'a Configuration, - &'a str, + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { @@ -284,13 +257,13 @@ where &'a Arc, &'a S, &'a Configuration, - &'a str, + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { let fut = async { - let (job_id, string) = repo.pop(queue, worker_id).await?; + let (job_id, job) = repo.pop(queue, worker_id).await?; let span = tracing::info_span!("Running Job"); @@ -303,7 +276,7 @@ where queue, worker_id, job_id, - (callback)(repo, store, config, string.as_ref()), + (callback)(repo, store, config, job), ) }) .instrument(span) @@ -337,7 +310,7 @@ async fn process_image_jobs( &'a S, &'a ProcessMap, &'a Configuration, - &'a str, + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { @@ -373,13 +346,13 @@ where &'a S, &'a ProcessMap, &'a Configuration, - &'a str, + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { let fut = async { - let (job_id, string) = repo.pop(queue, worker_id).await?; + let (job_id, job) = repo.pop(queue, worker_id).await?; let span = tracing::info_span!("Running Job"); @@ -392,7 +365,7 @@ where queue, worker_id, job_id, - (callback)(repo, store, process_map, config, string.as_ref()), + (callback)(repo, store, process_map, config, job), ) }) .instrument(span) diff --git a/src/queue/cleanup.rs b/src/queue/cleanup.rs index a6b76fd..aa31814 100644 --- a/src/queue/cleanup.rs +++ b/src/queue/cleanup.rs @@ -14,13 +14,13 @@ pub(super) fn perform<'a, S>( repo: &'a ArcRepo, store: &'a S, configuration: &'a Configuration, - job: &'a str, + job: serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> where S: Store, { Box::pin(async move { - match serde_json::from_str(job) { + match serde_json::from_value(job) { Ok(job) => match job { Cleanup::Hash { hash: in_hash } => hash(repo, in_hash).await?, Cleanup::Identifier { diff --git a/src/queue/process.rs b/src/queue/process.rs index a5f8e43..e257f10 100644 --- a/src/queue/process.rs +++ b/src/queue/process.rs @@ -17,13 +17,13 @@ pub(super) fn perform<'a, S>( store: &'a S, process_map: &'a ProcessMap, config: &'a Configuration, - job: &'a str, + job: serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> where S: Store + 'static, { Box::pin(async move { - match serde_json::from_str(job) { + match serde_json::from_value(job) { Ok(job) => match job { Process::Ingest { identifier, diff --git a/src/repo.rs b/src/repo.rs index 22318d4..e49f84a 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -9,11 +9,15 @@ use std::{fmt::Debug, sync::Arc}; use url::Url; use uuid::Uuid; +mod alias; +mod delete_token; mod hash; mod migrate; pub(crate) mod postgres; pub(crate) mod sled; +pub(crate) use alias::Alias; +pub(crate) use delete_token::DeleteToken; pub(crate) use hash::Hash; pub(crate) use migrate::{migrate_04, migrate_repo}; @@ -31,17 +35,6 @@ enum MaybeUuid { Name(String), } -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct Alias { - id: MaybeUuid, - extension: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct DeleteToken { - id: MaybeUuid, -} - #[derive(Debug)] pub(crate) struct HashAlreadyExists; #[derive(Debug)] @@ -372,13 +365,13 @@ impl JobId { #[async_trait::async_trait(?Send)] pub(crate) trait QueueRepo: BaseRepo { - async fn push(&self, queue: &'static str, job: Arc) -> Result; + async fn push(&self, queue: &'static str, job: serde_json::Value) -> Result; async fn pop( &self, queue: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc), RepoError>; + ) -> Result<(JobId, serde_json::Value), RepoError>; async fn heartbeat( &self, @@ -400,7 +393,7 @@ impl QueueRepo for Arc where T: QueueRepo, { - async fn push(&self, queue: &'static str, job: Arc) -> Result { + async fn push(&self, queue: &'static str, job: serde_json::Value) -> Result { T::push(self, queue, job).await } @@ -408,7 +401,7 @@ where &self, queue: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc), RepoError> { + ) -> Result<(JobId, serde_json::Value), RepoError> { T::pop(self, queue, worker_id).await } @@ -903,106 +896,6 @@ impl MaybeUuid { } } -fn split_at_dot(s: &str) -> Option<(&str, &str)> { - let index = s.find('.')?; - - Some(s.split_at(index)) -} - -impl Alias { - pub(crate) fn generate(extension: String) -> Self { - Alias { - id: MaybeUuid::Uuid(Uuid::new_v4()), - extension: Some(extension), - } - } - - pub(crate) fn from_existing(alias: &str) -> Self { - if let Some((start, end)) = split_at_dot(alias) { - Alias { - id: MaybeUuid::from_str(start), - extension: Some(end.into()), - } - } else { - Alias { - id: MaybeUuid::from_str(alias), - extension: None, - } - } - } - - pub(crate) fn extension(&self) -> Option<&str> { - self.extension.as_deref() - } - - pub(crate) fn to_bytes(&self) -> Vec { - let mut v = self.id.as_bytes().to_vec(); - - if let Some(ext) = self.extension() { - v.extend_from_slice(ext.as_bytes()); - } - - v - } - - pub(crate) fn from_slice(bytes: &[u8]) -> Option { - if let Ok(s) = std::str::from_utf8(bytes) { - Some(Self::from_existing(s)) - } else if bytes.len() >= 16 { - let id = Uuid::from_slice(&bytes[0..16]).expect("Already checked length"); - - let extension = if bytes.len() > 16 { - Some(String::from_utf8_lossy(&bytes[16..]).to_string()) - } else { - None - }; - - Some(Self { - id: MaybeUuid::Uuid(id), - extension, - }) - } else { - None - } - } -} - -impl DeleteToken { - pub(crate) fn from_existing(existing: &str) -> Self { - if let Ok(uuid) = Uuid::parse_str(existing) { - DeleteToken { - id: MaybeUuid::Uuid(uuid), - } - } else { - DeleteToken { - id: MaybeUuid::Name(existing.into()), - } - } - } - - pub(crate) fn generate() -> Self { - Self { - id: MaybeUuid::Uuid(Uuid::new_v4()), - } - } - - pub(crate) fn to_bytes(&self) -> Vec { - self.id.as_bytes().to_vec() - } - - pub(crate) fn from_slice(bytes: &[u8]) -> Option { - if let Ok(s) = std::str::from_utf8(bytes) { - Some(DeleteToken::from_existing(s)) - } else if bytes.len() == 16 { - Some(DeleteToken { - id: MaybeUuid::Uuid(Uuid::from_slice(bytes).ok()?), - }) - } else { - None - } - } -} - impl UploadId { pub(crate) fn generate() -> Self { Self { id: Uuid::new_v4() } @@ -1036,38 +929,6 @@ impl std::fmt::Display for MaybeUuid { } } -impl std::str::FromStr for DeleteToken { - type Err = std::convert::Infallible; - - fn from_str(s: &str) -> Result { - Ok(DeleteToken::from_existing(s)) - } -} - -impl std::fmt::Display for DeleteToken { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.id) - } -} - -impl std::str::FromStr for Alias { - type Err = std::convert::Infallible; - - fn from_str(s: &str) -> Result { - Ok(Alias::from_existing(s)) - } -} - -impl std::fmt::Display for Alias { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - if let Some(ext) = self.extension() { - write!(f, "{}{ext}", self.id) - } else { - write!(f, "{}", self.id) - } - } -} - #[cfg(test)] mod tests { use super::{Alias, DeleteToken, MaybeUuid, Uuid}; diff --git a/src/repo/alias.rs b/src/repo/alias.rs new file mode 100644 index 0000000..48156e1 --- /dev/null +++ b/src/repo/alias.rs @@ -0,0 +1,121 @@ +use diesel::{backend::Backend, sql_types::VarChar, AsExpression, FromSqlRow}; +use uuid::Uuid; + +use super::MaybeUuid; + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, AsExpression, FromSqlRow)] +#[diesel(sql_type = VarChar)] +pub(crate) struct Alias { + id: MaybeUuid, + extension: Option, +} + +impl diesel::serialize::ToSql for Alias { + fn to_sql<'b>( + &'b self, + out: &mut diesel::serialize::Output<'b, '_, diesel::pg::Pg>, + ) -> diesel::serialize::Result { + let s = self.to_string(); + + >::to_sql( + &s, + &mut out.reborrow(), + ) + } +} + +impl diesel::deserialize::FromSql for Alias +where + B: Backend, + String: diesel::deserialize::FromSql, +{ + fn from_sql( + bytes: ::RawValue<'_>, + ) -> diesel::deserialize::Result { + let s = String::from_sql(bytes)?; + + s.parse().map_err(From::from) + } +} + +impl Alias { + pub(crate) fn generate(extension: String) -> Self { + Alias { + id: MaybeUuid::Uuid(Uuid::new_v4()), + extension: Some(extension), + } + } + + pub(crate) fn from_existing(alias: &str) -> Self { + if let Some((start, end)) = split_at_dot(alias) { + Alias { + id: MaybeUuid::from_str(start), + extension: Some(end.into()), + } + } else { + Alias { + id: MaybeUuid::from_str(alias), + extension: None, + } + } + } + + pub(crate) fn extension(&self) -> Option<&str> { + self.extension.as_deref() + } + + pub(crate) fn to_bytes(&self) -> Vec { + let mut v = self.id.as_bytes().to_vec(); + + if let Some(ext) = self.extension() { + v.extend_from_slice(ext.as_bytes()); + } + + v + } + + pub(crate) fn from_slice(bytes: &[u8]) -> Option { + if let Ok(s) = std::str::from_utf8(bytes) { + Some(Self::from_existing(s)) + } else if bytes.len() >= 16 { + let id = Uuid::from_slice(&bytes[0..16]).expect("Already checked length"); + + let extension = if bytes.len() > 16 { + Some(String::from_utf8_lossy(&bytes[16..]).to_string()) + } else { + None + }; + + Some(Self { + id: MaybeUuid::Uuid(id), + extension, + }) + } else { + None + } + } +} + +fn split_at_dot(s: &str) -> Option<(&str, &str)> { + let index = s.find('.')?; + + Some(s.split_at(index)) +} + +impl std::str::FromStr for Alias { + type Err = std::convert::Infallible; + + fn from_str(s: &str) -> Result { + Ok(Alias::from_existing(s)) + } +} + +impl std::fmt::Display for Alias { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(ext) = self.extension() { + write!(f, "{}{ext}", self.id) + } else { + write!(f, "{}", self.id) + } + } +} diff --git a/src/repo/delete_token.rs b/src/repo/delete_token.rs new file mode 100644 index 0000000..7c06b09 --- /dev/null +++ b/src/repo/delete_token.rs @@ -0,0 +1,88 @@ +use diesel::{backend::Backend, sql_types::VarChar, AsExpression, FromSqlRow}; +use uuid::Uuid; + +use super::MaybeUuid; + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, AsExpression, FromSqlRow)] +#[diesel(sql_type = VarChar)] +pub(crate) struct DeleteToken { + id: MaybeUuid, +} + +impl diesel::serialize::ToSql for DeleteToken { + fn to_sql<'b>( + &'b self, + out: &mut diesel::serialize::Output<'b, '_, diesel::pg::Pg>, + ) -> diesel::serialize::Result { + let s = self.to_string(); + + >::to_sql( + &s, + &mut out.reborrow(), + ) + } +} + +impl diesel::deserialize::FromSql for DeleteToken +where + B: Backend, + String: diesel::deserialize::FromSql, +{ + fn from_sql( + bytes: ::RawValue<'_>, + ) -> diesel::deserialize::Result { + let s = String::from_sql(bytes)?; + + s.parse().map_err(From::from) + } +} + +impl DeleteToken { + pub(crate) fn from_existing(existing: &str) -> Self { + if let Ok(uuid) = Uuid::parse_str(existing) { + DeleteToken { + id: MaybeUuid::Uuid(uuid), + } + } else { + DeleteToken { + id: MaybeUuid::Name(existing.into()), + } + } + } + + pub(crate) fn generate() -> Self { + Self { + id: MaybeUuid::Uuid(Uuid::new_v4()), + } + } + + pub(crate) fn to_bytes(&self) -> Vec { + self.id.as_bytes().to_vec() + } + + pub(crate) fn from_slice(bytes: &[u8]) -> Option { + if let Ok(s) = std::str::from_utf8(bytes) { + Some(DeleteToken::from_existing(s)) + } else if bytes.len() == 16 { + Some(DeleteToken { + id: MaybeUuid::Uuid(Uuid::from_slice(bytes).ok()?), + }) + } else { + None + } + } +} + +impl std::str::FromStr for DeleteToken { + type Err = std::convert::Infallible; + + fn from_str(s: &str) -> Result { + Ok(DeleteToken::from_existing(s)) + } +} + +impl std::fmt::Display for DeleteToken { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.id) + } +} diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index a3720d5..eeea24d 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -1,8 +1,10 @@ mod embedded; +mod job_status; mod schema; use std::sync::Arc; +use dashmap::{DashMap, DashSet}; use diesel::prelude::*; use diesel_async::{ pooled_connection::{ @@ -11,20 +13,58 @@ use diesel_async::{ }, AsyncConnection, AsyncPgConnection, RunQueryDsl, }; +use tokio::sync::Notify; use tokio_postgres::{AsyncMessage, Notification}; use url::Url; +use uuid::Uuid; -use crate::error_code::ErrorCode; +use crate::{details::Details, error_code::ErrorCode}; + +use self::job_status::JobStatus; use super::{ - BaseRepo, Hash, HashAlreadyExists, HashPage, HashRepo, OrderedHash, RepoError, - VariantAlreadyExists, + Alias, AliasAlreadyExists, AliasRepo, BaseRepo, DeleteToken, DetailsRepo, Hash, + HashAlreadyExists, HashPage, HashRepo, JobId, OrderedHash, QueueRepo, RepoError, SettingsRepo, + StoreMigrationRepo, UploadId, VariantAlreadyExists, }; #[derive(Clone)] pub(crate) struct PostgresRepo { + inner: Arc, + notifications: Arc>, +} + +struct Inner { pool: Pool, - notifications: flume::Receiver, + queue_notifications: DashMap>, + completed_uploads: DashSet, + upload_notifier: Notify, +} + +async fn delegate_notifications(receiver: flume::Receiver, inner: Arc) { + while let Ok(notification) = receiver.recv_async().await { + match notification.channel() { + "queue_status_channel" => { + // new job inserted for queue + let queue_name = notification.payload().to_string(); + + inner + .queue_notifications + .entry(queue_name) + .or_insert_with(|| Arc::new(Notify::new())) + .notify_waiters(); + } + channel => { + tracing::info!( + "Unhandled postgres notification: {channel}: {}", + notification.payload() + ); + } + } + todo!() + } + + tracing::warn!("Notification delegator shutting down"); } #[derive(Debug, thiserror::Error)] @@ -46,6 +86,15 @@ pub(crate) enum PostgresError { #[error("Error in database")] Diesel(#[source] diesel::result::Error), + + #[error("Error deserializing hex value")] + Hex(#[source] hex::FromHexError), + + #[error("Error serializing details")] + SerializeDetails(#[source] serde_json::Error), + + #[error("Error deserializing details")] + DeserializeDetails(#[source] serde_json::Error), } impl PostgresError { @@ -71,7 +120,7 @@ impl PostgresRepo { handle.abort(); let _ = handle.await; - let (tx, notifications) = flume::bounded(10); + let (tx, rx) = flume::bounded(10); let mut config = ManagerConfig::default(); config.custom_setup = build_handler(tx); @@ -84,8 +133,17 @@ impl PostgresRepo { .build() .map_err(ConnectPostgresError::BuildPool)?; - Ok(PostgresRepo { + let inner = Arc::new(Inner { pool, + queue_notifications: DashMap::new(), + completed_uploads: DashSet::new(), + upload_notifier: Notify::new(), + }); + + let notifications = Arc::new(actix_rt::spawn(delegate_notifications(rx, inner.clone()))); + + Ok(PostgresRepo { + inner, notifications, }) } @@ -147,7 +205,7 @@ impl HashRepo for PostgresRepo { async fn size(&self) -> Result { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let count = hashes .count() @@ -161,7 +219,7 @@ impl HashRepo for PostgresRepo { async fn bound(&self, input_hash: Hash) -> Result, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let timestamp = hashes .select(created_at) @@ -185,7 +243,7 @@ impl HashRepo for PostgresRepo { ) -> Result { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let timestamp = to_primitive(date); @@ -212,7 +270,7 @@ impl HashRepo for PostgresRepo { ) -> Result { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let (mut page, prev) = if let Some(OrderedHash { timestamp, @@ -276,7 +334,7 @@ impl HashRepo for PostgresRepo { ) -> Result, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let timestamp = to_primitive(timestamp); @@ -306,7 +364,7 @@ impl HashRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; diesel::update(hashes) .filter(hash.eq(&input_hash)) @@ -321,7 +379,7 @@ impl HashRepo for PostgresRepo { async fn identifier(&self, input_hash: Hash) -> Result>, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let opt = hashes .select(identifier) @@ -342,7 +400,7 @@ impl HashRepo for PostgresRepo { ) -> Result, RepoError> { use schema::variants::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let res = diesel::insert_into(variants) .values(( @@ -370,7 +428,7 @@ impl HashRepo for PostgresRepo { ) -> Result>, RepoError> { use schema::variants::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let opt = variants .select(identifier) @@ -388,7 +446,7 @@ impl HashRepo for PostgresRepo { async fn variants(&self, input_hash: Hash) -> Result)>, RepoError> { use schema::variants::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let vec = variants .select((variant, identifier)) @@ -410,7 +468,7 @@ impl HashRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::variants::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; diesel::delete(variants) .filter(hash.eq(&input_hash)) @@ -429,7 +487,7 @@ impl HashRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; diesel::update(hashes) .filter(hash.eq(&input_hash)) @@ -444,7 +502,7 @@ impl HashRepo for PostgresRepo { async fn motion_identifier(&self, input_hash: Hash) -> Result>, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; let opt = hashes .select(motion_identifier) @@ -460,7 +518,7 @@ impl HashRepo for PostgresRepo { } async fn cleanup_hash(&self, input_hash: Hash) -> Result<(), RepoError> { - let mut conn = self.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; conn.transaction(|conn| { Box::pin(async move { @@ -482,6 +540,425 @@ impl HashRepo for PostgresRepo { } } +#[async_trait::async_trait(?Send)] +impl AliasRepo for PostgresRepo { + async fn create_alias( + &self, + input_alias: &Alias, + delete_token: &DeleteToken, + input_hash: Hash, + ) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let res = diesel::insert_into(aliases) + .values(( + alias.eq(input_alias), + hash.eq(&input_hash), + token.eq(delete_token), + )) + .execute(&mut conn) + .await; + + match res { + Ok(_) => Ok(Ok(())), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => Ok(Err(AliasAlreadyExists)), + Err(e) => Err(PostgresError::Diesel(e).into()), + } + } + + async fn delete_token(&self, input_alias: &Alias) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = aliases + .select(token) + .filter(alias.eq(input_alias)) + .get_result(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt) + } + + async fn hash(&self, input_alias: &Alias) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = aliases + .select(hash) + .filter(alias.eq(input_alias)) + .get_result(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt) + } + + async fn aliases_for_hash(&self, input_hash: Hash) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let vec = aliases + .select(alias) + .filter(hash.eq(&input_hash)) + .get_results(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(vec) + } + + async fn cleanup_alias(&self, input_alias: &Alias) -> Result<(), RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(aliases) + .filter(alias.eq(input_alias)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl SettingsRepo for PostgresRepo { + async fn set(&self, input_key: &'static str, input_value: Arc<[u8]>) -> Result<(), RepoError> { + use schema::settings::dsl::*; + + let input_value = hex::encode(input_value); + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::insert_into(settings) + .values((key.eq(input_key), value.eq(input_value))) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn get(&self, input_key: &'static str) -> Result>, RepoError> { + use schema::settings::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = settings + .select(value) + .filter(key.eq(input_key)) + .get_result::(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .map(hex::decode) + .transpose() + .map_err(PostgresError::Hex)? + .map(Arc::from); + + Ok(opt) + } + + async fn remove(&self, input_key: &'static str) -> Result<(), RepoError> { + use schema::settings::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(settings) + .filter(key.eq(input_key)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl DetailsRepo for PostgresRepo { + async fn relate_details( + &self, + input_identifier: &Arc, + input_details: &Details, + ) -> Result<(), RepoError> { + use schema::details::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let value = + serde_json::to_value(&input_details.inner).map_err(PostgresError::SerializeDetails)?; + + diesel::insert_into(details) + .values((identifier.eq(input_identifier.as_ref()), json.eq(&value))) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn details(&self, input_identifier: &Arc) -> Result, RepoError> { + use schema::details::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = details + .select(json) + .filter(identifier.eq(input_identifier.as_ref())) + .get_result::(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .map(serde_json::from_value) + .transpose() + .map_err(PostgresError::DeserializeDetails)? + .map(|inner| Details { inner }); + + Ok(opt) + } + + async fn cleanup_details(&self, input_identifier: &Arc) -> Result<(), RepoError> { + use schema::details::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(details) + .filter(identifier.eq(input_identifier.as_ref())) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl QueueRepo for PostgresRepo { + async fn push( + &self, + queue_name: &'static str, + job_json: serde_json::Value, + ) -> Result { + use schema::job_queue::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let job_id = diesel::insert_into(job_queue) + .values((queue.eq(queue_name), job.eq(job_json))) + .returning(id) + .get_result::(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(JobId(job_id)) + } + + async fn pop( + &self, + queue_name: &'static str, + worker_id: Uuid, + ) -> Result<(JobId, serde_json::Value), RepoError> { + use schema::job_queue::dsl::*; + + loop { + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let notifier: Arc = self + .inner + .queue_notifications + .entry(String::from(queue_name)) + .or_insert_with(|| Arc::new(Notify::new())) + .clone(); + + diesel::sql_query("LISTEN queue_status_channel;") + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + let timestamp = to_primitive(time::OffsetDateTime::now_utc()); + + diesel::update(job_queue) + .filter(heartbeat.le(timestamp.saturating_sub(time::Duration::minutes(2)))) + .set(( + heartbeat.eq(Option::::None), + status.eq(JobStatus::New), + )) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + // TODO: for_update().skip_locked() + let id_query = job_queue + .select(id) + .filter(status.eq(JobStatus::New).and(queue.eq(queue_name))) + .order(queue_time) + .into_boxed() + .single_value(); + + let opt = diesel::update(job_queue) + .filter(id.nullable().eq(id_query)) + .set(( + heartbeat.eq(timestamp), + status.eq(JobStatus::Running), + worker.eq(worker_id), + )) + .returning((id, job)) + .get_result(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)?; + + if let Some((job_id, job_json)) = opt { + diesel::sql_query("UNLISTEN queue_status_channel;") + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + return Ok((JobId(job_id), job_json)); + } + + let _ = actix_rt::time::timeout(std::time::Duration::from_secs(5), notifier.notified()) + .await; + + diesel::sql_query("UNLISTEN queue_status_channel;") + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + drop(conn); + } + } + + async fn heartbeat( + &self, + queue_name: &'static str, + worker_id: Uuid, + job_id: JobId, + ) -> Result<(), RepoError> { + use schema::job_queue::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let timestamp = to_primitive(time::OffsetDateTime::now_utc()); + + diesel::update(job_queue) + .filter( + id.eq(job_id.0) + .and(queue.eq(queue_name)) + .and(worker.eq(worker_id)), + ) + .set(heartbeat.eq(timestamp)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn complete_job( + &self, + queue_name: &'static str, + worker_id: Uuid, + job_id: JobId, + ) -> Result<(), RepoError> { + use schema::job_queue::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(job_queue) + .filter( + id.eq(job_id.0) + .and(queue.eq(queue_name)) + .and(worker.eq(worker_id)), + ) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl StoreMigrationRepo for PostgresRepo { + async fn is_continuing_migration(&self) -> Result { + use schema::store_migrations::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let count = store_migrations + .count() + .get_result::(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(count > 0) + } + + async fn mark_migrated( + &self, + input_old_identifier: &Arc, + input_new_identifier: &Arc, + ) -> Result<(), RepoError> { + use schema::store_migrations::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::insert_into(store_migrations) + .values(( + old_identifier.eq(input_old_identifier.as_ref()), + new_identifier.eq(input_new_identifier.as_ref()), + )) + .on_conflict((old_identifier, new_identifier)) + .do_nothing() + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn is_migrated(&self, input_old_identifier: &Arc) -> Result { + use schema::store_migrations::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let b = diesel::select(diesel::dsl::exists( + store_migrations.filter(old_identifier.eq(input_old_identifier.as_ref())), + )) + .get_result(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(b) + } + + async fn clear(&self) -> Result<(), RepoError> { + use schema::store_migrations::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(store_migrations) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + impl std::fmt::Debug for PostgresRepo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PostgresRepo") diff --git a/src/repo/postgres/job_status.rs b/src/repo/postgres/job_status.rs new file mode 100644 index 0000000..f9cee02 --- /dev/null +++ b/src/repo/postgres/job_status.rs @@ -0,0 +1,6 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, diesel_derive_enum::DbEnum)] +#[ExistingTypePath = "crate::repo::postgres::schema::sql_types::JobStatus"] +pub(super) enum JobStatus { + New, + Running, +} diff --git a/src/repo/postgres/migrations/V0006__create_queue.rs b/src/repo/postgres/migrations/V0006__create_queue.rs index d5161bf..ace111f 100644 --- a/src/repo/postgres/migrations/V0006__create_queue.rs +++ b/src/repo/postgres/migrations/V0006__create_queue.rs @@ -11,6 +11,7 @@ pub(crate) fn migration() -> String { t.inject_custom(r#""id" UUID PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL UNIQUE"#); t.add_column("queue", types::text().size(50).nullable(false)); t.add_column("job", types::custom("jsonb").nullable(false)); + t.add_column("worker", types::uuid().nullable(true)); t.add_column("status", types::custom("job_status").nullable(false)); t.add_column( "queue_time", @@ -18,7 +19,7 @@ pub(crate) fn migration() -> String { .nullable(false) .default(AutogenFunction::CurrentTimestamp), ); - t.add_column("heartbeat", types::datetime()); + t.add_column("heartbeat", types::datetime().nullable(true)); t.add_index("queue_status_index", types::index(["queue", "status"])); t.add_index("heartbeat_index", types::index(["heartbeat"])); @@ -30,7 +31,7 @@ CREATE OR REPLACE FUNCTION queue_status_notify() RETURNS trigger AS $$ BEGIN - PERFORM pg_notify('queue_status_channel', NEW.id::text); + PERFORM pg_notify('queue_status_channel', NEW.queue::text); RETURN NEW; END; $$ LANGUAGE plpgsql; diff --git a/src/repo/postgres/migrations/V0007__create_store_migrations.rs b/src/repo/postgres/migrations/V0007__create_store_migrations.rs index fe957b0..db106fc 100644 --- a/src/repo/postgres/migrations/V0007__create_store_migrations.rs +++ b/src/repo/postgres/migrations/V0007__create_store_migrations.rs @@ -7,9 +7,10 @@ pub(crate) fn migration() -> String { m.create_table("store_migrations", |t| { t.add_column( - "identifier", + "old_identifier", types::text().primary(true).nullable(false).unique(true), ); + t.add_column("new_identifier", types::text().nullable(false).unique(true)); }); m.make::().to_string() diff --git a/src/repo/postgres/schema.rs b/src/repo/postgres/schema.rs index 1b6aba4..da98a1c 100644 --- a/src/repo/postgres/schema.rs +++ b/src/repo/postgres/schema.rs @@ -38,9 +38,10 @@ diesel::table! { id -> Uuid, queue -> Text, job -> Jsonb, + worker -> Nullable, status -> JobStatus, queue_time -> Timestamp, - heartbeat -> Timestamp, + heartbeat -> Nullable, } } @@ -72,8 +73,9 @@ diesel::table! { } diesel::table! { - store_migrations (identifier) { - identifier -> Text, + store_migrations (old_identifier) { + old_identifier -> Text, + new_identifier -> Text, } } diff --git a/src/repo/sled.rs b/src/repo/sled.rs index 33cf611..671e186 100644 --- a/src/repo/sled.rs +++ b/src/repo/sled.rs @@ -46,10 +46,10 @@ pub(crate) enum SledError { Sled(#[from] sled::Error), #[error("Invalid details json")] - Details(serde_json::Error), + Details(#[source] serde_json::Error), #[error("Invalid upload result json")] - UploadResult(serde_json::Error), + UploadResult(#[source] serde_json::Error), #[error("Error parsing variant key")] VariantKey(#[from] VariantKeyError), @@ -57,6 +57,9 @@ pub(crate) enum SledError { #[error("Invalid string data in db")] Utf8(#[source] std::str::Utf8Error), + #[error("Invalid job json")] + Job(#[source] serde_json::Error), + #[error("Operation panicked")] Panic, @@ -70,6 +73,7 @@ impl SledError { Self::Sled(_) | Self::VariantKey(_) | Self::Utf8(_) => ErrorCode::SLED_ERROR, Self::Details(_) => ErrorCode::EXTRACT_DETAILS, Self::UploadResult(_) => ErrorCode::EXTRACT_UPLOAD_RESULT, + Self::Job(_) => ErrorCode::EXTRACT_JOB, Self::Panic => ErrorCode::PANIC, Self::Conflict => ErrorCode::CONFLICTED_RECORD, } @@ -660,11 +664,16 @@ fn try_into_arc_str(ivec: IVec) -> Result, SledError> { #[async_trait::async_trait(?Send)] impl QueueRepo for SledRepo { #[tracing::instrument(skip(self))] - async fn push(&self, queue_name: &'static str, job: Arc) -> Result { + async fn push( + &self, + queue_name: &'static str, + job: serde_json::Value, + ) -> Result { let metrics_guard = PushMetricsGuard::guard(queue_name); let id = JobId::gen(); let key = job_key(queue_name, id); + let job = serde_json::to_vec(&job).map_err(SledError::Job)?; let queue = self.queue.clone(); let job_state = self.job_state.clone(); @@ -709,7 +718,7 @@ impl QueueRepo for SledRepo { &self, queue_name: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc), RepoError> { + ) -> Result<(JobId, serde_json::Value), RepoError> { let metrics_guard = PopMetricsGuard::guard(queue_name); let now = time::OffsetDateTime::now_utc(); @@ -762,10 +771,14 @@ impl QueueRepo for SledRepo { tracing::Span::current().record("job_id", &format!("{job_id:?}")); - let opt = queue.get(&key)?.map(try_into_arc_str).transpose()?; + let opt = queue + .get(&key)? + .map(|ivec| serde_json::from_slice(&ivec[..])) + .transpose() + .map_err(SledError::Job)?; return Ok(opt.map(|job| (job_id, job))) - as Result)>, SledError>; + as Result, SledError>; } Ok(None) From 94cb2a9ef3f8c58169dff5b9d4a8c868108c1c3c Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 16:59:41 -0500 Subject: [PATCH 10/27] Pass clippy --- src/queue.rs | 18 +- src/repo.rs | 7 +- src/repo/hash.rs | 2 +- src/repo/postgres.rs | 458 +++++++++++++++++- .../migrations/V0009__create_uploads.rs | 8 +- src/repo/postgres/schema.rs | 2 +- 6 files changed, 470 insertions(+), 25 deletions(-) diff --git a/src/queue.rs b/src/queue.rs index 3f84fae..a7d3f25 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -64,7 +64,7 @@ pub(crate) async fn cleanup_alias( alias: Alias, token: DeleteToken, ) -> Result<(), Error> { - let job = serde_json::to_value(&Cleanup::Alias { + let job = serde_json::to_value(Cleanup::Alias { alias: Serde::new(alias), token: Serde::new(token), }) @@ -74,7 +74,7 @@ pub(crate) async fn cleanup_alias( } pub(crate) async fn cleanup_hash(repo: &Arc, hash: Hash) -> Result<(), Error> { - let job = serde_json::to_value(&Cleanup::Hash { hash }).map_err(UploadError::PushJob)?; + let job = serde_json::to_value(Cleanup::Hash { hash }).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } @@ -83,7 +83,7 @@ pub(crate) async fn cleanup_identifier( repo: &Arc, identifier: &Arc, ) -> Result<(), Error> { - let job = serde_json::to_value(&Cleanup::Identifier { + let job = serde_json::to_value(Cleanup::Identifier { identifier: identifier.to_string(), }) .map_err(UploadError::PushJob)?; @@ -97,25 +97,25 @@ async fn cleanup_variants( variant: Option, ) -> Result<(), Error> { let job = - serde_json::to_value(&Cleanup::Variant { hash, variant }).map_err(UploadError::PushJob)?; + serde_json::to_value(Cleanup::Variant { hash, variant }).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } pub(crate) async fn cleanup_outdated_proxies(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_value(&Cleanup::OutdatedProxies).map_err(UploadError::PushJob)?; + let job = serde_json::to_value(Cleanup::OutdatedProxies).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } pub(crate) async fn cleanup_outdated_variants(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_value(&Cleanup::OutdatedVariants).map_err(UploadError::PushJob)?; + let job = serde_json::to_value(Cleanup::OutdatedVariants).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } pub(crate) async fn cleanup_all_variants(repo: &Arc) -> Result<(), Error> { - let job = serde_json::to_value(&Cleanup::AllVariants).map_err(UploadError::PushJob)?; + let job = serde_json::to_value(Cleanup::AllVariants).map_err(UploadError::PushJob)?; repo.push(CLEANUP_QUEUE, job).await?; Ok(()) } @@ -126,7 +126,7 @@ pub(crate) async fn queue_ingest( upload_id: UploadId, declared_alias: Option, ) -> Result<(), Error> { - let job = serde_json::to_value(&Process::Ingest { + let job = serde_json::to_value(Process::Ingest { identifier: identifier.to_string(), declared_alias: declared_alias.map(Serde::new), upload_id: Serde::new(upload_id), @@ -143,7 +143,7 @@ pub(crate) async fn queue_generate( process_path: PathBuf, process_args: Vec, ) -> Result<(), Error> { - let job = serde_json::to_value(&Process::Generate { + let job = serde_json::to_value(Process::Generate { target_format, source: Serde::new(source), process_path, diff --git a/src/repo.rs b/src/repo.rs index e49f84a..6d3b2e3 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -562,10 +562,13 @@ impl HashPage { } } +type LocalBoxFuture<'a, T> = std::pin::Pin + 'a>>; + +type PageFuture = LocalBoxFuture<'static, Result>; + pub(crate) struct HashStream { repo: Option, - page_future: - Option>>>>, + page_future: Option, page: Option, } diff --git a/src/repo/hash.rs b/src/repo/hash.rs index b3c8036..23d9f8b 100644 --- a/src/repo/hash.rs +++ b/src/repo/hash.rs @@ -33,7 +33,7 @@ where fn from_sql(bytes: ::RawValue<'_>) -> diesel::deserialize::Result { let s = String::from_sql(bytes)?; - Self::from_base64(s).ok_or_else(|| format!("Invalid base64 hash").into()) + Self::from_base64(s).ok_or_else(|| "Invalid base64 hash".to_string().into()) } } diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index eeea24d..61835d9 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -2,9 +2,15 @@ mod embedded; mod job_status; mod schema; -use std::sync::Arc; +use std::{ + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + time::Duration, +}; -use dashmap::{DashMap, DashSet}; +use dashmap::DashMap; use diesel::prelude::*; use diesel_async::{ pooled_connection::{ @@ -18,26 +24,33 @@ use tokio_postgres::{AsyncMessage, Notification}; use url::Url; use uuid::Uuid; -use crate::{details::Details, error_code::ErrorCode}; +use crate::{ + details::Details, + error_code::{ErrorCode, OwnedErrorCode}, + serde_str::Serde, + stream::LocalBoxStream, +}; use self::job_status::JobStatus; use super::{ - Alias, AliasAlreadyExists, AliasRepo, BaseRepo, DeleteToken, DetailsRepo, Hash, - HashAlreadyExists, HashPage, HashRepo, JobId, OrderedHash, QueueRepo, RepoError, SettingsRepo, - StoreMigrationRepo, UploadId, VariantAlreadyExists, + Alias, AliasAccessRepo, AliasAlreadyExists, AliasRepo, BaseRepo, DeleteToken, DetailsRepo, + FullRepo, Hash, HashAlreadyExists, HashPage, HashRepo, JobId, OrderedHash, ProxyRepo, + QueueRepo, RepoError, SettingsRepo, StoreMigrationRepo, UploadId, UploadRepo, UploadResult, + VariantAccessRepo, VariantAlreadyExists, }; #[derive(Clone)] pub(crate) struct PostgresRepo { inner: Arc, + #[allow(dead_code)] notifications: Arc>, } struct Inner { + health_count: AtomicU64, pool: Pool, queue_notifications: DashMap>, - completed_uploads: DashSet, upload_notifier: Notify, } @@ -54,6 +67,9 @@ async fn delegate_notifications(receiver: flume::Receiver, inner: .or_insert_with(|| Arc::new(Notify::new())) .notify_waiters(); } + "upload_completion_channel" => { + inner.upload_notifier.notify_waiters(); + } channel => { tracing::info!( "Unhandled postgres notification: {channel}: {}", @@ -61,7 +77,6 @@ async fn delegate_notifications(receiver: flume::Receiver, inner: ); } } - todo!() } tracing::warn!("Notification delegator shutting down"); @@ -95,6 +110,12 @@ pub(crate) enum PostgresError { #[error("Error deserializing details")] DeserializeDetails(#[source] serde_json::Error), + + #[error("Error serializing upload result")] + SerializeUploadResult(#[source] serde_json::Error), + + #[error("Error deserializing upload result")] + DeserializeUploadResult(#[source] serde_json::Error), } impl PostgresError { @@ -134,9 +155,9 @@ impl PostgresRepo { .map_err(ConnectPostgresError::BuildPool)?; let inner = Arc::new(Inner { + health_count: AtomicU64::new(0), pool, queue_notifications: DashMap::new(), - completed_uploads: DashSet::new(), upload_notifier: Notify::new(), }); @@ -829,8 +850,7 @@ impl QueueRepo for PostgresRepo { return Ok((JobId(job_id), job_json)); } - let _ = actix_rt::time::timeout(std::time::Duration::from_secs(5), notifier.notified()) - .await; + let _ = actix_rt::time::timeout(Duration::from_secs(5), notifier.notified()).await; diesel::sql_query("UNLISTEN queue_status_channel;") .execute(&mut conn) @@ -959,6 +979,422 @@ impl StoreMigrationRepo for PostgresRepo { } } +#[async_trait::async_trait(?Send)] +impl ProxyRepo for PostgresRepo { + async fn relate_url(&self, input_url: Url, input_alias: Alias) -> Result<(), RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::insert_into(proxies) + .values((url.eq(input_url.as_str()), alias.eq(&input_alias))) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn related(&self, input_url: Url) -> Result, RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = proxies + .select(alias) + .filter(url.eq(input_url.as_str())) + .get_result(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt) + } + + async fn remove_relation(&self, input_alias: Alias) -> Result<(), RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(proxies) + .filter(alias.eq(&input_alias)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl AliasAccessRepo for PostgresRepo { + async fn set_accessed_alias( + &self, + input_alias: Alias, + timestamp: time::OffsetDateTime, + ) -> Result<(), RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let timestamp = to_primitive(timestamp); + + diesel::update(proxies) + .filter(alias.eq(&input_alias)) + .set(accessed.eq(timestamp)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn alias_accessed_at( + &self, + input_alias: Alias, + ) -> Result, RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = proxies + .select(accessed) + .filter(alias.eq(&input_alias)) + .get_result::(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .map(time::PrimitiveDateTime::assume_utc); + + Ok(opt) + } + + async fn older_aliases( + &self, + timestamp: time::OffsetDateTime, + ) -> Result>, RepoError> { + Ok(Box::pin(PageStream { + inner: self.inner.clone(), + future: None, + current: Vec::new(), + older_than: to_primitive(timestamp), + next: Box::new(|inner, older_than| { + Box::pin(async move { + use schema::proxies::dsl::*; + + let mut conn = inner.pool.get().await.map_err(PostgresError::Pool)?; + + let vec = proxies + .select((accessed, alias)) + .filter(accessed.lt(older_than)) + .order(accessed.desc()) + .limit(100) + .get_results(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(vec) + }) + }), + })) + } + + async fn remove_alias_access(&self, _: Alias) -> Result<(), RepoError> { + // Noop - handled by ProxyRepo::remove_relation + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl VariantAccessRepo for PostgresRepo { + async fn set_accessed_variant( + &self, + input_hash: Hash, + input_variant: String, + input_accessed: time::OffsetDateTime, + ) -> Result<(), RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let timestamp = to_primitive(input_accessed); + + diesel::update(variants) + .filter(hash.eq(&input_hash).and(variant.eq(&input_variant))) + .set(accessed.eq(timestamp)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn variant_accessed_at( + &self, + input_hash: Hash, + input_variant: String, + ) -> Result, RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let opt = variants + .select(accessed) + .filter(hash.eq(&input_hash).and(variant.eq(&input_variant))) + .get_result(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .map(time::PrimitiveDateTime::assume_utc); + + Ok(opt) + } + + async fn older_variants( + &self, + timestamp: time::OffsetDateTime, + ) -> Result>, RepoError> { + Ok(Box::pin(PageStream { + inner: self.inner.clone(), + future: None, + current: Vec::new(), + older_than: to_primitive(timestamp), + next: Box::new(|inner, older_than| { + Box::pin(async move { + use schema::variants::dsl::*; + + let mut conn = inner.pool.get().await.map_err(PostgresError::Pool)?; + + let vec = variants + .select((accessed, (hash, variant))) + .filter(accessed.lt(older_than)) + .order(accessed.desc()) + .limit(100) + .get_results(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(vec) + }) + }), + })) + } + + async fn remove_variant_access(&self, _: Hash, _: String) -> Result<(), RepoError> { + // Noop - handled by HashRepo::remove_variant + Ok(()) + } +} + +#[derive(serde::Deserialize, serde::Serialize)] +enum InnerUploadResult { + Success { + alias: Serde, + token: Serde, + }, + Failure { + message: String, + code: OwnedErrorCode, + }, +} + +impl From for InnerUploadResult { + fn from(u: UploadResult) -> Self { + match u { + UploadResult::Success { alias, token } => InnerUploadResult::Success { + alias: Serde::new(alias), + token: Serde::new(token), + }, + UploadResult::Failure { message, code } => InnerUploadResult::Failure { message, code }, + } + } +} + +impl From for UploadResult { + fn from(i: InnerUploadResult) -> Self { + match i { + InnerUploadResult::Success { alias, token } => UploadResult::Success { + alias: Serde::into_inner(alias), + token: Serde::into_inner(token), + }, + InnerUploadResult::Failure { message, code } => UploadResult::Failure { message, code }, + } + } +} + +#[async_trait::async_trait(?Send)] +impl UploadRepo for PostgresRepo { + async fn create_upload(&self) -> Result { + use schema::uploads::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let uuid = diesel::insert_into(uploads) + .default_values() + .returning(id) + .get_result(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(UploadId { id: uuid }) + } + + async fn wait(&self, upload_id: UploadId) -> Result { + use schema::uploads::dsl::*; + + loop { + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::sql_query("LISTEN upload_completion_channel;") + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + let opt = uploads + .select(result) + .filter(id.eq(upload_id.id)) + .get_result(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)? + .flatten(); + + if let Some(upload_result) = opt { + diesel::sql_query("UNLISTEN upload_completion_channel;") + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + let upload_result: InnerUploadResult = serde_json::from_value(upload_result) + .map_err(PostgresError::DeserializeUploadResult)?; + + return Ok(upload_result.into()); + } + + let _ = actix_rt::time::timeout( + Duration::from_secs(2), + self.inner.upload_notifier.notified(), + ) + .await; + + diesel::sql_query("UNLISTEN upload_completion_channel;") + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + drop(conn); + } + } + + async fn claim(&self, upload_id: UploadId) -> Result<(), RepoError> { + use schema::uploads::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + diesel::delete(uploads) + .filter(id.eq(upload_id.id)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + async fn complete_upload( + &self, + upload_id: UploadId, + upload_result: UploadResult, + ) -> Result<(), RepoError> { + use schema::uploads::dsl::*; + + let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + + let upload_result: InnerUploadResult = upload_result.into(); + let upload_result = + serde_json::to_value(&upload_result).map_err(PostgresError::SerializeUploadResult)?; + + diesel::update(uploads) + .filter(id.eq(upload_id.id)) + .set(result.eq(upload_result)) + .execute(&mut conn) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl FullRepo for PostgresRepo { + async fn health_check(&self) -> Result<(), RepoError> { + let next = self.inner.health_count.fetch_add(1, Ordering::Relaxed); + + self.set("health-value", Arc::from(next.to_be_bytes())) + .await?; + + Ok(()) + } +} + +type LocalBoxFuture<'a, T> = std::pin::Pin + 'a>>; + +type NextFuture = LocalBoxFuture<'static, Result, RepoError>>; + +struct PageStream { + inner: Arc, + future: Option>, + current: Vec, + older_than: time::PrimitiveDateTime, + next: Box, time::PrimitiveDateTime) -> NextFuture>, +} + +impl futures_core::Stream for PageStream +where + I: Unpin, +{ + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + + loop { + // Pop because we reversed the list + if let Some(alias) = this.current.pop() { + return std::task::Poll::Ready(Some(Ok(alias))); + } + + if let Some(future) = this.future.as_mut() { + let res = std::task::ready!(future.as_mut().poll(cx)); + + this.future.take(); + + match res { + Ok(page) if page.is_empty() => { + return std::task::Poll::Ready(None); + } + Ok(page) => { + let (mut timestamps, mut aliases): (Vec<_>, Vec<_>) = + page.into_iter().unzip(); + // reverse because we use .pop() to get next + aliases.reverse(); + + this.current = aliases; + this.older_than = timestamps.pop().expect("Verified nonempty"); + } + Err(e) => return std::task::Poll::Ready(Some(Err(e))), + } + } else { + let inner = this.inner.clone(); + let older_than = this.older_than; + + this.future = Some((this.next)(inner, older_than)); + } + } + } +} + impl std::fmt::Debug for PostgresRepo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PostgresRepo") diff --git a/src/repo/postgres/migrations/V0009__create_uploads.rs b/src/repo/postgres/migrations/V0009__create_uploads.rs index 9214ec4..4cfbafa 100644 --- a/src/repo/postgres/migrations/V0009__create_uploads.rs +++ b/src/repo/postgres/migrations/V0009__create_uploads.rs @@ -7,7 +7,13 @@ pub(crate) fn migration() -> String { m.create_table("uploads", |t| { t.inject_custom(r#""id" UUID PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL UNIQUE"#); - t.add_column("result", types::custom("jsonb")); + t.add_column("result", types::custom("jsonb").nullable(true)); + t.add_column( + "created_at", + types::datetime() + .nullable(false) + .default(AutogenFunction::CurrentTimestamp), + ); }); m.inject_custom( diff --git a/src/repo/postgres/schema.rs b/src/repo/postgres/schema.rs index da98a1c..80c5d4e 100644 --- a/src/repo/postgres/schema.rs +++ b/src/repo/postgres/schema.rs @@ -82,7 +82,7 @@ diesel::table! { diesel::table! { uploads (id) { id -> Uuid, - result -> Jsonb, + result -> Nullable, } } From a2933dbebcb72d5739e27218b0727a65385fe9d0 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 17:11:34 -0500 Subject: [PATCH 11/27] Implement all the todos --- src/error.rs | 2 +- src/error_code.rs | 4 ++++ src/lib.rs | 18 +++++++++--------- src/repo.rs | 2 +- src/repo/postgres.rs | 10 +++++++++- src/repo/postgres/schema.rs | 1 + 6 files changed, 25 insertions(+), 12 deletions(-) diff --git a/src/error.rs b/src/error.rs index 38c481e..de2d9a3 100644 --- a/src/error.rs +++ b/src/error.rs @@ -166,7 +166,7 @@ impl UploadError { Self::InvalidToken => ErrorCode::INVALID_DELETE_TOKEN, Self::UnsupportedProcessExtension => ErrorCode::INVALID_FILE_EXTENSION, Self::DuplicateAlias => ErrorCode::DUPLICATE_ALIAS, - Self::PushJob(_) => todo!(), + Self::PushJob(_) => ErrorCode::PUSH_JOB, Self::Range => ErrorCode::RANGE_NOT_SATISFIABLE, Self::Limit(_) => ErrorCode::VALIDATE_FILE_SIZE, Self::Timeout(_) => ErrorCode::STREAM_TOO_SLOW, diff --git a/src/error_code.rs b/src/error_code.rs index 365b261..0a940ad 100644 --- a/src/error_code.rs +++ b/src/error_code.rs @@ -56,12 +56,16 @@ impl ErrorCode { code: "already-claimed", }; pub(crate) const SLED_ERROR: ErrorCode = ErrorCode { code: "sled-error" }; + pub(crate) const POSTGRES_ERROR: ErrorCode = ErrorCode { + code: "postgres-error", + }; pub(crate) const EXTRACT_DETAILS: ErrorCode = ErrorCode { code: "extract-details", }; pub(crate) const EXTRACT_UPLOAD_RESULT: ErrorCode = ErrorCode { code: "extract-upload-result", }; + pub(crate) const PUSH_JOB: ErrorCode = ErrorCode { code: "push-job" }; pub(crate) const EXTRACT_JOB: ErrorCode = ErrorCode { code: "extract-job", }; diff --git a/src/lib.rs b/src/lib.rs index e6476c7..f55e7de 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2038,16 +2038,14 @@ impl PictRsConfiguration { match repo { Repo::Sled(sled_repo) => { - launch_file_store( - Arc::new(sled_repo.clone()), - store, - client, - config, - move |sc| sled_extra_config(sc, sled_repo.clone()), - ) + launch_file_store(arc_repo, store, client, config, move |sc| { + sled_extra_config(sc, sled_repo.clone()) + }) .await?; } - Repo::Postgres(_) => todo!(), + Repo::Postgres(_) => { + launch_file_store(arc_repo, store, client, config, |_| {}).await?; + } } } config::Store::ObjectStorage(config::ObjectStorage { @@ -2101,7 +2099,9 @@ impl PictRsConfiguration { }) .await?; } - Repo::Postgres(_) => todo!(), + Repo::Postgres(_) => { + launch_object_store(arc_repo, store, client, config, |_| {}).await?; + } } } } diff --git a/src/repo.rs b/src/repo.rs index 6d3b2e3..a177745 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -877,7 +877,7 @@ impl Repo { pub(crate) fn to_arc(&self) -> ArcRepo { match self { Self::Sled(sled_repo) => Arc::new(sled_repo.clone()), - Self::Postgres(_) => todo!(), + Self::Postgres(postgres_repo) => Arc::new(postgres_repo.clone()), } } } diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 61835d9..d03a7c5 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -120,7 +120,15 @@ pub(crate) enum PostgresError { impl PostgresError { pub(super) const fn error_code(&self) -> ErrorCode { - todo!() + match self { + Self::Pool(_) + | Self::Diesel(_) + | Self::SerializeDetails(_) + | Self::SerializeUploadResult(_) + | Self::Hex(_) => ErrorCode::POSTGRES_ERROR, + Self::DeserializeDetails(_) => ErrorCode::EXTRACT_DETAILS, + Self::DeserializeUploadResult(_) => ErrorCode::EXTRACT_UPLOAD_RESULT, + } } } diff --git a/src/repo/postgres/schema.rs b/src/repo/postgres/schema.rs index 80c5d4e..964937b 100644 --- a/src/repo/postgres/schema.rs +++ b/src/repo/postgres/schema.rs @@ -83,6 +83,7 @@ diesel::table! { uploads (id) { id -> Uuid, result -> Nullable, + created_at -> Timestamp, } } From 8a892ba62224006f6a2feaec22d07ac6dad67c22 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 18:21:46 -0500 Subject: [PATCH 12/27] it might work --- dev.toml | 4 ++ src/queue/process.rs | 4 +- src/repo.rs | 2 +- src/repo/postgres.rs | 60 +++++++++++++++++-- .../migrations/V0000__enable_pgcrypto.rs | 2 +- .../migrations/V0006__create_queue.rs | 5 +- src/store/file_store.rs | 25 +++++--- 7 files changed, 84 insertions(+), 18 deletions(-) diff --git a/dev.toml b/dev.toml index ae32642..6cdd71b 100644 --- a/dev.toml +++ b/dev.toml @@ -62,6 +62,10 @@ crf_max = 12 type = 'postgres' url = 'postgres://postgres:1234@localhost:5432/postgres' +# [repo] +# type = 'sled' +# path = 'data/sled-repo-local' + [store] type = 'filesystem' path = 'data/files-local' diff --git a/src/queue/process.rs b/src/queue/process.rs index e257f10..5454dc8 100644 --- a/src/queue/process.rs +++ b/src/queue/process.rs @@ -68,7 +68,7 @@ where }) } -#[tracing::instrument(skip_all)] +#[tracing::instrument(skip(repo, store, media))] async fn process_ingest( repo: &ArcRepo, store: &S, @@ -126,7 +126,7 @@ where } #[allow(clippy::too_many_arguments)] -#[tracing::instrument(skip_all)] +#[tracing::instrument(skip(repo, store, process_map, process_path, process_args, config))] async fn generate( repo: &ArcRepo, store: &S, diff --git a/src/repo.rs b/src/repo.rs index a177745..fdf5ce1 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -524,7 +524,7 @@ where } } -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] pub(crate) struct OrderedHash { timestamp: time::OffsetDateTime, hash: Hash, diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index d03a7c5..ab08fef 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -231,6 +231,7 @@ impl BaseRepo for PostgresRepo {} #[async_trait::async_trait(?Send)] impl HashRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn size(&self) -> Result { use schema::hashes::dsl::*; @@ -245,6 +246,7 @@ impl HashRepo for PostgresRepo { Ok(count.try_into().expect("non-negative count")) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn bound(&self, input_hash: Hash) -> Result, RepoError> { use schema::hashes::dsl::*; @@ -265,6 +267,7 @@ impl HashRepo for PostgresRepo { })) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn hash_page_by_date( &self, date: time::OffsetDateTime, @@ -292,6 +295,7 @@ impl HashRepo for PostgresRepo { self.hashes_ordered(ordered_hash, limit).await } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn hashes_ordered( &self, bound: Option, @@ -355,6 +359,7 @@ impl HashRepo for PostgresRepo { }) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn create_hash_with_timestamp( &self, input_hash: Hash, @@ -386,6 +391,7 @@ impl HashRepo for PostgresRepo { } } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn update_identifier( &self, input_hash: Hash, @@ -405,6 +411,7 @@ impl HashRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn identifier(&self, input_hash: Hash) -> Result>, RepoError> { use schema::hashes::dsl::*; @@ -421,6 +428,7 @@ impl HashRepo for PostgresRepo { Ok(opt.map(Arc::from)) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn relate_variant_identifier( &self, input_hash: Hash, @@ -450,6 +458,7 @@ impl HashRepo for PostgresRepo { } } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn variant_identifier( &self, input_hash: Hash, @@ -472,6 +481,7 @@ impl HashRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn variants(&self, input_hash: Hash) -> Result)>, RepoError> { use schema::variants::dsl::*; @@ -490,6 +500,7 @@ impl HashRepo for PostgresRepo { Ok(vec) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn remove_variant( &self, input_hash: Hash, @@ -509,6 +520,7 @@ impl HashRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn relate_motion_identifier( &self, input_hash: Hash, @@ -528,6 +540,7 @@ impl HashRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn motion_identifier(&self, input_hash: Hash) -> Result>, RepoError> { use schema::hashes::dsl::*; @@ -546,18 +559,19 @@ impl HashRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn cleanup_hash(&self, input_hash: Hash) -> Result<(), RepoError> { let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; conn.transaction(|conn| { Box::pin(async move { - diesel::delete(schema::hashes::dsl::hashes) - .filter(schema::hashes::dsl::hash.eq(&input_hash)) + diesel::delete(schema::variants::dsl::variants) + .filter(schema::variants::dsl::hash.eq(&input_hash)) .execute(conn) .await?; - diesel::delete(schema::variants::dsl::variants) - .filter(schema::variants::dsl::hash.eq(&input_hash)) + diesel::delete(schema::hashes::dsl::hashes) + .filter(schema::hashes::dsl::hash.eq(&input_hash)) .execute(conn) .await }) @@ -571,6 +585,7 @@ impl HashRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl AliasRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn create_alias( &self, input_alias: &Alias, @@ -600,6 +615,7 @@ impl AliasRepo for PostgresRepo { } } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn delete_token(&self, input_alias: &Alias) -> Result, RepoError> { use schema::aliases::dsl::*; @@ -616,6 +632,7 @@ impl AliasRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn hash(&self, input_alias: &Alias) -> Result, RepoError> { use schema::aliases::dsl::*; @@ -632,6 +649,7 @@ impl AliasRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn aliases_for_hash(&self, input_hash: Hash) -> Result, RepoError> { use schema::aliases::dsl::*; @@ -647,6 +665,7 @@ impl AliasRepo for PostgresRepo { Ok(vec) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn cleanup_alias(&self, input_alias: &Alias) -> Result<(), RepoError> { use schema::aliases::dsl::*; @@ -664,6 +683,7 @@ impl AliasRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl SettingsRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self, input_value))] async fn set(&self, input_key: &'static str, input_value: Arc<[u8]>) -> Result<(), RepoError> { use schema::settings::dsl::*; @@ -672,7 +692,10 @@ impl SettingsRepo for PostgresRepo { let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; diesel::insert_into(settings) - .values((key.eq(input_key), value.eq(input_value))) + .values((key.eq(input_key), value.eq(&input_value))) + .on_conflict(key) + .do_update() + .set(value.eq(&input_value)) .execute(&mut conn) .await .map_err(PostgresError::Diesel)?; @@ -680,6 +703,7 @@ impl SettingsRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn get(&self, input_key: &'static str) -> Result>, RepoError> { use schema::settings::dsl::*; @@ -700,6 +724,7 @@ impl SettingsRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn remove(&self, input_key: &'static str) -> Result<(), RepoError> { use schema::settings::dsl::*; @@ -717,6 +742,7 @@ impl SettingsRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl DetailsRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self, input_details))] async fn relate_details( &self, input_identifier: &Arc, @@ -738,6 +764,7 @@ impl DetailsRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn details(&self, input_identifier: &Arc) -> Result, RepoError> { use schema::details::dsl::*; @@ -758,6 +785,7 @@ impl DetailsRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn cleanup_details(&self, input_identifier: &Arc) -> Result<(), RepoError> { use schema::details::dsl::*; @@ -775,6 +803,7 @@ impl DetailsRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl QueueRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self, job_json))] async fn push( &self, queue_name: &'static str, @@ -794,6 +823,7 @@ impl QueueRepo for PostgresRepo { Ok(JobId(job_id)) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn pop( &self, queue_name: &'static str, @@ -869,6 +899,7 @@ impl QueueRepo for PostgresRepo { } } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn heartbeat( &self, queue_name: &'static str, @@ -895,6 +926,7 @@ impl QueueRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn complete_job( &self, queue_name: &'static str, @@ -921,6 +953,7 @@ impl QueueRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl StoreMigrationRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn is_continuing_migration(&self) -> Result { use schema::store_migrations::dsl::*; @@ -935,6 +968,7 @@ impl StoreMigrationRepo for PostgresRepo { Ok(count > 0) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn mark_migrated( &self, input_old_identifier: &Arc, @@ -958,6 +992,7 @@ impl StoreMigrationRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn is_migrated(&self, input_old_identifier: &Arc) -> Result { use schema::store_migrations::dsl::*; @@ -973,6 +1008,7 @@ impl StoreMigrationRepo for PostgresRepo { Ok(b) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn clear(&self) -> Result<(), RepoError> { use schema::store_migrations::dsl::*; @@ -989,6 +1025,7 @@ impl StoreMigrationRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl ProxyRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn relate_url(&self, input_url: Url, input_alias: Alias) -> Result<(), RepoError> { use schema::proxies::dsl::*; @@ -1003,6 +1040,7 @@ impl ProxyRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn related(&self, input_url: Url) -> Result, RepoError> { use schema::proxies::dsl::*; @@ -1019,6 +1057,7 @@ impl ProxyRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn remove_relation(&self, input_alias: Alias) -> Result<(), RepoError> { use schema::proxies::dsl::*; @@ -1036,6 +1075,7 @@ impl ProxyRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl AliasAccessRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn set_accessed_alias( &self, input_alias: Alias, @@ -1057,6 +1097,7 @@ impl AliasAccessRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn alias_accessed_at( &self, input_alias: Alias, @@ -1077,6 +1118,7 @@ impl AliasAccessRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn older_aliases( &self, timestamp: time::OffsetDateTime, @@ -1115,6 +1157,7 @@ impl AliasAccessRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl VariantAccessRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn set_accessed_variant( &self, input_hash: Hash, @@ -1137,6 +1180,7 @@ impl VariantAccessRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn variant_accessed_at( &self, input_hash: Hash, @@ -1158,6 +1202,7 @@ impl VariantAccessRepo for PostgresRepo { Ok(opt) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn older_variants( &self, timestamp: time::OffsetDateTime, @@ -1232,6 +1277,7 @@ impl From for UploadResult { #[async_trait::async_trait(?Send)] impl UploadRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn create_upload(&self) -> Result { use schema::uploads::dsl::*; @@ -1247,6 +1293,7 @@ impl UploadRepo for PostgresRepo { Ok(UploadId { id: uuid }) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn wait(&self, upload_id: UploadId) -> Result { use schema::uploads::dsl::*; @@ -1293,6 +1340,7 @@ impl UploadRepo for PostgresRepo { } } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn claim(&self, upload_id: UploadId) -> Result<(), RepoError> { use schema::uploads::dsl::*; @@ -1307,6 +1355,7 @@ impl UploadRepo for PostgresRepo { Ok(()) } + #[tracing::instrument(level = "DEBUG", skip(self))] async fn complete_upload( &self, upload_id: UploadId, @@ -1333,6 +1382,7 @@ impl UploadRepo for PostgresRepo { #[async_trait::async_trait(?Send)] impl FullRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn health_check(&self) -> Result<(), RepoError> { let next = self.inner.health_count.fetch_add(1, Ordering::Relaxed); diff --git a/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs b/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs index 44eb85c..4c8bd57 100644 --- a/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs +++ b/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs @@ -5,7 +5,7 @@ use barrel::{types, Migration}; pub(crate) fn migration() -> String { let mut m = Migration::new(); - m.inject_custom("CREATE EXTENSION pgcrypto;"); + m.inject_custom("CREATE EXTENSION IF NOT EXISTS pgcrypto;"); m.make::().to_string() } diff --git a/src/repo/postgres/migrations/V0006__create_queue.rs b/src/repo/postgres/migrations/V0006__create_queue.rs index ace111f..f9a96e0 100644 --- a/src/repo/postgres/migrations/V0006__create_queue.rs +++ b/src/repo/postgres/migrations/V0006__create_queue.rs @@ -12,7 +12,10 @@ pub(crate) fn migration() -> String { t.add_column("queue", types::text().size(50).nullable(false)); t.add_column("job", types::custom("jsonb").nullable(false)); t.add_column("worker", types::uuid().nullable(true)); - t.add_column("status", types::custom("job_status").nullable(false)); + t.add_column( + "status", + types::custom("job_status").nullable(false).default("new"), + ); t.add_column( "queue_time", types::datetime() diff --git a/src/store/file_store.rs b/src/store/file_store.rs index dafa8f7..e208261 100644 --- a/src/store/file_store.rs +++ b/src/store/file_store.rs @@ -27,6 +27,9 @@ pub(crate) enum FileError { #[error("Failed to generate path")] PathGenerator(#[from] storage_path_generator::PathError), + #[error("Couldn't strip root dir")] + PrefixError, + #[error("Couldn't convert Path to String")] StringError, @@ -40,7 +43,7 @@ impl FileError { Self::Io(_) => ErrorCode::FILE_IO_ERROR, Self::PathGenerator(_) => ErrorCode::PARSE_PATH_ERROR, Self::FileExists => ErrorCode::FILE_EXISTS, - Self::StringError => ErrorCode::FORMAT_FILE_ID_ERROR, + Self::StringError | Self::PrefixError => ErrorCode::FORMAT_FILE_ID_ERROR, } } } @@ -54,6 +57,7 @@ pub(crate) struct FileStore { #[async_trait::async_trait(?Send)] impl Store for FileStore { + #[tracing::instrument(level = "DEBUG", skip(self))] async fn health_check(&self) -> Result<(), StoreError> { tokio::fs::metadata(&self.root_dir) .await @@ -62,7 +66,7 @@ impl Store for FileStore { Ok(()) } - #[tracing::instrument(skip(reader))] + #[tracing::instrument(skip(self, reader))] async fn save_async_read( &self, mut reader: Reader, @@ -93,7 +97,7 @@ impl Store for FileStore { .await } - #[tracing::instrument(skip(bytes))] + #[tracing::instrument(skip(self, bytes))] async fn save_bytes( &self, bytes: Bytes, @@ -113,7 +117,7 @@ impl Store for FileStore { None } - #[tracing::instrument] + #[tracing::instrument(skip(self))] async fn to_stream( &self, identifier: &Arc, @@ -137,7 +141,7 @@ impl Store for FileStore { Ok(Box::pin(stream)) } - #[tracing::instrument(skip(writer))] + #[tracing::instrument(skip(self, writer))] async fn read_into( &self, identifier: &Arc, @@ -153,7 +157,7 @@ impl Store for FileStore { Ok(()) } - #[tracing::instrument] + #[tracing::instrument(skip(self))] async fn len(&self, identifier: &Arc) -> Result { let path = self.path_from_file_id(identifier); @@ -165,7 +169,7 @@ impl Store for FileStore { Ok(len) } - #[tracing::instrument] + #[tracing::instrument(skip(self))] async fn remove(&self, identifier: &Arc) -> Result<(), StoreError> { let path = self.path_from_file_id(identifier); @@ -190,7 +194,11 @@ impl FileStore { } fn file_id_from_path(&self, path: PathBuf) -> Result, FileError> { - path.to_str().ok_or(FileError::StringError).map(Into::into) + path.strip_prefix(&self.root_dir) + .map_err(|_| FileError::PrefixError)? + .to_str() + .ok_or(FileError::StringError) + .map(Into::into) } fn path_from_file_id(&self, file_id: &Arc) -> PathBuf { @@ -219,6 +227,7 @@ impl FileStore { Ok(target_path.join(filename)) } + #[tracing::instrument(level = "DEBUG", skip(self, path), fields(path = ?path.as_ref()))] async fn safe_remove_file>(&self, path: P) -> Result<(), FileError> { tokio::fs::remove_file(&path).await?; self.try_remove_parents(path.as_ref()).await; From 31caea438e6b5a1c55d2a33d2b0e7298702c7366 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 20:05:29 -0500 Subject: [PATCH 13/27] Fix slow connection pool access --- dev.toml | 3 +- docker/object-storage/docker-compose.yml | 37 ++-- src/repo.rs | 218 ----------------------- src/repo/alias.rs | 153 ++++++++++++++++ src/repo/delete_token.rs | 72 ++++++++ src/repo/postgres.rs | 206 +++++++++++---------- src/repo/postgres/schema.rs | 2 +- 7 files changed, 365 insertions(+), 326 deletions(-) diff --git a/dev.toml b/dev.toml index 6cdd71b..6c72d3e 100644 --- a/dev.toml +++ b/dev.toml @@ -11,6 +11,7 @@ targets = 'warn,tracing_actix_web=info,actix_server=info,actix_web=info' buffer_capacity = 102400 [tracing.opentelemetry] +url = 'http://127.0.0.1:4317' service_name = 'pict-rs' targets = 'info' @@ -60,7 +61,7 @@ crf_max = 12 [repo] type = 'postgres' -url = 'postgres://postgres:1234@localhost:5432/postgres' +url = 'postgres://pictrs:1234@localhost:5432/pictrs' # [repo] # type = 'sled' diff --git a/docker/object-storage/docker-compose.yml b/docker/object-storage/docker-compose.yml index a4e5d5a..3a028a9 100644 --- a/docker/object-storage/docker-compose.yml +++ b/docker/object-storage/docker-compose.yml @@ -13,7 +13,7 @@ services: # - "6669:6669" # environment: # - PICTRS__TRACING__CONSOLE__ADDRESS=0.0.0.0:6669 - # - PICTRS__TRACING__OPENTELEMETRY__URL=http://otel:4137 + # - PICTRS__TRACING__OPENTELEMETRY__URL=http://jaeger:4317 # - RUST_BACKTRACE=1 # stdin_open: true # tty: true @@ -27,7 +27,7 @@ services: # - "8081:8081" # environment: # - PICTRS_PROXY_UPSTREAM=http://pictrs:8080 - # - PICTRS_PROXY_OPENTELEMETRY_URL=http://otel:4137 + # - PICTRS_PROXY_OPENTELEMETRY_URL=http://jaeger:4317 minio: image: quay.io/minio/minio @@ -39,7 +39,7 @@ services: - ./storage/minio:/mnt garage: - image: dxflrs/garage:v0.8.1 + image: dxflrs/garage:v0.8.3 ports: - "3900:3900" - "3901:3901" @@ -47,26 +47,35 @@ services: - "3903:3903" - "3904:3904" environment: - - RUST_LOG=debug + - RUST_LOG=info volumes: - ./storage/garage:/mnt - ./garage.toml:/etc/garage.toml - otel: - image: otel/opentelemetry-collector:latest - command: --config otel-local-config.yaml + postgres: + image: postgres:15-alpine + ports: + - "5432:5432" + environment: + - PGDATA=/var/lib/postgresql/data + - POSTGRES_DB=pictrs + - POSTGRES_USER=pictrs + - POSTGRES_PASSWORD=1234 volumes: - - type: bind - source: ./otel.yml - target: /otel-local-config.yaml - restart: always - depends_on: - - jaeger + - ./storage/postgres:/var/lib/postgresql/data jaeger: - image: jaegertracing/all-in-one:1 + image: jaegertracing/all-in-one:1.48 ports: + - "6831:6831/udp" + - "6832:6832/udp" + - "5778:5778" + - "4317:4317" + - "4138:4138" - "14250:14250" + - "14268:14268" + - "14269:14269" + - "9411:9411" # To view traces, visit http://localhost:16686 - "16686:16686" restart: always diff --git a/src/repo.rs b/src/repo.rs index fdf5ce1..0c8c312 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -931,221 +931,3 @@ impl std::fmt::Display for MaybeUuid { } } } - -#[cfg(test)] -mod tests { - use super::{Alias, DeleteToken, MaybeUuid, Uuid}; - - #[test] - fn string_delete_token() { - let delete_token = DeleteToken::from_existing("blah"); - - assert_eq!( - delete_token, - DeleteToken { - id: MaybeUuid::Name(String::from("blah")) - } - ) - } - - #[test] - fn uuid_string_delete_token() { - let uuid = Uuid::new_v4(); - - let delete_token = DeleteToken::from_existing(&uuid.to_string()); - - assert_eq!( - delete_token, - DeleteToken { - id: MaybeUuid::Uuid(uuid), - } - ) - } - - #[test] - fn bytes_delete_token() { - let delete_token = DeleteToken::from_slice(b"blah").unwrap(); - - assert_eq!( - delete_token, - DeleteToken { - id: MaybeUuid::Name(String::from("blah")) - } - ) - } - - #[test] - fn uuid_bytes_delete_token() { - let uuid = Uuid::new_v4(); - - let delete_token = DeleteToken::from_slice(&uuid.as_bytes()[..]).unwrap(); - - assert_eq!( - delete_token, - DeleteToken { - id: MaybeUuid::Uuid(uuid), - } - ) - } - - #[test] - fn uuid_bytes_string_delete_token() { - let uuid = Uuid::new_v4(); - - let delete_token = DeleteToken::from_slice(uuid.to_string().as_bytes()).unwrap(); - - assert_eq!( - delete_token, - DeleteToken { - id: MaybeUuid::Uuid(uuid), - } - ) - } - - #[test] - fn string_alias() { - let alias = Alias::from_existing("blah"); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Name(String::from("blah")), - extension: None - } - ); - } - - #[test] - fn string_alias_ext() { - let alias = Alias::from_existing("blah.mp4"); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Name(String::from("blah")), - extension: Some(String::from(".mp4")), - } - ); - } - - #[test] - fn uuid_string_alias() { - let uuid = Uuid::new_v4(); - - let alias = Alias::from_existing(&uuid.to_string()); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Uuid(uuid), - extension: None, - } - ) - } - - #[test] - fn uuid_string_alias_ext() { - let uuid = Uuid::new_v4(); - - let alias_str = format!("{uuid}.mp4"); - let alias = Alias::from_existing(&alias_str); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Uuid(uuid), - extension: Some(String::from(".mp4")), - } - ) - } - - #[test] - fn bytes_alias() { - let alias = Alias::from_slice(b"blah").unwrap(); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Name(String::from("blah")), - extension: None - } - ); - } - - #[test] - fn bytes_alias_ext() { - let alias = Alias::from_slice(b"blah.mp4").unwrap(); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Name(String::from("blah")), - extension: Some(String::from(".mp4")), - } - ); - } - - #[test] - fn uuid_bytes_alias() { - let uuid = Uuid::new_v4(); - - let alias = Alias::from_slice(&uuid.as_bytes()[..]).unwrap(); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Uuid(uuid), - extension: None, - } - ) - } - - #[test] - fn uuid_bytes_string_alias() { - let uuid = Uuid::new_v4(); - - let alias = Alias::from_slice(uuid.to_string().as_bytes()).unwrap(); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Uuid(uuid), - extension: None, - } - ) - } - - #[test] - fn uuid_bytes_alias_ext() { - let uuid = Uuid::new_v4(); - - let mut alias_bytes = uuid.as_bytes().to_vec(); - alias_bytes.extend_from_slice(b".mp4"); - - let alias = Alias::from_slice(&alias_bytes).unwrap(); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Uuid(uuid), - extension: Some(String::from(".mp4")), - } - ) - } - - #[test] - fn uuid_bytes_string_alias_ext() { - let uuid = Uuid::new_v4(); - - let alias_str = format!("{uuid}.mp4"); - let alias = Alias::from_slice(alias_str.as_bytes()).unwrap(); - - assert_eq!( - alias, - Alias { - id: MaybeUuid::Uuid(uuid), - extension: Some(String::from(".mp4")), - } - ) - } -} diff --git a/src/repo/alias.rs b/src/repo/alias.rs index 48156e1..a0d021c 100644 --- a/src/repo/alias.rs +++ b/src/repo/alias.rs @@ -119,3 +119,156 @@ impl std::fmt::Display for Alias { } } } + +#[cfg(test)] +mod tests { + use super::{Alias, MaybeUuid}; + use uuid::Uuid; + + #[test] + fn string_alias() { + let alias = Alias::from_existing("blah"); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Name(String::from("blah")), + extension: None + } + ); + } + + #[test] + fn string_alias_ext() { + let alias = Alias::from_existing("blah.mp4"); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Name(String::from("blah")), + extension: Some(String::from(".mp4")), + } + ); + } + + #[test] + fn uuid_string_alias() { + let uuid = Uuid::new_v4(); + + let alias = Alias::from_existing(&uuid.to_string()); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Uuid(uuid), + extension: None, + } + ) + } + + #[test] + fn uuid_string_alias_ext() { + let uuid = Uuid::new_v4(); + + let alias_str = format!("{uuid}.mp4"); + let alias = Alias::from_existing(&alias_str); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Uuid(uuid), + extension: Some(String::from(".mp4")), + } + ) + } + + #[test] + fn bytes_alias() { + let alias = Alias::from_slice(b"blah").unwrap(); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Name(String::from("blah")), + extension: None + } + ); + } + + #[test] + fn bytes_alias_ext() { + let alias = Alias::from_slice(b"blah.mp4").unwrap(); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Name(String::from("blah")), + extension: Some(String::from(".mp4")), + } + ); + } + + #[test] + fn uuid_bytes_alias() { + let uuid = Uuid::new_v4(); + + let alias = Alias::from_slice(&uuid.as_bytes()[..]).unwrap(); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Uuid(uuid), + extension: None, + } + ) + } + + #[test] + fn uuid_bytes_string_alias() { + let uuid = Uuid::new_v4(); + + let alias = Alias::from_slice(uuid.to_string().as_bytes()).unwrap(); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Uuid(uuid), + extension: None, + } + ) + } + + #[test] + fn uuid_bytes_alias_ext() { + let uuid = Uuid::new_v4(); + + let mut alias_bytes = uuid.as_bytes().to_vec(); + alias_bytes.extend_from_slice(b".mp4"); + + let alias = Alias::from_slice(&alias_bytes).unwrap(); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Uuid(uuid), + extension: Some(String::from(".mp4")), + } + ) + } + + #[test] + fn uuid_bytes_string_alias_ext() { + let uuid = Uuid::new_v4(); + + let alias_str = format!("{uuid}.mp4"); + let alias = Alias::from_slice(alias_str.as_bytes()).unwrap(); + + assert_eq!( + alias, + Alias { + id: MaybeUuid::Uuid(uuid), + extension: Some(String::from(".mp4")), + } + ) + } +} diff --git a/src/repo/delete_token.rs b/src/repo/delete_token.rs index 7c06b09..ca1a50f 100644 --- a/src/repo/delete_token.rs +++ b/src/repo/delete_token.rs @@ -86,3 +86,75 @@ impl std::fmt::Display for DeleteToken { write!(f, "{}", self.id) } } + +#[cfg(test)] +mod tests { + use super::{DeleteToken, MaybeUuid}; + use uuid::Uuid; + + #[test] + fn string_delete_token() { + let delete_token = DeleteToken::from_existing("blah"); + + assert_eq!( + delete_token, + DeleteToken { + id: MaybeUuid::Name(String::from("blah")) + } + ) + } + + #[test] + fn uuid_string_delete_token() { + let uuid = Uuid::new_v4(); + + let delete_token = DeleteToken::from_existing(&uuid.to_string()); + + assert_eq!( + delete_token, + DeleteToken { + id: MaybeUuid::Uuid(uuid), + } + ) + } + + #[test] + fn bytes_delete_token() { + let delete_token = DeleteToken::from_slice(b"blah").unwrap(); + + assert_eq!( + delete_token, + DeleteToken { + id: MaybeUuid::Name(String::from("blah")) + } + ) + } + + #[test] + fn uuid_bytes_delete_token() { + let uuid = Uuid::new_v4(); + + let delete_token = DeleteToken::from_slice(&uuid.as_bytes()[..]).unwrap(); + + assert_eq!( + delete_token, + DeleteToken { + id: MaybeUuid::Uuid(uuid), + } + ) + } + + #[test] + fn uuid_bytes_string_delete_token() { + let uuid = Uuid::new_v4(); + + let delete_token = DeleteToken::from_slice(uuid.to_string().as_bytes()).unwrap(); + + assert_eq!( + delete_token, + DeleteToken { + id: MaybeUuid::Uuid(uuid), + } + ) + } +} diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index ab08fef..f63331d 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -14,7 +14,7 @@ use dashmap::DashMap; use diesel::prelude::*; use diesel_async::{ pooled_connection::{ - deadpool::{BuildError, Pool, PoolError}, + deadpool::{BuildError, Object, Pool, PoolError}, AsyncDieselConnectionManager, ManagerConfig, }, AsyncConnection, AsyncPgConnection, RunQueryDsl, @@ -176,6 +176,17 @@ impl PostgresRepo { notifications, }) } + + async fn get_connection(&self) -> Result, PostgresError> { + self.inner.get_connection().await + } +} + +impl Inner { + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn get_connection(&self) -> Result, PostgresError> { + self.pool.get().await.map_err(PostgresError::Pool) + } } type BoxFuture<'a, T> = std::pin::Pin + Send + 'a>>; @@ -235,7 +246,7 @@ impl HashRepo for PostgresRepo { async fn size(&self) -> Result { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let count = hashes .count() @@ -250,7 +261,7 @@ impl HashRepo for PostgresRepo { async fn bound(&self, input_hash: Hash) -> Result, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let timestamp = hashes .select(created_at) @@ -275,7 +286,7 @@ impl HashRepo for PostgresRepo { ) -> Result { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let timestamp = to_primitive(date); @@ -303,7 +314,7 @@ impl HashRepo for PostgresRepo { ) -> Result { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let (mut page, prev) = if let Some(OrderedHash { timestamp, @@ -368,7 +379,7 @@ impl HashRepo for PostgresRepo { ) -> Result, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let timestamp = to_primitive(timestamp); @@ -399,7 +410,7 @@ impl HashRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::update(hashes) .filter(hash.eq(&input_hash)) @@ -415,7 +426,7 @@ impl HashRepo for PostgresRepo { async fn identifier(&self, input_hash: Hash) -> Result>, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = hashes .select(identifier) @@ -437,7 +448,7 @@ impl HashRepo for PostgresRepo { ) -> Result, RepoError> { use schema::variants::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let res = diesel::insert_into(variants) .values(( @@ -466,7 +477,7 @@ impl HashRepo for PostgresRepo { ) -> Result>, RepoError> { use schema::variants::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = variants .select(identifier) @@ -485,7 +496,7 @@ impl HashRepo for PostgresRepo { async fn variants(&self, input_hash: Hash) -> Result)>, RepoError> { use schema::variants::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let vec = variants .select((variant, identifier)) @@ -508,7 +519,7 @@ impl HashRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::variants::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(variants) .filter(hash.eq(&input_hash)) @@ -528,7 +539,7 @@ impl HashRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::update(hashes) .filter(hash.eq(&input_hash)) @@ -544,7 +555,7 @@ impl HashRepo for PostgresRepo { async fn motion_identifier(&self, input_hash: Hash) -> Result>, RepoError> { use schema::hashes::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = hashes .select(motion_identifier) @@ -561,7 +572,7 @@ impl HashRepo for PostgresRepo { #[tracing::instrument(level = "DEBUG", skip(self))] async fn cleanup_hash(&self, input_hash: Hash) -> Result<(), RepoError> { - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; conn.transaction(|conn| { Box::pin(async move { @@ -594,7 +605,7 @@ impl AliasRepo for PostgresRepo { ) -> Result, RepoError> { use schema::aliases::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let res = diesel::insert_into(aliases) .values(( @@ -619,7 +630,7 @@ impl AliasRepo for PostgresRepo { async fn delete_token(&self, input_alias: &Alias) -> Result, RepoError> { use schema::aliases::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = aliases .select(token) @@ -636,7 +647,7 @@ impl AliasRepo for PostgresRepo { async fn hash(&self, input_alias: &Alias) -> Result, RepoError> { use schema::aliases::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = aliases .select(hash) @@ -653,7 +664,7 @@ impl AliasRepo for PostgresRepo { async fn aliases_for_hash(&self, input_hash: Hash) -> Result, RepoError> { use schema::aliases::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let vec = aliases .select(alias) @@ -669,7 +680,7 @@ impl AliasRepo for PostgresRepo { async fn cleanup_alias(&self, input_alias: &Alias) -> Result<(), RepoError> { use schema::aliases::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(aliases) .filter(alias.eq(input_alias)) @@ -689,7 +700,7 @@ impl SettingsRepo for PostgresRepo { let input_value = hex::encode(input_value); - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::insert_into(settings) .values((key.eq(input_key), value.eq(&input_value))) @@ -707,7 +718,7 @@ impl SettingsRepo for PostgresRepo { async fn get(&self, input_key: &'static str) -> Result>, RepoError> { use schema::settings::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = settings .select(value) @@ -728,7 +739,7 @@ impl SettingsRepo for PostgresRepo { async fn remove(&self, input_key: &'static str) -> Result<(), RepoError> { use schema::settings::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(settings) .filter(key.eq(input_key)) @@ -750,7 +761,7 @@ impl DetailsRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::details::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let value = serde_json::to_value(&input_details.inner).map_err(PostgresError::SerializeDetails)?; @@ -768,7 +779,7 @@ impl DetailsRepo for PostgresRepo { async fn details(&self, input_identifier: &Arc) -> Result, RepoError> { use schema::details::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = details .select(json) @@ -789,7 +800,7 @@ impl DetailsRepo for PostgresRepo { async fn cleanup_details(&self, input_identifier: &Arc) -> Result<(), RepoError> { use schema::details::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(details) .filter(identifier.eq(input_identifier.as_ref())) @@ -811,7 +822,7 @@ impl QueueRepo for PostgresRepo { ) -> Result { use schema::job_queue::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let job_id = diesel::insert_into(job_queue) .values((queue.eq(queue_name), job.eq(job_json))) @@ -832,7 +843,7 @@ impl QueueRepo for PostgresRepo { use schema::job_queue::dsl::*; loop { - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let notifier: Arc = self .inner @@ -848,7 +859,7 @@ impl QueueRepo for PostgresRepo { let timestamp = to_primitive(time::OffsetDateTime::now_utc()); - diesel::update(job_queue) + let count = diesel::update(job_queue) .filter(heartbeat.le(timestamp.saturating_sub(time::Duration::minutes(2)))) .set(( heartbeat.eq(Option::::None), @@ -858,44 +869,58 @@ impl QueueRepo for PostgresRepo { .await .map_err(PostgresError::Diesel)?; - // TODO: for_update().skip_locked() - let id_query = job_queue - .select(id) - .filter(status.eq(JobStatus::New).and(queue.eq(queue_name))) - .order(queue_time) - .into_boxed() - .single_value(); + if count > 0 { + tracing::info!("Reset {count} jobs"); + } - let opt = diesel::update(job_queue) - .filter(id.nullable().eq(id_query)) - .set(( - heartbeat.eq(timestamp), - status.eq(JobStatus::Running), - worker.eq(worker_id), - )) - .returning((id, job)) - .get_result(&mut conn) - .await - .optional() - .map_err(PostgresError::Diesel)?; - - if let Some((job_id, job_json)) = opt { - diesel::sql_query("UNLISTEN queue_status_channel;") - .execute(&mut conn) + // TODO: combine into 1 query + let opt = loop { + let id_opt = job_queue + .select(id) + .filter(status.eq(JobStatus::New).and(queue.eq(queue_name))) + .order(queue_time) + .limit(1) + .get_result::(&mut conn) .await + .optional() .map_err(PostgresError::Diesel)?; + let Some(id_val) = id_opt else { + break None; + }; + + let opt = diesel::update(job_queue) + .filter(id.eq(id_val)) + .filter(status.eq(JobStatus::New)) + .set(( + heartbeat.eq(timestamp), + status.eq(JobStatus::Running), + worker.eq(worker_id), + )) + .returning((id, job)) + .get_result(&mut conn) + .await + .optional() + .map_err(PostgresError::Diesel)?; + + if let Some(tup) = opt { + break Some(tup); + } + }; + + if let Some((job_id, job_json)) = opt { return Ok((JobId(job_id), job_json)); } - let _ = actix_rt::time::timeout(Duration::from_secs(5), notifier.notified()).await; - - diesel::sql_query("UNLISTEN queue_status_channel;") - .execute(&mut conn) - .await - .map_err(PostgresError::Diesel)?; - drop(conn); + if actix_rt::time::timeout(Duration::from_secs(5), notifier.notified()) + .await + .is_ok() + { + tracing::debug!("Notified"); + } else { + tracing::debug!("Timed out"); + } } } @@ -908,7 +933,7 @@ impl QueueRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::job_queue::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let timestamp = to_primitive(time::OffsetDateTime::now_utc()); @@ -935,7 +960,7 @@ impl QueueRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::job_queue::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(job_queue) .filter( @@ -957,7 +982,7 @@ impl StoreMigrationRepo for PostgresRepo { async fn is_continuing_migration(&self) -> Result { use schema::store_migrations::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let count = store_migrations .count() @@ -976,7 +1001,7 @@ impl StoreMigrationRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::store_migrations::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::insert_into(store_migrations) .values(( @@ -996,7 +1021,7 @@ impl StoreMigrationRepo for PostgresRepo { async fn is_migrated(&self, input_old_identifier: &Arc) -> Result { use schema::store_migrations::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let b = diesel::select(diesel::dsl::exists( store_migrations.filter(old_identifier.eq(input_old_identifier.as_ref())), @@ -1012,7 +1037,7 @@ impl StoreMigrationRepo for PostgresRepo { async fn clear(&self) -> Result<(), RepoError> { use schema::store_migrations::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(store_migrations) .execute(&mut conn) @@ -1029,7 +1054,7 @@ impl ProxyRepo for PostgresRepo { async fn relate_url(&self, input_url: Url, input_alias: Alias) -> Result<(), RepoError> { use schema::proxies::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::insert_into(proxies) .values((url.eq(input_url.as_str()), alias.eq(&input_alias))) @@ -1044,7 +1069,7 @@ impl ProxyRepo for PostgresRepo { async fn related(&self, input_url: Url) -> Result, RepoError> { use schema::proxies::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = proxies .select(alias) @@ -1061,7 +1086,7 @@ impl ProxyRepo for PostgresRepo { async fn remove_relation(&self, input_alias: Alias) -> Result<(), RepoError> { use schema::proxies::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(proxies) .filter(alias.eq(&input_alias)) @@ -1083,7 +1108,7 @@ impl AliasAccessRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::proxies::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let timestamp = to_primitive(timestamp); @@ -1104,7 +1129,7 @@ impl AliasAccessRepo for PostgresRepo { ) -> Result, RepoError> { use schema::proxies::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = proxies .select(accessed) @@ -1132,7 +1157,7 @@ impl AliasAccessRepo for PostgresRepo { Box::pin(async move { use schema::proxies::dsl::*; - let mut conn = inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = inner.get_connection().await?; let vec = proxies .select((accessed, alias)) @@ -1166,7 +1191,7 @@ impl VariantAccessRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::variants::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let timestamp = to_primitive(input_accessed); @@ -1188,7 +1213,7 @@ impl VariantAccessRepo for PostgresRepo { ) -> Result, RepoError> { use schema::variants::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let opt = variants .select(accessed) @@ -1216,7 +1241,7 @@ impl VariantAccessRepo for PostgresRepo { Box::pin(async move { use schema::variants::dsl::*; - let mut conn = inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = inner.get_connection().await?; let vec = variants .select((accessed, (hash, variant))) @@ -1281,7 +1306,7 @@ impl UploadRepo for PostgresRepo { async fn create_upload(&self) -> Result { use schema::uploads::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let uuid = diesel::insert_into(uploads) .default_values() @@ -1298,7 +1323,7 @@ impl UploadRepo for PostgresRepo { use schema::uploads::dsl::*; loop { - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::sql_query("LISTEN upload_completion_channel;") .execute(&mut conn) @@ -1315,28 +1340,25 @@ impl UploadRepo for PostgresRepo { .flatten(); if let Some(upload_result) = opt { - diesel::sql_query("UNLISTEN upload_completion_channel;") - .execute(&mut conn) - .await - .map_err(PostgresError::Diesel)?; - let upload_result: InnerUploadResult = serde_json::from_value(upload_result) .map_err(PostgresError::DeserializeUploadResult)?; return Ok(upload_result.into()); } - let _ = actix_rt::time::timeout( - Duration::from_secs(2), + drop(conn); + + if actix_rt::time::timeout( + Duration::from_secs(5), self.inner.upload_notifier.notified(), ) - .await; - - diesel::sql_query("UNLISTEN upload_completion_channel;") - .execute(&mut conn) - .await - .map_err(PostgresError::Diesel)?; - drop(conn); + .await + .is_ok() + { + tracing::debug!("Notified"); + } else { + tracing::debug!("Timed out"); + } } } @@ -1344,7 +1366,7 @@ impl UploadRepo for PostgresRepo { async fn claim(&self, upload_id: UploadId) -> Result<(), RepoError> { use schema::uploads::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; diesel::delete(uploads) .filter(id.eq(upload_id.id)) @@ -1363,7 +1385,7 @@ impl UploadRepo for PostgresRepo { ) -> Result<(), RepoError> { use schema::uploads::dsl::*; - let mut conn = self.inner.pool.get().await.map_err(PostgresError::Pool)?; + let mut conn = self.get_connection().await?; let upload_result: InnerUploadResult = upload_result.into(); let upload_result = diff --git a/src/repo/postgres/schema.rs b/src/repo/postgres/schema.rs index 964937b..72ba9aa 100644 --- a/src/repo/postgres/schema.rs +++ b/src/repo/postgres/schema.rs @@ -1,7 +1,7 @@ // @generated automatically by Diesel CLI. pub mod sql_types { - #[derive(diesel::sql_types::SqlType)] + #[derive(diesel::query_builder::QueryId, diesel::sql_types::SqlType)] #[diesel(postgres_type(name = "job_status"))] pub struct JobStatus; } From a43de122f9050b8ef7a978f7a6a3b792b7caa5b2 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 21:30:47 -0500 Subject: [PATCH 14/27] postgres: add already-claimed case, general: tracing paranoia --- src/backgrounded.rs | 28 ++++------ src/file.rs | 6 +- src/ingest.rs | 42 ++++++-------- src/ingest/hasher.rs | 7 +-- src/lib.rs | 59 ++++++++++---------- src/migrate_store.rs | 4 +- src/process.rs | 54 +++++++++--------- src/queue/process.rs | 2 +- src/repo/postgres.rs | 113 ++++++++++++++++++++++---------------- src/repo/sled.rs | 34 +++++------- src/repo_04/sled.rs | 4 +- src/store/object_store.rs | 2 +- src/stream.rs | 27 ++++++--- src/sync.rs | 38 +++++++++++++ src/tmp_file.rs | 3 +- 15 files changed, 234 insertions(+), 189 deletions(-) create mode 100644 src/sync.rs diff --git a/src/backgrounded.rs b/src/backgrounded.rs index 94839a1..a1aa87a 100644 --- a/src/backgrounded.rs +++ b/src/backgrounded.rs @@ -82,14 +82,12 @@ impl Drop for Backgrounded { let cleanup_span = tracing::info_span!(parent: &cleanup_parent_span, "Backgrounded cleanup Identifier", identifier = ?identifier); - tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { - actix_rt::spawn( - async move { - let _ = crate::queue::cleanup_identifier(&repo, &identifier).await; - } - .instrument(cleanup_span), - ) - }); + crate::sync::spawn( + async move { + let _ = crate::queue::cleanup_identifier(&repo, &identifier).await; + } + .instrument(cleanup_span), + ); } if let Some(upload_id) = self.upload_id { @@ -97,14 +95,12 @@ impl Drop for Backgrounded { let cleanup_span = tracing::info_span!(parent: &cleanup_parent_span, "Backgrounded cleanup Upload ID", upload_id = ?upload_id); - tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { - actix_rt::spawn( - async move { - let _ = repo.claim(upload_id).await; - } - .instrument(cleanup_span), - ) - }); + crate::sync::spawn( + async move { + let _ = repo.claim(upload_id).await; + } + .instrument(cleanup_span), + ); } } } diff --git a/src/file.rs b/src/file.rs index 6706ca7..bb753b4 100644 --- a/src/file.rs +++ b/src/file.rs @@ -446,15 +446,15 @@ mod io_uring { actix_rt::System::new().block_on(async move { let arbiter = actix_rt::Arbiter::new(); - let (tx, rx) = tokio::sync::oneshot::channel(); + let (tx, rx) = crate::sync::channel(1); arbiter.spawn(async move { - let handle = actix_rt::spawn($fut); + let handle = crate::sync::spawn($fut); let _ = tx.send(handle.await.unwrap()); }); - rx.await.unwrap() + rx.into_recv_async().await.unwrap() }) }; } diff --git a/src/ingest.rs b/src/ingest.rs index e0cd34b..75f27b3 100644 --- a/src/ingest.rs +++ b/src/ingest.rs @@ -217,14 +217,12 @@ impl Drop for Session { let cleanup_span = tracing::info_span!(parent: &cleanup_parent_span, "Session cleanup hash", hash = ?hash); - tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { - actix_rt::spawn( - async move { - let _ = crate::queue::cleanup_hash(&repo, hash).await; - } - .instrument(cleanup_span), - ) - }); + crate::sync::spawn( + async move { + let _ = crate::queue::cleanup_hash(&repo, hash).await; + } + .instrument(cleanup_span), + ); } if let Some(alias) = self.alias.take() { @@ -233,14 +231,12 @@ impl Drop for Session { let cleanup_span = tracing::info_span!(parent: &cleanup_parent_span, "Session cleanup alias", alias = ?alias); - tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { - actix_rt::spawn( - async move { - let _ = crate::queue::cleanup_alias(&repo, alias, token).await; - } - .instrument(cleanup_span), - ) - }); + crate::sync::spawn( + async move { + let _ = crate::queue::cleanup_alias(&repo, alias, token).await; + } + .instrument(cleanup_span), + ); } if let Some(identifier) = self.identifier.take() { @@ -248,14 +244,12 @@ impl Drop for Session { let cleanup_span = tracing::info_span!(parent: &cleanup_parent_span, "Session cleanup identifier", identifier = ?identifier); - tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { - actix_rt::spawn( - async move { - let _ = crate::queue::cleanup_identifier(&repo, &identifier).await; - } - .instrument(cleanup_span), - ) - }); + crate::sync::spawn( + async move { + let _ = crate::queue::cleanup_identifier(&repo, &identifier).await; + } + .instrument(cleanup_span), + ); } } } diff --git a/src/ingest/hasher.rs b/src/ingest/hasher.rs index b5b809c..6466e10 100644 --- a/src/ingest/hasher.rs +++ b/src/ingest/hasher.rs @@ -82,16 +82,15 @@ mod test { actix_rt::System::new().block_on(async move { let arbiter = actix_rt::Arbiter::new(); - let (tx, rx) = tracing::trace_span!(parent: None, "Create channel") - .in_scope(|| tokio::sync::oneshot::channel()); + let (tx, rx) = crate::sync::channel(1); arbiter.spawn(async move { - let handle = actix_rt::spawn($fut); + let handle = crate::sync::spawn($fut); let _ = tx.send(handle.await.unwrap()); }); - rx.await.unwrap() + rx.into_recv_async().await.unwrap() }) }; } diff --git a/src/lib.rs b/src/lib.rs index f55e7de..cbb3835 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,6 +26,7 @@ mod repo_04; mod serde_str; mod store; mod stream; +mod sync; mod tmp_file; mod validate; @@ -84,8 +85,9 @@ const DAYS: u32 = 24 * HOURS; const NOT_FOUND_KEY: &str = "404-alias"; static PROCESS_SEMAPHORE: Lazy = Lazy::new(|| { - tracing::trace_span!(parent: None, "Initialize semaphore") - .in_scope(|| Semaphore::new(num_cpus::get().saturating_sub(1).max(1))) + let permits = num_cpus::get().saturating_sub(1).max(1); + + crate::sync::bare_semaphore(permits) }); async fn ensure_details( @@ -1671,44 +1673,39 @@ fn spawn_cleanup(repo: ArcRepo, config: &Configuration) { return; } - tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { - actix_rt::spawn(async move { - let mut interval = actix_rt::time::interval(Duration::from_secs(30)); + crate::sync::spawn(async move { + let mut interval = actix_rt::time::interval(Duration::from_secs(30)); - loop { - interval.tick().await; + loop { + interval.tick().await; - if let Err(e) = queue::cleanup_outdated_variants(&repo).await { - tracing::warn!( - "Failed to spawn cleanup for outdated variants:{}", - format!("\n{e}\n{e:?}") - ); - } - - if let Err(e) = queue::cleanup_outdated_proxies(&repo).await { - tracing::warn!( - "Failed to spawn cleanup for outdated proxies:{}", - format!("\n{e}\n{e:?}") - ); - } + if let Err(e) = queue::cleanup_outdated_variants(&repo).await { + tracing::warn!( + "Failed to spawn cleanup for outdated variants:{}", + format!("\n{e}\n{e:?}") + ); } - }); - }) + + if let Err(e) = queue::cleanup_outdated_proxies(&repo).await { + tracing::warn!( + "Failed to spawn cleanup for outdated proxies:{}", + format!("\n{e}\n{e:?}") + ); + } + } + }); } fn spawn_workers(repo: ArcRepo, store: S, config: Configuration, process_map: ProcessMap) where S: Store + 'static, { - tracing::trace_span!(parent: None, "Spawn task").in_scope(|| { - actix_rt::spawn(queue::process_cleanup( - repo.clone(), - store.clone(), - config.clone(), - )) - }); - tracing::trace_span!(parent: None, "Spawn task") - .in_scope(|| actix_rt::spawn(queue::process_images(repo, store, process_map, config))); + crate::sync::spawn(queue::process_cleanup( + repo.clone(), + store.clone(), + config.clone(), + )); + crate::sync::spawn(queue::process_images(repo, store, process_map, config)); } async fn launch_file_store( diff --git a/src/migrate_store.rs b/src/migrate_store.rs index 78cb0b0..d2ce902 100644 --- a/src/migrate_store.rs +++ b/src/migrate_store.rs @@ -61,7 +61,7 @@ where tracing::warn!("Retrying migration +{failure_count}"); } - tokio::time::sleep(Duration::from_secs(3)).await; + actix_rt::time::sleep(Duration::from_secs(3)).await; } Ok(()) @@ -364,7 +364,7 @@ where tracing::warn!("Failed moving file. Retrying +{failure_count}"); } - tokio::time::sleep(Duration::from_secs(3)).await; + actix_rt::time::sleep(Duration::from_secs(3)).await; } } } diff --git a/src/process.rs b/src/process.rs index 9621444..1cb58d5 100644 --- a/src/process.rs +++ b/src/process.rs @@ -1,5 +1,6 @@ use actix_rt::task::JoinHandle; use actix_web::web::Bytes; +use flume::r#async::RecvFut; use std::{ future::Future, pin::Pin, @@ -10,7 +11,6 @@ use std::{ use tokio::{ io::{AsyncRead, AsyncWriteExt, ReadBuf}, process::{Child, ChildStdin, ChildStdout, Command}, - sync::oneshot::{channel, Receiver}, }; use tracing::{Instrument, Span}; @@ -73,7 +73,7 @@ struct DropHandle { pub(crate) struct ProcessRead { inner: I, - err_recv: Receiver, + err_recv: RecvFut<'static, std::io::Error>, err_closed: bool, #[allow(dead_code)] handle: DropHandle, @@ -206,39 +206,37 @@ impl Process { let stdin = child.stdin.take().expect("stdin exists"); let stdout = child.stdout.take().expect("stdout exists"); - let (tx, rx) = tracing::trace_span!(parent: None, "Create channel", %command) - .in_scope(channel::); + let (tx, rx) = crate::sync::channel::(1); + let rx = rx.into_recv_async(); let span = tracing::info_span!(parent: None, "Background process task", %command); span.follows_from(Span::current()); - let handle = tracing::trace_span!(parent: None, "Spawn task", %command).in_scope(|| { - actix_rt::spawn( - async move { - let child_fut = async { - (f)(stdin).await?; + let handle = crate::sync::spawn( + async move { + let child_fut = async { + (f)(stdin).await?; - child.wait().await - }; + child.wait().await + }; - let error = match actix_rt::time::timeout(timeout, child_fut).await { - Ok(Ok(status)) if status.success() => { - guard.disarm(); - return; - } - Ok(Ok(status)) => { - std::io::Error::new(std::io::ErrorKind::Other, StatusError(status)) - } - Ok(Err(e)) => e, - Err(_) => std::io::ErrorKind::TimedOut.into(), - }; + let error = match actix_rt::time::timeout(timeout, child_fut).await { + Ok(Ok(status)) if status.success() => { + guard.disarm(); + return; + } + Ok(Ok(status)) => { + std::io::Error::new(std::io::ErrorKind::Other, StatusError(status)) + } + Ok(Err(e)) => e, + Err(_) => std::io::ErrorKind::TimedOut.into(), + }; - let _ = tx.send(error); - let _ = child.kill().await; - } - .instrument(span), - ) - }); + let _ = tx.send(error); + let _ = child.kill().await; + } + .instrument(span), + ); let sleep = actix_rt::time::sleep(timeout); diff --git a/src/queue/process.rs b/src/queue/process.rs index 5454dc8..13e8b78 100644 --- a/src/queue/process.rs +++ b/src/queue/process.rs @@ -86,7 +86,7 @@ where let repo = repo.clone(); let media = media.clone(); - let error_boundary = actix_rt::spawn(async move { + let error_boundary = crate::sync::spawn(async move { let stream = store2 .to_stream(&ident, None, None) .await? diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index f63331d..addba2b 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -20,7 +20,8 @@ use diesel_async::{ AsyncConnection, AsyncPgConnection, RunQueryDsl, }; use tokio::sync::Notify; -use tokio_postgres::{AsyncMessage, Notification}; +use tokio_postgres::{tls::NoTlsStream, AsyncMessage, Connection, NoTls, Notification, Socket}; +use tracing::Instrument; use url::Url; use uuid::Uuid; @@ -64,7 +65,7 @@ async fn delegate_notifications(receiver: flume::Receiver, inner: inner .queue_notifications .entry(queue_name) - .or_insert_with(|| Arc::new(Notify::new())) + .or_insert_with(crate::sync::notify) .notify_waiters(); } "upload_completion_channel" => { @@ -134,12 +135,11 @@ impl PostgresError { impl PostgresRepo { pub(crate) async fn connect(postgres_url: Url) -> Result { - let (mut client, conn) = - tokio_postgres::connect(postgres_url.as_str(), tokio_postgres::tls::NoTls) - .await - .map_err(ConnectPostgresError::ConnectForMigration)?; + let (mut client, conn) = tokio_postgres::connect(postgres_url.as_str(), NoTls) + .await + .map_err(ConnectPostgresError::ConnectForMigration)?; - let handle = actix_rt::spawn(conn); + let handle = crate::sync::spawn(conn); embedded::migrations::runner() .run_async(&mut client) @@ -166,10 +166,12 @@ impl PostgresRepo { health_count: AtomicU64::new(0), pool, queue_notifications: DashMap::new(), - upload_notifier: Notify::new(), + upload_notifier: crate::sync::bare_notify(), }); - let notifications = Arc::new(actix_rt::spawn(delegate_notifications(rx, inner.clone()))); + let handle = crate::sync::spawn(delegate_notifications(rx, inner.clone())); + + let notifications = Arc::new(handle); Ok(PostgresRepo { inner, @@ -198,41 +200,53 @@ fn build_handler(sender: flume::Sender) -> ConfigFn { move |config: &str| -> BoxFuture<'_, ConnectionResult> { let sender = sender.clone(); - Box::pin(async move { - let (client, mut conn) = - tokio_postgres::connect(config, tokio_postgres::tls::NoTls) - .await - .map_err(|e| ConnectionError::BadConnection(e.to_string()))?; + let connect_span = tracing::trace_span!(parent: None, "connect future"); - // not very cash money (structured concurrency) of me - actix_rt::spawn(async move { - while let Some(res) = std::future::poll_fn(|cx| conn.poll_message(cx)).await { - match res { - Err(e) => { - tracing::error!("Database Connection {e:?}"); - return; - } - Ok(AsyncMessage::Notice(e)) => { - tracing::warn!("Database Notice {e:?}"); - } - Ok(AsyncMessage::Notification(notification)) => { - if sender.send_async(notification).await.is_err() { - tracing::warn!("Missed notification. Are we shutting down?"); - } - } - Ok(_) => { - tracing::warn!("Unhandled AsyncMessage!!! Please contact the developer of this application"); - } - } - } - }); + Box::pin( + async move { + let (client, conn) = + tokio_postgres::connect(config, tokio_postgres::tls::NoTls) + .await + .map_err(|e| ConnectionError::BadConnection(e.to_string()))?; - AsyncPgConnection::try_from(client).await - }) + // not very cash money (structured concurrency) of me + spawn_db_notification_task(sender, conn); + + AsyncPgConnection::try_from(client).await + } + .instrument(connect_span), + ) }, ) } +fn spawn_db_notification_task( + sender: flume::Sender, + mut conn: Connection, +) { + crate::sync::spawn(async move { + while let Some(res) = std::future::poll_fn(|cx| conn.poll_message(cx)).await { + match res { + Err(e) => { + tracing::error!("Database Connection {e:?}"); + return; + } + Ok(AsyncMessage::Notice(e)) => { + tracing::warn!("Database Notice {e:?}"); + } + Ok(AsyncMessage::Notification(notification)) => { + if sender.send_async(notification).await.is_err() { + tracing::warn!("Missed notification. Are we shutting down?"); + } + } + Ok(_) => { + tracing::warn!("Unhandled AsyncMessage!!! Please contact the developer of this application"); + } + } + } + }); +} + fn to_primitive(timestamp: time::OffsetDateTime) -> time::PrimitiveDateTime { let timestamp = timestamp.to_offset(time::UtcOffset::UTC); time::PrimitiveDateTime::new(timestamp.date(), timestamp.time()) @@ -849,7 +863,7 @@ impl QueueRepo for PostgresRepo { .inner .queue_notifications .entry(String::from(queue_name)) - .or_insert_with(|| Arc::new(Notify::new())) + .or_insert_with(crate::sync::notify) .clone(); diesel::sql_query("LISTEN queue_status_channel;") @@ -1330,20 +1344,27 @@ impl UploadRepo for PostgresRepo { .await .map_err(PostgresError::Diesel)?; - let opt = uploads + let nested_opt = uploads .select(result) .filter(id.eq(upload_id.id)) .get_result(&mut conn) .await .optional() - .map_err(PostgresError::Diesel)? - .flatten(); + .map_err(PostgresError::Diesel)?; - if let Some(upload_result) = opt { - let upload_result: InnerUploadResult = serde_json::from_value(upload_result) - .map_err(PostgresError::DeserializeUploadResult)?; + match nested_opt { + Some(opt) => { + if let Some(upload_result) = opt { + let upload_result: InnerUploadResult = + serde_json::from_value(upload_result) + .map_err(PostgresError::DeserializeUploadResult)?; - return Ok(upload_result.into()); + return Ok(upload_result.into()); + } + } + None => { + return Err(RepoError::AlreadyClaimed); + } } drop(conn); diff --git a/src/repo/sled.rs b/src/repo/sled.rs index 671e186..aa42e9a 100644 --- a/src/repo/sled.rs +++ b/src/repo/sled.rs @@ -29,9 +29,7 @@ macro_rules! b { ($self:ident.$ident:ident, $expr:expr) => {{ let $ident = $self.$ident.clone(); - let span = tracing::Span::current(); - - actix_rt::task::spawn_blocking(move || span.in_scope(|| $expr)) + crate::sync::spawn_blocking(move || $expr) .await .map_err(SledError::from) .map_err(RepoError::from)? @@ -175,7 +173,7 @@ impl SledRepo { let this = self.db.clone(); - actix_rt::task::spawn_blocking(move || { + crate::sync::spawn_blocking(move || { let export = this.export(); export_db.import(export); }) @@ -258,7 +256,7 @@ impl AliasAccessRepo for SledRepo { let alias_access = self.alias_access.clone(); let inverse_alias_access = self.inverse_alias_access.clone(); - let res = actix_rt::task::spawn_blocking(move || { + let res = crate::sync::spawn_blocking(move || { (&alias_access, &inverse_alias_access).transaction( |(alias_access, inverse_alias_access)| { if let Some(old) = alias_access.insert(alias.to_bytes(), &value_bytes)? { @@ -324,7 +322,7 @@ impl AliasAccessRepo for SledRepo { let alias_access = self.alias_access.clone(); let inverse_alias_access = self.inverse_alias_access.clone(); - let res = actix_rt::task::spawn_blocking(move || { + let res = crate::sync::spawn_blocking(move || { (&alias_access, &inverse_alias_access).transaction( |(alias_access, inverse_alias_access)| { if let Some(old) = alias_access.remove(alias.to_bytes())? { @@ -364,7 +362,7 @@ impl VariantAccessRepo for SledRepo { let variant_access = self.variant_access.clone(); let inverse_variant_access = self.inverse_variant_access.clone(); - let res = actix_rt::task::spawn_blocking(move || { + let res = crate::sync::spawn_blocking(move || { (&variant_access, &inverse_variant_access).transaction( |(variant_access, inverse_variant_access)| { if let Some(old) = variant_access.insert(&key, &value_bytes)? { @@ -434,7 +432,7 @@ impl VariantAccessRepo for SledRepo { let variant_access = self.variant_access.clone(); let inverse_variant_access = self.inverse_variant_access.clone(); - let res = actix_rt::task::spawn_blocking(move || { + let res = crate::sync::spawn_blocking(move || { (&variant_access, &inverse_variant_access).transaction( |(variant_access, inverse_variant_access)| { if let Some(old) = variant_access.remove(&key)? { @@ -678,7 +676,7 @@ impl QueueRepo for SledRepo { let queue = self.queue.clone(); let job_state = self.job_state.clone(); - let res = actix_rt::task::spawn_blocking(move || { + let res = crate::sync::spawn_blocking(move || { (&queue, &job_state).transaction(|(queue, job_state)| { let state = JobState::pending(); @@ -705,7 +703,7 @@ impl QueueRepo for SledRepo { .write() .unwrap() .entry(queue_name) - .or_insert_with(|| Arc::new(Notify::new())) + .or_insert_with(crate::sync::notify) .notify_one(); metrics_guard.disarm(); @@ -728,7 +726,7 @@ impl QueueRepo for SledRepo { let job_state = self.job_state.clone(); let span = tracing::Span::current(); - let opt = actix_rt::task::spawn_blocking(move || { + let opt = crate::sync::spawn_blocking(move || { let _guard = span.enter(); // Job IDs are generated with Uuid version 7 - defining their first bits as a // timestamp. Scanning a prefix should give us jobs in the order they were queued. @@ -802,9 +800,7 @@ impl QueueRepo for SledRepo { notify } else { let mut guard = self.queue_notifier.write().unwrap(); - let entry = guard - .entry(queue_name) - .or_insert_with(|| Arc::new(Notify::new())); + let entry = guard.entry(queue_name).or_insert_with(crate::sync::notify); Arc::clone(entry) }; @@ -823,7 +819,7 @@ impl QueueRepo for SledRepo { let job_state = self.job_state.clone(); - actix_rt::task::spawn_blocking(move || { + crate::sync::spawn_blocking(move || { if let Some(state) = job_state.get(&key)? { let new_state = JobState::running(worker_id); @@ -853,7 +849,7 @@ impl QueueRepo for SledRepo { let queue = self.queue.clone(); let job_state = self.job_state.clone(); - let res = actix_rt::task::spawn_blocking(move || { + let res = crate::sync::spawn_blocking(move || { (&queue, &job_state).transaction(|(queue, job_state)| { queue.remove(&key[..])?; job_state.remove(&key[..])?; @@ -1112,7 +1108,7 @@ impl HashRepo for SledRepo { None => (self.hashes_inverse.iter(), None), }; - actix_rt::task::spawn_blocking(move || { + crate::sync::spawn_blocking(move || { let page_iter = page_iter .keys() .rev() @@ -1164,7 +1160,7 @@ impl HashRepo for SledRepo { let page_iter = self.hashes_inverse.range(..=date_nanos); let prev_iter = Some(self.hashes_inverse.range(date_nanos..)); - actix_rt::task::spawn_blocking(move || { + crate::sync::spawn_blocking(move || { let page_iter = page_iter .keys() .rev() @@ -1292,7 +1288,7 @@ impl HashRepo for SledRepo { let hash_variant_identifiers = self.hash_variant_identifiers.clone(); - actix_rt::task::spawn_blocking(move || { + crate::sync::spawn_blocking(move || { hash_variant_identifiers .compare_and_swap(key, Option::<&[u8]>::None, Some(value.as_bytes())) .map(|res| res.map_err(|_| VariantAlreadyExists)) diff --git a/src/repo_04/sled.rs b/src/repo_04/sled.rs index 09a684a..07ccd9e 100644 --- a/src/repo_04/sled.rs +++ b/src/repo_04/sled.rs @@ -34,9 +34,7 @@ macro_rules! b { ($self:ident.$ident:ident, $expr:expr) => {{ let $ident = $self.$ident.clone(); - let span = tracing::Span::current(); - - actix_rt::task::spawn_blocking(move || span.in_scope(|| $expr)) + crate::sync::spawn_blocking(move || $expr) .await .map_err(SledError::from) .map_err(RepoError::from)? diff --git a/src/store/object_store.rs b/src/store/object_store.rs index 4c1aab7..324a0eb 100644 --- a/src/store/object_store.rs +++ b/src/store/object_store.rs @@ -277,7 +277,7 @@ impl Store for ObjectStore { let object_id2 = object_id.clone(); let upload_id2 = upload_id.clone(); - let handle = actix_rt::spawn( + let handle = crate::sync::spawn( async move { let response = this .create_upload_part_request( diff --git a/src/stream.rs b/src/stream.rs index 9833050..1309e15 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,5 +1,6 @@ use actix_rt::{task::JoinHandle, time::Sleep}; use actix_web::web::Bytes; +use flume::r#async::RecvStream; use futures_core::Stream; use std::{ future::Future, @@ -174,19 +175,25 @@ pin_project_lite::pin_project! { } } -enum IterStreamState { +enum IterStreamState +where + T: 'static, +{ New { iterator: I, buffer: usize, }, Running { handle: JoinHandle<()>, - receiver: tokio::sync::mpsc::Receiver, + receiver: RecvStream<'static, T>, }, Pending, } -pub(crate) struct IterStream { +pub(crate) struct IterStream +where + T: 'static, +{ state: IterStreamState, } @@ -287,14 +294,13 @@ where match std::mem::replace(&mut this.state, IterStreamState::Pending) { IterStreamState::New { iterator, buffer } => { - let (sender, receiver) = tracing::trace_span!(parent: None, "Create channel") - .in_scope(|| tokio::sync::mpsc::channel(buffer)); + let (sender, receiver) = crate::sync::channel(buffer); - let mut handle = actix_rt::task::spawn_blocking(move || { + let mut handle = crate::sync::spawn_blocking(move || { let iterator = iterator.into_iter(); for item in iterator { - if sender.blocking_send(item).is_err() { + if sender.send(item).is_err() { break; } } @@ -304,14 +310,17 @@ where return Poll::Ready(None); } - this.state = IterStreamState::Running { handle, receiver }; + this.state = IterStreamState::Running { + handle, + receiver: receiver.into_stream(), + }; self.poll_next(cx) } IterStreamState::Running { mut handle, mut receiver, - } => match Pin::new(&mut receiver).poll_recv(cx) { + } => match Pin::new(&mut receiver).poll_next(cx) { Poll::Ready(Some(item)) => { this.state = IterStreamState::Running { handle, receiver }; diff --git a/src/sync.rs b/src/sync.rs new file mode 100644 index 0000000..3111daf --- /dev/null +++ b/src/sync.rs @@ -0,0 +1,38 @@ +use std::sync::Arc; + +use tokio::sync::{Notify, Semaphore}; + +pub(crate) fn channel(bound: usize) -> (flume::Sender, flume::Receiver) { + tracing::trace_span!(parent: None, "make channel").in_scope(|| flume::bounded(bound)) +} + +pub(crate) fn notify() -> Arc { + Arc::new(bare_notify()) +} + +pub(crate) fn bare_notify() -> Notify { + tracing::trace_span!(parent: None, "make notifier").in_scope(Notify::new) +} + +pub(crate) fn bare_semaphore(permits: usize) -> Semaphore { + tracing::trace_span!(parent: None, "make semaphore").in_scope(|| Semaphore::new(permits)) +} + +pub(crate) fn spawn(future: F) -> actix_rt::task::JoinHandle +where + F: std::future::Future + 'static, + F::Output: 'static, +{ + tracing::trace_span!(parent: None, "spawn task").in_scope(|| actix_rt::spawn(future)) +} + +pub(crate) fn spawn_blocking(function: F) -> actix_rt::task::JoinHandle +where + F: FnOnce() -> Out + Send + 'static, + Out: Send + 'static, +{ + let outer_span = tracing::Span::current(); + + tracing::trace_span!(parent: None, "spawn blocking task") + .in_scope(|| actix_rt::task::spawn_blocking(move || outer_span.in_scope(function))) +} diff --git a/src/tmp_file.rs b/src/tmp_file.rs index 1335abc..edcb861 100644 --- a/src/tmp_file.rs +++ b/src/tmp_file.rs @@ -13,8 +13,7 @@ struct TmpFile(PathBuf); impl Drop for TmpFile { fn drop(&mut self) { - tracing::trace_span!(parent: None, "Spawn task") - .in_scope(|| actix_rt::spawn(tokio::fs::remove_file(self.0.clone()))); + crate::sync::spawn(tokio::fs::remove_file(self.0.clone())); } } From b9ec1f88badb62dde4931bf3c9ede69f0a93d818 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 21:32:31 -0500 Subject: [PATCH 15/27] Update tracing-opentelemetry, deps minor & point --- Cargo.lock | 73 +++++++++++++++++++++++++++--------------------------- Cargo.toml | 4 ++- 2 files changed, 40 insertions(+), 37 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fae5c89..922a62e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -77,7 +77,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e01ed3140b2f8d422c68afa1ed2e85d996ea619c988ac834d255db32138655cb" dependencies = [ "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -313,7 +313,7 @@ checksum = "bc00ceb34980c03614e35a3a4e218276a0a824e911d07651cd0d858a51e8c0f0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -500,7 +500,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -759,7 +759,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -771,7 +771,7 @@ dependencies = [ "diesel_table_macro_syntax", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -780,7 +780,7 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc5557efc453706fed5e4fa85006fe9817c224c3f480a34c7e5959fd700921c5" dependencies = [ - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -939,7 +939,7 @@ checksum = "89ca545a94061b6365f2c7355b4b32bd20df3ff95f02da9329b34ccc3bd6ee72" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -1377,9 +1377,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.6.2" +version = "2.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5486aed0026218e61b8a01d5fbd5a0a134649abb71a0e53b7bc088529dced86e" +checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" [[package]] name = "memoffset" @@ -1426,7 +1426,7 @@ checksum = "ddece26afd34c31585c74a4db0630c376df271c285d682d1e55012197830b6df" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -1537,9 +1537,9 @@ dependencies = [ [[package]] name = "object" -version = "0.32.0" +version = "0.32.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77ac5bbd07aea88c60a577a1ce218075ffd59208b2d7ca97adf9bfc5aeb21ebe" +checksum = "9cf5f9dd3933bd50a9e1f149ec995f39ae2c496d31fd772c1fd45ebc27e902b0" dependencies = [ "memchr", ] @@ -1767,7 +1767,7 @@ dependencies = [ "pest_meta", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -1881,7 +1881,7 @@ checksum = "4359fd9c9171ec6e8c62926d6faaf553a8dc3f64e1507e76da7911b4f6a04405" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -2141,18 +2141,18 @@ dependencies = [ "quote", "refinery-core", "regex", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] name = "regex" -version = "1.9.4" +version = "1.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12de2eff854e5fa4b1295edd650e227e9d8fb0c9e90b12e7f36d6a6811791a29" +checksum = "697061221ea1b4a94a624f67d0ae2bfe4e22b8a17b6a192afb11046542cc8c47" dependencies = [ "aho-corasick", "memchr", - "regex-automata 0.3.7", + "regex-automata 0.3.8", "regex-syntax 0.7.5", ] @@ -2167,9 +2167,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.3.7" +version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49530408a136e16e5b486e883fbb6ba058e8e4e8ae6621a77b048b314336e629" +checksum = "c2f401f4955220693b56f8ec66ee9c78abffd8d1c4f23dc41a23839eb88f0795" dependencies = [ "aho-corasick", "memchr", @@ -2472,7 +2472,7 @@ checksum = "4eca7ac642d82aa35b60049a6eccb4be6be75e599bd2e9adb5f875a737654af2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -2676,9 +2676,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.29" +version = "2.0.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c324c494eba9d92503e6f1ef2e6df781e78f6a7705a0202d9801b198807d518a" +checksum = "718fa2415bcb8d8bd775917a1bf12a7931b6dfa890753378538118181e0cb398" dependencies = [ "proc-macro2", "quote", @@ -2702,22 +2702,22 @@ dependencies = [ [[package]] name = "thiserror" -version = "1.0.47" +version = "1.0.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97a802ec30afc17eee47b2855fc72e0c4cd62be9b4efe6591edde0ec5bd68d8f" +checksum = "9d6d7a740b8a666a7e828dd00da9c0dc290dff53154ea77ac109281de90589b7" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.47" +version = "1.0.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6bb623b56e39ab7dcd4b1b98bb6c8f8d907ed255b18de254088016b27a8ee19b" +checksum = "49922ecae66cc8a249b77e68d1d0623c1b2c514f0060c27cdc68bd62a1219d35" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -2811,7 +2811,7 @@ checksum = "630bdcf245f78637c13ec01ffae6187cca34625e8c63150d424b59e55af2675e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -3009,8 +3009,7 @@ dependencies = [ [[package]] name = "tracing-actix-web" version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c0b08ce08cbde6a96fc1e4ebb8132053e53ec7a5cd27eef93ede6b73ebbda06" +source = "git+https://github.com/asonix/tracing-actix-web?branch=asonix/tracing-opentelemetry-021#0b337cc17fb88efc76d913751eee3ac66e4cd27f" dependencies = [ "actix-web", "opentelemetry", @@ -3028,7 +3027,7 @@ checksum = "5f4f31f56159e98206da9efd823404b79b6ef3143b4a7ab76e67b1751b25a4ab" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -3074,12 +3073,14 @@ dependencies = [ [[package]] name = "tracing-opentelemetry" -version = "0.20.0" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc09e402904a5261e42cf27aea09ccb7d5318c6717a9eec3d8e2e65c56b18f19" +checksum = "75327c6b667828ddc28f5e3f169036cb793c3f588d83bf0f262a7f062ffed3c8" dependencies = [ "once_cell", "opentelemetry", + "opentelemetry_sdk", + "smallvec", "tracing", "tracing-core", "tracing-log", @@ -3264,7 +3265,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", "wasm-bindgen-shared", ] @@ -3298,7 +3299,7 @@ checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", "wasm-bindgen-backend", "wasm-bindgen-shared", ] diff --git a/Cargo.toml b/Cargo.toml index abf9469..c07f167 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,7 +71,7 @@ tracing = "0.1.15" tracing-error = "0.2.0" tracing-futures = "0.2.4" tracing-log = "0.1.2" -tracing-opentelemetry = "0.20" +tracing-opentelemetry = "0.21" tracing-subscriber = { version = "0.3.0", features = [ "ansi", "env-filter", @@ -87,3 +87,5 @@ uuid = { version = "1", features = ["serde", "std", "v4", "v7"] } version = "0.7.5" default-features = false features = ["opentelemetry_0_20"] +git = "https://github.com/asonix/tracing-actix-web" +branch = "asonix/tracing-opentelemetry-021" From e961bdc3313befa04ea21cce35a5179f19fb99ff Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 21:32:54 -0500 Subject: [PATCH 16/27] Increase opentelemetry output --- dev.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev.toml b/dev.toml index 6c72d3e..c5df2cc 100644 --- a/dev.toml +++ b/dev.toml @@ -13,7 +13,7 @@ buffer_capacity = 102400 [tracing.opentelemetry] url = 'http://127.0.0.1:4317' service_name = 'pict-rs' -targets = 'info' +targets = 'info,pict_rs=debug' [old_repo] path = 'data/sled-repo-local' From 00aa00c55caa2be5af392ea05492754344dc9349 Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 21:58:31 -0500 Subject: [PATCH 17/27] postgres: Enable 'previous page' when less than full limit fit on previous page --- src/repo/postgres.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index addba2b..cf36b9c 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -354,11 +354,11 @@ impl HashRepo for PostgresRepo { .or_filter(created_at.eq(timestamp).and(hash.gt(&bound_hash))) .order(created_at) .then_order_by(hash) - .offset(limit.saturating_sub(1) as i64) - .first::(&mut conn) + .limit(limit as i64) + .get_results::(&mut conn) .await - .optional() - .map_err(PostgresError::Diesel)?; + .map_err(PostgresError::Diesel)? + .pop(); (page, prev) } else { From ac9777782a7f8570a5e758ca0cc06e73b51e47ed Mon Sep 17 00:00:00 2001 From: asonix Date: Sun, 3 Sep 2023 22:02:33 -0500 Subject: [PATCH 18/27] Add script for updating diesel schema --- scripts/update-schema.sh | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100755 scripts/update-schema.sh diff --git a/scripts/update-schema.sh b/scripts/update-schema.sh new file mode 100755 index 0000000..f22c84a --- /dev/null +++ b/scripts/update-schema.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash + +diesel \ + --database-url 'postgres://pictrs:1234@localhost:5432/pictrs' \ + print-schema \ + --custom-type-derives "diesel::query_builder::QueryId" \ + > src/repo/postgres/schema.rs From 37448722de996c245ce52bf63b887041e5c4318f Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 16:19:46 -0500 Subject: [PATCH 19/27] postgres: Rework job & upload notifications (more) postgres: Add metrics to job push & pop, upload wait sled: add upload wait metrics --- src/details.rs | 1 + src/discover/exiftool.rs | 1 + src/discover/magick.rs | 2 + src/repo.rs | 10 +- src/repo/metrics.rs | 74 ++++++ src/repo/postgres.rs | 225 +++++++++++++++--- .../migrations/V0006__create_queue.rs | 2 +- src/repo/sled.rs | 63 +---- 8 files changed, 283 insertions(+), 95 deletions(-) create mode 100644 src/repo/metrics.rs diff --git a/src/details.rs b/src/details.rs index d632e32..f5b9a86 100644 --- a/src/details.rs +++ b/src/details.rs @@ -103,6 +103,7 @@ impl Details { )) } + #[tracing::instrument(level = "DEBUG")] pub(crate) async fn from_store( store: &S, identifier: &Arc, diff --git a/src/discover/exiftool.rs b/src/discover/exiftool.rs index 1530eac..672ce10 100644 --- a/src/discover/exiftool.rs +++ b/src/discover/exiftool.rs @@ -9,6 +9,7 @@ use crate::{ use super::Discovery; +#[tracing::instrument(level = "DEBUG", skip_all)] pub(super) async fn check_reorient( Discovery { input, diff --git a/src/discover/magick.rs b/src/discover/magick.rs index 554c812..48c33fd 100644 --- a/src/discover/magick.rs +++ b/src/discover/magick.rs @@ -97,6 +97,7 @@ pub(super) async fn confirm_bytes( .await } +#[tracing::instrument(level = "DEBUG", skip(f))] async fn count_avif_frames(f: F, timeout: u64) -> Result where F: FnOnce(crate::file::File) -> Fut, @@ -147,6 +148,7 @@ where Ok(lines) } +#[tracing::instrument(level = "DEBUG", skip(f))] async fn discover_file(f: F, timeout: u64) -> Result where F: FnOnce(crate::file::File) -> Fut, diff --git a/src/repo.rs b/src/repo.rs index 0c8c312..dd32c19 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -1,3 +1,9 @@ +mod alias; +mod delete_token; +mod hash; +mod metrics; +mod migrate; + use crate::{ config, details::Details, @@ -9,10 +15,6 @@ use std::{fmt::Debug, sync::Arc}; use url::Url; use uuid::Uuid; -mod alias; -mod delete_token; -mod hash; -mod migrate; pub(crate) mod postgres; pub(crate) mod sled; diff --git a/src/repo/metrics.rs b/src/repo/metrics.rs new file mode 100644 index 0000000..3c07efa --- /dev/null +++ b/src/repo/metrics.rs @@ -0,0 +1,74 @@ +use std::time::Instant; + +pub(super) struct PushMetricsGuard { + queue: &'static str, + armed: bool, +} + +pub(super) struct PopMetricsGuard { + queue: &'static str, + start: Instant, + armed: bool, +} + +pub(super) struct WaitMetricsGuard { + start: Instant, + armed: bool, +} + +impl PushMetricsGuard { + pub(super) fn guard(queue: &'static str) -> Self { + Self { queue, armed: true } + } + + pub(super) fn disarm(mut self) { + self.armed = false; + } +} + +impl PopMetricsGuard { + pub(super) fn guard(queue: &'static str) -> Self { + Self { + queue, + start: Instant::now(), + armed: true, + } + } + + pub(super) fn disarm(mut self) { + self.armed = false; + } +} + +impl WaitMetricsGuard { + pub(super) fn guard() -> Self { + Self { + start: Instant::now(), + armed: true, + } + } + + pub(super) fn disarm(mut self) { + self.armed = false; + } +} + +impl Drop for PushMetricsGuard { + fn drop(&mut self) { + metrics::increment_counter!("pict-rs.queue.push", "completed" => (!self.armed).to_string(), "queue" => self.queue); + } +} + +impl Drop for PopMetricsGuard { + fn drop(&mut self) { + metrics::histogram!("pict-rs.queue.pop.duration", self.start.elapsed().as_secs_f64(), "completed" => (!self.armed).to_string(), "queue" => self.queue); + metrics::increment_counter!("pict-rs.queue.pop", "completed" => (!self.armed).to_string(), "queue" => self.queue); + } +} + +impl Drop for WaitMetricsGuard { + fn drop(&mut self) { + metrics::histogram!("pict-rs.upload.wait.duration", self.start.elapsed().as_secs_f64(), "completed" => (!self.armed).to_string()); + metrics::increment_counter!("pict-rs.upload.wait", "completed" => (!self.armed).to_string()); + } +} diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index cf36b9c..659eb6a 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -3,11 +3,12 @@ mod job_status; mod schema; use std::{ + collections::{BTreeSet, VecDeque}, sync::{ atomic::{AtomicU64, Ordering}, - Arc, + Arc, Weak, }, - time::Duration, + time::{Duration, Instant}, }; use dashmap::DashMap; @@ -35,6 +36,7 @@ use crate::{ use self::job_status::JobStatus; use super::{ + metrics::{PopMetricsGuard, PushMetricsGuard, WaitMetricsGuard}, Alias, AliasAccessRepo, AliasAlreadyExists, AliasRepo, BaseRepo, DeleteToken, DetailsRepo, FullRepo, Hash, HashAlreadyExists, HashPage, HashRepo, JobId, OrderedHash, ProxyRepo, QueueRepo, RepoError, SettingsRepo, StoreMigrationRepo, UploadId, UploadRepo, UploadResult, @@ -52,35 +54,24 @@ struct Inner { health_count: AtomicU64, pool: Pool, queue_notifications: DashMap>, - upload_notifier: Notify, + upload_notifications: DashMap>, } -async fn delegate_notifications(receiver: flume::Receiver, inner: Arc) { - while let Ok(notification) = receiver.recv_async().await { - match notification.channel() { - "queue_status_channel" => { - // new job inserted for queue - let queue_name = notification.payload().to_string(); +struct UploadInterest { + inner: Arc, + interest: Option>, + upload_id: UploadId, +} - inner - .queue_notifications - .entry(queue_name) - .or_insert_with(crate::sync::notify) - .notify_waiters(); - } - "upload_completion_channel" => { - inner.upload_notifier.notify_waiters(); - } - channel => { - tracing::info!( - "Unhandled postgres notification: {channel}: {}", - notification.payload() - ); - } - } - } +struct JobNotifierState<'a> { + inner: &'a Inner, + capacity: usize, + jobs: BTreeSet, + jobs_ordered: VecDeque, +} - tracing::warn!("Notification delegator shutting down"); +struct UploadNotifierState<'a> { + inner: &'a Inner, } #[derive(Debug, thiserror::Error)] @@ -166,7 +157,7 @@ impl PostgresRepo { health_count: AtomicU64::new(0), pool, queue_notifications: DashMap::new(), - upload_notifier: crate::sync::bare_notify(), + upload_notifications: DashMap::new(), }); let handle = crate::sync::spawn(delegate_notifications(rx, inner.clone())); @@ -184,10 +175,125 @@ impl PostgresRepo { } } +struct GetConnectionMetricsGuard { + start: Instant, + armed: bool, +} + +impl GetConnectionMetricsGuard { + fn guard() -> Self { + GetConnectionMetricsGuard { + start: Instant::now(), + armed: true, + } + } + + fn disarm(mut self) { + self.armed = false; + } +} + +impl Drop for GetConnectionMetricsGuard { + fn drop(&mut self) { + metrics::increment_counter!("pict-rs.postgres.pool.get.end", "completed" => (!self.armed).to_string()); + metrics::histogram!("pict-rs.postgres.pool.get.duration", self.start.elapsed().as_secs_f64(), "completed" => (!self.armed).to_string()); + } +} + impl Inner { - #[tracing::instrument(level = "DEBUG", skip(self))] + #[tracing::instrument(level = "TRACE", skip(self))] async fn get_connection(&self) -> Result, PostgresError> { - self.pool.get().await.map_err(PostgresError::Pool) + let guard = GetConnectionMetricsGuard::guard(); + let res = self.pool.get().await.map_err(PostgresError::Pool); + guard.disarm(); + res + } + + fn interest(self: &Arc, upload_id: UploadId) -> UploadInterest { + let notify = crate::sync::notify(); + + self.upload_notifications + .insert(upload_id, Arc::downgrade(¬ify)); + + UploadInterest { + inner: self.clone(), + interest: Some(notify), + upload_id, + } + } +} + +impl UploadInterest { + async fn notified_timeout(&self, timeout: Duration) -> Result<(), tokio::time::error::Elapsed> { + actix_rt::time::timeout( + timeout, + self.interest.as_ref().expect("interest exists").notified(), + ) + .await + } +} + +impl Drop for UploadInterest { + fn drop(&mut self) { + if let Some(interest) = self.interest.take() { + if Arc::into_inner(interest).is_some() { + self.inner.upload_notifications.remove(&self.upload_id); + } + } + } +} + +impl<'a> JobNotifierState<'a> { + fn handle(&mut self, payload: &str) { + let Some((job_id, queue_name)) = payload.split_once(' ') else { + tracing::warn!("Invalid queue payload {payload}"); + return; + }; + + let Ok(job_id) = job_id.parse::().map(JobId) else { + tracing::warn!("Invalid job ID {job_id}"); + return; + }; + + if !self.jobs.insert(job_id) { + // duplicate job + return; + } + + self.jobs_ordered.push_back(job_id); + + if self.jobs_ordered.len() > self.capacity { + if let Some(job_id) = self.jobs_ordered.pop_front() { + self.jobs.remove(&job_id); + } + } + + self.inner + .queue_notifications + .entry(queue_name.to_string()) + .or_insert_with(crate::sync::notify) + .notify_one(); + + metrics::increment_counter!("pict-rs.postgres.job-notifier.notified", "queue" => queue_name.to_string()); + } +} + +impl<'a> UploadNotifierState<'a> { + fn handle(&self, payload: &str) { + let Ok(upload_id) = payload.parse::() else { + tracing::warn!("Invalid upload id {payload}"); + return; + }; + + if let Some(notifier) = self + .inner + .upload_notifications + .get(&upload_id) + .and_then(|weak| weak.upgrade()) + { + notifier.notify_waiters(); + metrics::increment_counter!("pict-rs.postgres.upload-notifier.notified"); + } } } @@ -195,6 +301,44 @@ type BoxFuture<'a, T> = std::pin::Pin + type ConfigFn = Box BoxFuture<'_, ConnectionResult> + Send + Sync + 'static>; +async fn delegate_notifications(receiver: flume::Receiver, inner: Arc) { + let parallelism = std::thread::available_parallelism() + .map(|u| u.into()) + .unwrap_or(1_usize); + + let mut job_notifier_state = JobNotifierState { + inner: &inner, + capacity: parallelism * 8, + jobs: BTreeSet::new(), + jobs_ordered: VecDeque::new(), + }; + + let upload_notifier_state = UploadNotifierState { inner: &inner }; + + while let Ok(notification) = receiver.recv_async().await { + metrics::increment_counter!("pict-rs.postgres.notification"); + + match notification.channel() { + "queue_status_channel" => { + // new job inserted for queue + job_notifier_state.handle(notification.payload()); + } + "upload_completion_channel" => { + // new upload finished + upload_notifier_state.handle(notification.payload()); + } + channel => { + tracing::info!( + "Unhandled postgres notification: {channel}: {}", + notification.payload() + ); + } + } + } + + tracing::warn!("Notification delegator shutting down"); +} + fn build_handler(sender: flume::Sender) -> ConfigFn { Box::new( move |config: &str| -> BoxFuture<'_, ConnectionResult> { @@ -834,6 +978,8 @@ impl QueueRepo for PostgresRepo { queue_name: &'static str, job_json: serde_json::Value, ) -> Result { + let guard = PushMetricsGuard::guard(queue_name); + use schema::job_queue::dsl::*; let mut conn = self.get_connection().await?; @@ -845,6 +991,8 @@ impl QueueRepo for PostgresRepo { .await .map_err(PostgresError::Diesel)?; + guard.disarm(); + Ok(JobId(job_id)) } @@ -854,6 +1002,8 @@ impl QueueRepo for PostgresRepo { queue_name: &'static str, worker_id: Uuid, ) -> Result<(JobId, serde_json::Value), RepoError> { + let guard = PopMetricsGuard::guard(queue_name); + use schema::job_queue::dsl::*; loop { @@ -923,6 +1073,7 @@ impl QueueRepo for PostgresRepo { }; if let Some((job_id, job_json)) = opt { + guard.disarm(); return Ok((JobId(job_id), job_json)); } @@ -1334,9 +1485,14 @@ impl UploadRepo for PostgresRepo { #[tracing::instrument(level = "DEBUG", skip(self))] async fn wait(&self, upload_id: UploadId) -> Result { + let guard = WaitMetricsGuard::guard(); use schema::uploads::dsl::*; + let interest = self.inner.interest(upload_id); + loop { + let interest_future = interest.notified_timeout(Duration::from_secs(5)); + let mut conn = self.get_connection().await?; diesel::sql_query("LISTEN upload_completion_channel;") @@ -1359,6 +1515,7 @@ impl UploadRepo for PostgresRepo { serde_json::from_value(upload_result) .map_err(PostgresError::DeserializeUploadResult)?; + guard.disarm(); return Ok(upload_result.into()); } } @@ -1369,13 +1526,7 @@ impl UploadRepo for PostgresRepo { drop(conn); - if actix_rt::time::timeout( - Duration::from_secs(5), - self.inner.upload_notifier.notified(), - ) - .await - .is_ok() - { + if interest_future.await.is_ok() { tracing::debug!("Notified"); } else { tracing::debug!("Timed out"); diff --git a/src/repo/postgres/migrations/V0006__create_queue.rs b/src/repo/postgres/migrations/V0006__create_queue.rs index f9a96e0..cfabb83 100644 --- a/src/repo/postgres/migrations/V0006__create_queue.rs +++ b/src/repo/postgres/migrations/V0006__create_queue.rs @@ -34,7 +34,7 @@ CREATE OR REPLACE FUNCTION queue_status_notify() RETURNS trigger AS $$ BEGIN - PERFORM pg_notify('queue_status_channel', NEW.queue::text); + PERFORM pg_notify('queue_status_channel', NEW.id::text || ' ' || NEW.queue::text); RETURN NEW; END; $$ LANGUAGE plpgsql; diff --git a/src/repo/sled.rs b/src/repo/sled.rs index aa42e9a..77eb054 100644 --- a/src/repo/sled.rs +++ b/src/repo/sled.rs @@ -12,17 +12,18 @@ use std::{ atomic::{AtomicU64, Ordering}, Arc, RwLock, }, - time::Instant, }; use tokio::sync::Notify; use url::Url; use uuid::Uuid; use super::{ - hash::Hash, Alias, AliasAccessRepo, AliasAlreadyExists, AliasRepo, BaseRepo, DeleteToken, - Details, DetailsRepo, FullRepo, HashAlreadyExists, HashPage, HashRepo, JobId, OrderedHash, - ProxyRepo, QueueRepo, RepoError, SettingsRepo, StoreMigrationRepo, UploadId, UploadRepo, - UploadResult, VariantAccessRepo, VariantAlreadyExists, + hash::Hash, + metrics::{PopMetricsGuard, PushMetricsGuard, WaitMetricsGuard}, + Alias, AliasAccessRepo, AliasAlreadyExists, AliasRepo, BaseRepo, DeleteToken, Details, + DetailsRepo, FullRepo, HashAlreadyExists, HashPage, HashRepo, JobId, OrderedHash, ProxyRepo, + QueueRepo, RepoError, SettingsRepo, StoreMigrationRepo, UploadId, UploadRepo, UploadResult, + VariantAccessRepo, VariantAlreadyExists, }; macro_rules! b { @@ -490,54 +491,6 @@ impl From for UploadResult { } } -struct PushMetricsGuard { - queue: &'static str, - armed: bool, -} - -struct PopMetricsGuard { - queue: &'static str, - start: Instant, - armed: bool, -} - -impl PushMetricsGuard { - fn guard(queue: &'static str) -> Self { - Self { queue, armed: true } - } - - fn disarm(mut self) { - self.armed = false; - } -} - -impl PopMetricsGuard { - fn guard(queue: &'static str) -> Self { - Self { - queue, - start: Instant::now(), - armed: true, - } - } - - fn disarm(mut self) { - self.armed = false; - } -} - -impl Drop for PushMetricsGuard { - fn drop(&mut self) { - metrics::increment_counter!("pict-rs.queue.push", "completed" => (!self.armed).to_string(), "queue" => self.queue); - } -} - -impl Drop for PopMetricsGuard { - fn drop(&mut self) { - metrics::histogram!("pict-rs.queue.pop.duration", self.start.elapsed().as_secs_f64(), "completed" => (!self.armed).to_string(), "queue" => self.queue); - metrics::increment_counter!("pict-rs.queue.pop", "completed" => (!self.armed).to_string(), "queue" => self.queue); - } -} - #[async_trait::async_trait(?Send)] impl UploadRepo for SledRepo { #[tracing::instrument(level = "trace", skip(self))] @@ -551,6 +504,7 @@ impl UploadRepo for SledRepo { #[tracing::instrument(skip(self))] async fn wait(&self, upload_id: UploadId) -> Result { + let guard = WaitMetricsGuard::guard(); let mut subscriber = self.uploads.watch_prefix(upload_id.as_bytes()); let bytes = upload_id.as_bytes().to_vec(); @@ -560,6 +514,7 @@ impl UploadRepo for SledRepo { if bytes != b"1" { let result: InnerUploadResult = serde_json::from_slice(&bytes).map_err(SledError::UploadResult)?; + guard.disarm(); return Ok(result.into()); } } else { @@ -575,6 +530,8 @@ impl UploadRepo for SledRepo { if value != b"1" { let result: InnerUploadResult = serde_json::from_slice(&value).map_err(SledError::UploadResult)?; + + guard.disarm(); return Ok(result.into()); } } From 33615672aed2d3d8f697b283bf24d49d3d8188e0 Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 16:28:20 -0500 Subject: [PATCH 20/27] Enable prometheus metrics in dev --- dev.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dev.toml b/dev.toml index c5df2cc..691b1fb 100644 --- a/dev.toml +++ b/dev.toml @@ -15,6 +15,9 @@ url = 'http://127.0.0.1:4317' service_name = 'pict-rs' targets = 'info,pict_rs=debug' +[metrics] +prometheus_address = "127.0.0.1:8070" + [old_repo] path = 'data/sled-repo-local' cache_capacity = 67108864 From a4b1ab7dfb0907fca35272cfc47867f6efcc323c Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 21:51:27 -0500 Subject: [PATCH 21/27] Instrument postgres db calls --- src/error.rs | 28 +++++- src/future.rs | 75 ++++++++++++++ src/init_tracing.rs | 7 +- src/lib.rs | 4 +- src/queue.rs | 15 ++- src/queue/cleanup.rs | 3 +- src/queue/process.rs | 3 +- src/repo.rs | 10 +- src/repo/postgres.rs | 234 +++++++++++++++++++++++++++++++++++++++---- src/store.rs | 8 ++ 10 files changed, 350 insertions(+), 37 deletions(-) create mode 100644 src/future.rs diff --git a/src/error.rs b/src/error.rs index de2d9a3..f4ff1ed 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use actix_web::{http::StatusCode, HttpResponse, ResponseError}; use color_eyre::Report; @@ -5,6 +7,8 @@ use crate::error_code::ErrorCode; pub(crate) struct Error { inner: color_eyre::Report, + debug: Arc, + display: Arc, } impl Error { @@ -21,17 +25,21 @@ impl Error { .map(|e| e.error_code()) .unwrap_or(ErrorCode::UNKNOWN_ERROR) } + + pub(crate) fn is_disconnected(&self) -> bool { + self.kind().map(|e| e.is_disconnected()).unwrap_or(false) + } } impl std::fmt::Debug for Error { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Debug::fmt(&self.inner, f) + f.write_str(&self.debug) } } impl std::fmt::Display for Error { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(&self.inner, f) + f.write_str(&self.display) } } @@ -46,8 +54,14 @@ where UploadError: From, { fn from(error: T) -> Self { + let inner = Report::from(UploadError::from(error)); + let debug = Arc::from(format!("{inner:?}")); + let display = Arc::from(format!("{inner}")); + Error { - inner: Report::from(UploadError::from(error)), + inner, + debug, + display, } } } @@ -172,6 +186,14 @@ impl UploadError { Self::Timeout(_) => ErrorCode::STREAM_TOO_SLOW, } } + + const fn is_disconnected(&self) -> bool { + match self { + Self::Repo(e) => e.is_disconnected(), + Self::Store(s) => s.is_disconnected(), + _ => false, + } + } } impl From for UploadError { diff --git a/src/future.rs b/src/future.rs new file mode 100644 index 0000000..aea858d --- /dev/null +++ b/src/future.rs @@ -0,0 +1,75 @@ +use std::{ + future::Future, + time::{Duration, Instant}, +}; + +pub(crate) type LocalBoxFuture<'a, T> = std::pin::Pin + 'a>>; + +pub(crate) trait WithTimeout: Future { + fn with_timeout(self, duration: Duration) -> actix_rt::time::Timeout + where + Self: Sized, + { + actix_rt::time::timeout(duration, self) + } +} + +pub(crate) trait WithMetrics: Future { + fn with_metrics(self, name: &'static str) -> MetricsFuture + where + Self: Sized, + { + MetricsFuture { + future: self, + metrics: Metrics { + name, + start: Instant::now(), + complete: false, + }, + } + } +} + +impl WithMetrics for F where F: Future {} +impl WithTimeout for F where F: Future {} + +pin_project_lite::pin_project! { + pub(crate) struct MetricsFuture { + #[pin] + future: F, + + metrics: Metrics, + } +} + +struct Metrics { + name: &'static str, + start: Instant, + complete: bool, +} + +impl Future for MetricsFuture +where + F: Future, +{ + type Output = F::Output; + + fn poll( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll { + let this = self.project(); + + let out = std::task::ready!(this.future.poll(cx)); + + this.metrics.complete = true; + + std::task::Poll::Ready(out) + } +} + +impl Drop for Metrics { + fn drop(&mut self) { + metrics::histogram!(self.name, self.start.elapsed().as_secs_f64(), "complete" => self.complete.to_string()); + } +} diff --git a/src/init_tracing.rs b/src/init_tracing.rs index 62bfe74..5750d8c 100644 --- a/src/init_tracing.rs +++ b/src/init_tracing.rs @@ -8,9 +8,7 @@ use opentelemetry_otlp::WithExportConfig; use tracing::subscriber::set_global_default; use tracing_error::ErrorLayer; use tracing_log::LogTracer; -use tracing_subscriber::{ - fmt::format::FmtSpan, layer::SubscriberExt, registry::LookupSpan, Layer, Registry, -}; +use tracing_subscriber::{layer::SubscriberExt, registry::LookupSpan, Layer, Registry}; pub(super) fn init_tracing(tracing: &Tracing) -> color_eyre::Result<()> { color_eyre::install()?; @@ -19,8 +17,7 @@ pub(super) fn init_tracing(tracing: &Tracing) -> color_eyre::Result<()> { opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new()); - let format_layer = - tracing_subscriber::fmt::layer().with_span_events(FmtSpan::NEW | FmtSpan::CLOSE); + let format_layer = tracing_subscriber::fmt::layer(); match tracing.logging.format { LogFormat::Compact => with_format(format_layer.compact(), tracing), diff --git a/src/lib.rs b/src/lib.rs index cbb3835..ad18041 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,7 @@ mod exiftool; mod ffmpeg; mod file; mod formats; +mod future; mod generate; mod ingest; mod init_tracing; @@ -53,8 +54,8 @@ use std::{ time::{Duration, SystemTime}, }; use tokio::sync::Semaphore; +use tracing::Instrument; use tracing_actix_web::TracingLogger; -use tracing_futures::Instrument; use self::{ backgrounded::Backgrounded, @@ -1550,6 +1551,7 @@ async fn identifier( }))) } +#[tracing::instrument(skip(repo, store))] async fn healthz( repo: web::Data, store: web::Data, diff --git a/src/queue.rs b/src/queue.rs index a7d3f25..a3cbe15 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -3,14 +3,13 @@ use crate::{ config::Configuration, error::{Error, UploadError}, formats::InputProcessableFormat, + future::LocalBoxFuture, repo::{Alias, DeleteToken, FullRepo, Hash, JobId, UploadId}, serde_str::Serde, store::Store, }; use std::{ - future::Future, path::PathBuf, - pin::Pin, sync::Arc, time::{Duration, Instant}, }; @@ -179,8 +178,6 @@ pub(crate) async fn process_images( .await } -type LocalBoxFuture<'a, T> = Pin + 'a>>; - async fn process_jobs( repo: &Arc, store: &S, @@ -205,6 +202,11 @@ async fn process_jobs( if let Err(e) = res { tracing::warn!("Error processing jobs: {}", format!("{e}")); tracing::warn!("{}", format!("{e:?}")); + + if e.is_disconnected() { + actix_rt::time::sleep(Duration::from_secs(3)).await; + } + continue; } @@ -323,6 +325,11 @@ async fn process_image_jobs( if let Err(e) = res { tracing::warn!("Error processing jobs: {}", format!("{e}")); tracing::warn!("{}", format!("{e:?}")); + + if e.is_disconnected() { + actix_rt::time::sleep(Duration::from_secs(3)).await; + } + continue; } diff --git a/src/queue/cleanup.rs b/src/queue/cleanup.rs index aa31814..ea3ff75 100644 --- a/src/queue/cleanup.rs +++ b/src/queue/cleanup.rs @@ -3,7 +3,8 @@ use std::sync::Arc; use crate::{ config::Configuration, error::{Error, UploadError}, - queue::{Cleanup, LocalBoxFuture}, + future::LocalBoxFuture, + queue::Cleanup, repo::{Alias, ArcRepo, DeleteToken, Hash}, serde_str::Serde, store::Store, diff --git a/src/queue/process.rs b/src/queue/process.rs index 13e8b78..14b0fce 100644 --- a/src/queue/process.rs +++ b/src/queue/process.rs @@ -3,8 +3,9 @@ use crate::{ config::Configuration, error::{Error, UploadError}, formats::InputProcessableFormat, + future::LocalBoxFuture, ingest::Session, - queue::{LocalBoxFuture, Process}, + queue::Process, repo::{Alias, ArcRepo, UploadId, UploadResult}, serde_str::Serde, store::Store, diff --git a/src/repo.rs b/src/repo.rs index dd32c19..e75fb5d 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -8,6 +8,7 @@ use crate::{ config, details::Details, error_code::{ErrorCode, OwnedErrorCode}, + future::LocalBoxFuture, stream::LocalBoxStream, }; use base64::Engine; @@ -85,6 +86,13 @@ impl RepoError { Self::Canceled => ErrorCode::PANIC, } } + + pub(crate) const fn is_disconnected(&self) -> bool { + match self { + Self::PostgresError(e) => e.is_disconnected(), + _ => false, + } + } } #[async_trait::async_trait(?Send)] @@ -564,8 +572,6 @@ impl HashPage { } } -type LocalBoxFuture<'a, T> = std::pin::Pin + 'a>>; - type PageFuture = LocalBoxFuture<'static, Result>; pub(crate) struct HashStream { diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 659eb6a..51b18a5 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -12,6 +12,7 @@ use std::{ }; use dashmap::DashMap; +use deadpool::managed::Hook; use diesel::prelude::*; use diesel_async::{ pooled_connection::{ @@ -29,6 +30,7 @@ use uuid::Uuid; use crate::{ details::Details, error_code::{ErrorCode, OwnedErrorCode}, + future::{LocalBoxFuture, WithMetrics, WithTimeout}, serde_str::Serde, stream::LocalBoxStream, }; @@ -108,6 +110,9 @@ pub(crate) enum PostgresError { #[error("Error deserializing upload result")] DeserializeUploadResult(#[source] serde_json::Error), + + #[error("Timed out waiting for postgres")] + DbTimeout, } impl PostgresError { @@ -117,11 +122,26 @@ impl PostgresError { | Self::Diesel(_) | Self::SerializeDetails(_) | Self::SerializeUploadResult(_) - | Self::Hex(_) => ErrorCode::POSTGRES_ERROR, + | Self::Hex(_) + | Self::DbTimeout => ErrorCode::POSTGRES_ERROR, Self::DeserializeDetails(_) => ErrorCode::EXTRACT_DETAILS, Self::DeserializeUploadResult(_) => ErrorCode::EXTRACT_UPLOAD_RESULT, } } + + pub(super) const fn is_disconnected(&self) -> bool { + match self { + Self::Pool( + PoolError::Closed + | PoolError::Backend(diesel_async::pooled_connection::PoolError::ConnectionError(_)), + ) + | Self::Diesel(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::ClosedConnection, + _, + )) => true, + _ => false, + } + } } impl PostgresRepo { @@ -140,6 +160,10 @@ impl PostgresRepo { handle.abort(); let _ = handle.await; + let parallelism = std::thread::available_parallelism() + .map(|u| u.into()) + .unwrap_or(1_usize); + let (tx, rx) = flume::bounded(10); let mut config = ManagerConfig::default(); @@ -149,7 +173,21 @@ impl PostgresRepo { postgres_url, config, ); + let pool = Pool::builder(mgr) + .runtime(deadpool::Runtime::Tokio1) + .wait_timeout(Some(Duration::from_secs(1))) + .create_timeout(Some(Duration::from_secs(2))) + .recycle_timeout(Some(Duration::from_secs(2))) + .post_create(Hook::sync_fn(|_, _| { + metrics::increment_counter!("pict-rs.postgres.pool.connection.create"); + Ok(()) + })) + .post_recycle(Hook::sync_fn(|_, _| { + metrics::increment_counter!("pict-rs.postgres.pool.connection.recycle"); + Ok(()) + })) + .max_size(parallelism * 8) .build() .map_err(ConnectPostgresError::BuildPool)?; @@ -160,7 +198,7 @@ impl PostgresRepo { upload_notifications: DashMap::new(), }); - let handle = crate::sync::spawn(delegate_notifications(rx, inner.clone())); + let handle = crate::sync::spawn(delegate_notifications(rx, inner.clone(), parallelism * 8)); let notifications = Arc::new(handle); @@ -195,7 +233,7 @@ impl GetConnectionMetricsGuard { impl Drop for GetConnectionMetricsGuard { fn drop(&mut self) { - metrics::increment_counter!("pict-rs.postgres.pool.get.end", "completed" => (!self.armed).to_string()); + metrics::increment_counter!("pict-rs.postgres.pool.get", "completed" => (!self.armed).to_string()); metrics::histogram!("pict-rs.postgres.pool.get.duration", self.start.elapsed().as_secs_f64(), "completed" => (!self.armed).to_string()); } } @@ -204,9 +242,12 @@ impl Inner { #[tracing::instrument(level = "TRACE", skip(self))] async fn get_connection(&self) -> Result, PostgresError> { let guard = GetConnectionMetricsGuard::guard(); - let res = self.pool.get().await.map_err(PostgresError::Pool); + + let obj = self.pool.get().await.map_err(PostgresError::Pool)?; + guard.disarm(); - res + + Ok(obj) } fn interest(self: &Arc, upload_id: UploadId) -> UploadInterest { @@ -301,14 +342,14 @@ type BoxFuture<'a, T> = std::pin::Pin + type ConfigFn = Box BoxFuture<'_, ConnectionResult> + Send + Sync + 'static>; -async fn delegate_notifications(receiver: flume::Receiver, inner: Arc) { - let parallelism = std::thread::available_parallelism() - .map(|u| u.into()) - .unwrap_or(1_usize); - +async fn delegate_notifications( + receiver: flume::Receiver, + inner: Arc, + capacity: usize, +) { let mut job_notifier_state = JobNotifierState { inner: &inner, - capacity: parallelism * 8, + capacity, jobs: BTreeSet::new(), jobs_ordered: VecDeque::new(), }; @@ -409,7 +450,10 @@ impl HashRepo for PostgresRepo { let count = hashes .count() .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.count") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(count.try_into().expect("non-negative count")) @@ -424,8 +468,11 @@ impl HashRepo for PostgresRepo { let timestamp = hashes .select(created_at) .filter(hash.eq(&input_hash)) - .first(&mut conn) + .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.bound") .await + .map_err(|_| PostgresError::DbTimeout)? .map(time::PrimitiveDateTime::assume_utc) .optional() .map_err(PostgresError::Diesel)?; @@ -452,8 +499,11 @@ impl HashRepo for PostgresRepo { .select((created_at, hash)) .filter(created_at.lt(timestamp)) .order(created_at.desc()) - .first::<(time::PrimitiveDateTime, Hash)>(&mut conn) + .get_result::<(time::PrimitiveDateTime, Hash)>(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.ordered-hash") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)? .map(|tup| OrderedHash { @@ -488,8 +538,11 @@ impl HashRepo for PostgresRepo { .order(created_at.desc()) .then_order_by(hash.desc()) .limit(limit as i64 + 1) - .load::(&mut conn) + .get_results::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.next-hashes") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; let prev = hashes @@ -500,7 +553,10 @@ impl HashRepo for PostgresRepo { .then_order_by(hash) .limit(limit as i64) .get_results::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.prev-hashes") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)? .pop(); @@ -511,8 +567,11 @@ impl HashRepo for PostgresRepo { .order(created_at.desc()) .then_order_by(hash.desc()) .limit(limit as i64 + 1) - .load::(&mut conn) + .get_results::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.first-hashes") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; (page, None) @@ -548,7 +607,10 @@ impl HashRepo for PostgresRepo { created_at.eq(×tamp), )) .execute(&mut conn) - .await; + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.create-hash") + .await + .map_err(|_| PostgresError::DbTimeout)?; match res { Ok(_) => Ok(Ok(())), @@ -574,7 +636,10 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .set(identifier.eq(input_identifier.as_ref())) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.update-identifier") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -590,7 +655,10 @@ impl HashRepo for PostgresRepo { .select(identifier) .filter(hash.eq(&input_hash)) .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.identifier") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)?; @@ -615,7 +683,10 @@ impl HashRepo for PostgresRepo { identifier.eq(input_identifier.as_ref()), )) .execute(&mut conn) - .await; + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.variants.relate-variant-identifier") + .await + .map_err(|_| PostgresError::DbTimeout)?; match res { Ok(_) => Ok(Ok(())), @@ -642,7 +713,10 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .filter(variant.eq(&input_variant)) .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.variants.identifier") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)? .map(Arc::from); @@ -660,7 +734,10 @@ impl HashRepo for PostgresRepo { .select((variant, identifier)) .filter(hash.eq(&input_hash)) .get_results::<(String, String)>(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.variants.for-hash") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)? .into_iter() .map(|(s, i)| (s, Arc::from(i))) @@ -683,7 +760,10 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .filter(variant.eq(&input_variant)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.variants.remove") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -703,7 +783,10 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .set(motion_identifier.eq(input_identifier.as_ref())) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.relate-motion-identifier") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -719,7 +802,10 @@ impl HashRepo for PostgresRepo { .select(motion_identifier) .filter(hash.eq(&input_hash)) .get_result::>(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.hashes.motion-identifier") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)? .flatten() @@ -737,11 +823,13 @@ impl HashRepo for PostgresRepo { diesel::delete(schema::variants::dsl::variants) .filter(schema::variants::dsl::hash.eq(&input_hash)) .execute(conn) + .with_metrics("pict-rs.postgres.variants.cleanup") .await?; diesel::delete(schema::hashes::dsl::hashes) .filter(schema::hashes::dsl::hash.eq(&input_hash)) .execute(conn) + .with_metrics("pict-rs.postgres.hashes.cleanup") .await }) }) @@ -772,7 +860,10 @@ impl AliasRepo for PostgresRepo { token.eq(delete_token), )) .execute(&mut conn) - .await; + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.aliases.create") + .await + .map_err(|_| PostgresError::DbTimeout)?; match res { Ok(_) => Ok(Ok(())), @@ -794,7 +885,10 @@ impl AliasRepo for PostgresRepo { .select(token) .filter(alias.eq(input_alias)) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.aliases.delete-token") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)?; @@ -811,7 +905,10 @@ impl AliasRepo for PostgresRepo { .select(hash) .filter(alias.eq(input_alias)) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.aliases.hash") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)?; @@ -828,7 +925,10 @@ impl AliasRepo for PostgresRepo { .select(alias) .filter(hash.eq(&input_hash)) .get_results(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.aliases.for-hash") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(vec) @@ -843,7 +943,10 @@ impl AliasRepo for PostgresRepo { diesel::delete(aliases) .filter(alias.eq(input_alias)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.aliases.cleanup") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -866,7 +969,10 @@ impl SettingsRepo for PostgresRepo { .do_update() .set(value.eq(&input_value)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.settings.set") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -882,7 +988,10 @@ impl SettingsRepo for PostgresRepo { .select(value) .filter(key.eq(input_key)) .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.settings.get") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)? .map(hex::decode) @@ -902,7 +1011,10 @@ impl SettingsRepo for PostgresRepo { diesel::delete(settings) .filter(key.eq(input_key)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.settings.remove") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -927,7 +1039,10 @@ impl DetailsRepo for PostgresRepo { diesel::insert_into(details) .values((identifier.eq(input_identifier.as_ref()), json.eq(&value))) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.details.relate") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -943,7 +1058,10 @@ impl DetailsRepo for PostgresRepo { .select(json) .filter(identifier.eq(input_identifier.as_ref())) .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.details.get") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)? .map(serde_json::from_value) @@ -963,7 +1081,10 @@ impl DetailsRepo for PostgresRepo { diesel::delete(details) .filter(identifier.eq(input_identifier.as_ref())) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.details.cleanup") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -988,7 +1109,10 @@ impl QueueRepo for PostgresRepo { .values((queue.eq(queue_name), job.eq(job_json))) .returning(id) .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.queue.push") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; guard.disarm(); @@ -1018,7 +1142,10 @@ impl QueueRepo for PostgresRepo { diesel::sql_query("LISTEN queue_status_channel;") .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.queue.listen") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; let timestamp = to_primitive(time::OffsetDateTime::now_utc()); @@ -1030,7 +1157,10 @@ impl QueueRepo for PostgresRepo { status.eq(JobStatus::New), )) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.queue.requeue") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; if count > 0 { @@ -1045,7 +1175,10 @@ impl QueueRepo for PostgresRepo { .order(queue_time) .limit(1) .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.queue.select") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)?; @@ -1063,7 +1196,10 @@ impl QueueRepo for PostgresRepo { )) .returning((id, job)) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.queue.claim") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)?; @@ -1110,7 +1246,10 @@ impl QueueRepo for PostgresRepo { ) .set(heartbeat.eq(timestamp)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.queue.heartbeat") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1134,7 +1273,10 @@ impl QueueRepo for PostgresRepo { .and(worker.eq(worker_id)), ) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.queue.complete") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1152,7 +1294,10 @@ impl StoreMigrationRepo for PostgresRepo { let count = store_migrations .count() .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.store-migration.count") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(count > 0) @@ -1176,7 +1321,10 @@ impl StoreMigrationRepo for PostgresRepo { .on_conflict((old_identifier, new_identifier)) .do_nothing() .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.store-migration.mark-migrated") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1192,7 +1340,10 @@ impl StoreMigrationRepo for PostgresRepo { store_migrations.filter(old_identifier.eq(input_old_identifier.as_ref())), )) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.store-migration.is-migrated") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(b) @@ -1206,7 +1357,10 @@ impl StoreMigrationRepo for PostgresRepo { diesel::delete(store_migrations) .execute(&mut conn) + .with_timeout(Duration::from_secs(20)) + .with_metrics("pict-rs.postgres.store-migration.clear") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1224,7 +1378,10 @@ impl ProxyRepo for PostgresRepo { diesel::insert_into(proxies) .values((url.eq(input_url.as_str()), alias.eq(&input_alias))) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.proxy.relate-url") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1240,7 +1397,10 @@ impl ProxyRepo for PostgresRepo { .select(alias) .filter(url.eq(input_url.as_str())) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.proxy.related") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)?; @@ -1256,7 +1416,10 @@ impl ProxyRepo for PostgresRepo { diesel::delete(proxies) .filter(alias.eq(&input_alias)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.proxy.remove-relation") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1281,7 +1444,10 @@ impl AliasAccessRepo for PostgresRepo { .filter(alias.eq(&input_alias)) .set(accessed.eq(timestamp)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.alias-access.set-accessed") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1300,7 +1466,10 @@ impl AliasAccessRepo for PostgresRepo { .select(accessed) .filter(alias.eq(&input_alias)) .get_result::(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.alias-access.accessed-at") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)? .map(time::PrimitiveDateTime::assume_utc); @@ -1330,7 +1499,10 @@ impl AliasAccessRepo for PostgresRepo { .order(accessed.desc()) .limit(100) .get_results(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.alias-access.older-aliases") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(vec) @@ -1364,7 +1536,10 @@ impl VariantAccessRepo for PostgresRepo { .filter(hash.eq(&input_hash).and(variant.eq(&input_variant))) .set(accessed.eq(timestamp)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.variant-access.set-accessed") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1384,7 +1559,10 @@ impl VariantAccessRepo for PostgresRepo { .select(accessed) .filter(hash.eq(&input_hash).and(variant.eq(&input_variant))) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.variant-access.accessed-at") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)? .map(time::PrimitiveDateTime::assume_utc); @@ -1414,7 +1592,10 @@ impl VariantAccessRepo for PostgresRepo { .order(accessed.desc()) .limit(100) .get_results(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.variant-access.older-variants") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(vec) @@ -1477,7 +1658,10 @@ impl UploadRepo for PostgresRepo { .default_values() .returning(id) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.uploads.create") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(UploadId { id: uuid }) @@ -1497,14 +1681,20 @@ impl UploadRepo for PostgresRepo { diesel::sql_query("LISTEN upload_completion_channel;") .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.uploads.listen") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; let nested_opt = uploads .select(result) .filter(id.eq(upload_id.id)) .get_result(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.uploads.wait") .await + .map_err(|_| PostgresError::DbTimeout)? .optional() .map_err(PostgresError::Diesel)?; @@ -1543,7 +1733,10 @@ impl UploadRepo for PostgresRepo { diesel::delete(uploads) .filter(id.eq(upload_id.id)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.uploads.claim") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1567,7 +1760,10 @@ impl UploadRepo for PostgresRepo { .filter(id.eq(upload_id.id)) .set(result.eq(upload_result)) .execute(&mut conn) + .with_timeout(Duration::from_secs(5)) + .with_metrics("pict-rs.postgres.uploads.complete") .await + .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; Ok(()) @@ -1587,8 +1783,6 @@ impl FullRepo for PostgresRepo { } } -type LocalBoxFuture<'a, T> = std::pin::Pin + 'a>>; - type NextFuture = LocalBoxFuture<'static, Result, RepoError>>; struct PageStream { diff --git a/src/store.rs b/src/store.rs index 3293237..b6f4d1e 100644 --- a/src/store.rs +++ b/src/store.rs @@ -39,9 +39,17 @@ impl StoreError { Self::FileNotFound(_) | Self::ObjectNotFound(_) => ErrorCode::NOT_FOUND, } } + pub(crate) const fn is_not_found(&self) -> bool { matches!(self, Self::FileNotFound(_)) || matches!(self, Self::ObjectNotFound(_)) } + + pub(crate) const fn is_disconnected(&self) -> bool { + match self { + Self::Repo(e) => e.is_disconnected(), + _ => false, + } + } } impl From for StoreError { From fa11c4853e0f21f930c5d8778aaa59444065c356 Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 21:51:51 -0500 Subject: [PATCH 22/27] Update tracing style --- Cargo.lock | 11 ----------- Cargo.toml | 3 +-- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 922a62e..39d3be6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1856,7 +1856,6 @@ dependencies = [ "tracing", "tracing-actix-web", "tracing-error", - "tracing-futures", "tracing-log", "tracing-opentelemetry", "tracing-subscriber", @@ -3050,16 +3049,6 @@ dependencies = [ "tracing-subscriber", ] -[[package]] -name = "tracing-futures" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" -dependencies = [ - "pin-project", - "tracing", -] - [[package]] name = "tracing-log" version = "0.1.3" diff --git a/Cargo.toml b/Cargo.toml index c07f167..071a322 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -69,7 +69,6 @@ tokio-util = { version = "0.7", default-features = false, features = [ toml = "0.7.0" tracing = "0.1.15" tracing-error = "0.2.0" -tracing-futures = "0.2.4" tracing-log = "0.1.2" tracing-opentelemetry = "0.21" tracing-subscriber = { version = "0.3.0", features = [ @@ -86,6 +85,6 @@ uuid = { version = "1", features = ["serde", "std", "v4", "v7"] } [dependencies.tracing-actix-web] version = "0.7.5" default-features = false -features = ["opentelemetry_0_20"] +features = ["emit_event_on_error", "opentelemetry_0_20"] git = "https://github.com/asonix/tracing-actix-web" branch = "asonix/tracing-opentelemetry-021" From 3bd0f78e75060a28a85963331d0ad90205e6708d Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 21:58:57 -0500 Subject: [PATCH 23/27] Use with_timeout for all timeouts --- src/lib.rs | 7 ++++++- src/middleware.rs | 12 +++++++++--- src/process.rs | 6 +++--- src/repo/postgres.rs | 15 +++++++++------ 4 files changed, 27 insertions(+), 13 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index ad18041..b0e5680 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -38,6 +38,7 @@ use actix_web::{ web, App, HttpRequest, HttpResponse, HttpResponseBuilder, HttpServer, }; use details::{ApiDetails, HumanDate}; +use future::WithTimeout; use futures_core::Stream; use metrics_exporter_prometheus::PrometheusBuilder; use middleware::Metrics; @@ -432,7 +433,11 @@ async fn claim_upload( ) -> Result { let upload_id = Serde::into_inner(query.into_inner().upload_id); - match actix_rt::time::timeout(Duration::from_secs(10), repo.wait(upload_id)).await { + match repo + .wait(upload_id) + .with_timeout(Duration::from_secs(10)) + .await + { Ok(wait_res) => { let upload_result = wait_res?; repo.claim(upload_id).await?; diff --git a/src/middleware.rs b/src/middleware.rs index bbc2bdf..95c8846 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -12,6 +12,8 @@ use std::{ task::{Context, Poll}, }; +use crate::future::WithTimeout; + pub(crate) use self::metrics::Metrics; pub(crate) struct Deadline; @@ -149,8 +151,12 @@ impl actix_web::error::ResponseError for DeadlineExceeded { HttpResponse::build(self.status_code()) .content_type("application/json") .body( - serde_json::to_string(&serde_json::json!({ "msg": self.to_string() })) - .unwrap_or_else(|_| r#"{"msg":"request timeout"}"#.to_string()), + serde_json::to_string( + &serde_json::json!({ "msg": self.to_string(), "code": "request-timeout" }), + ) + .unwrap_or_else(|_| { + r#"{"msg":"request timeout","code":"request-timeout"}"#.to_string() + }), ) } } @@ -163,7 +169,7 @@ where DeadlineFuture { inner: match timeout { Some(duration) => DeadlineFutureInner::Timed { - timeout: actix_rt::time::timeout(duration, future), + timeout: future.with_timeout(duration), }, None => DeadlineFutureInner::Untimed { future }, }, diff --git a/src/process.rs b/src/process.rs index 1cb58d5..61cd538 100644 --- a/src/process.rs +++ b/src/process.rs @@ -14,7 +14,7 @@ use tokio::{ }; use tracing::{Instrument, Span}; -use crate::error_code::ErrorCode; +use crate::{error_code::ErrorCode, future::WithTimeout}; struct MetricsGuard { start: Instant, @@ -159,7 +159,7 @@ impl Process { timeout, } = self; - let res = actix_rt::time::timeout(timeout, child.wait()).await; + let res = child.wait().with_timeout(timeout).await; match res { Ok(Ok(status)) if status.success() => { @@ -220,7 +220,7 @@ impl Process { child.wait().await }; - let error = match actix_rt::time::timeout(timeout, child_fut).await { + let error = match child_fut.with_timeout(timeout).await { Ok(Ok(status)) if status.success() => { guard.disarm(); return; diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 51b18a5..3117ecd 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -266,11 +266,12 @@ impl Inner { impl UploadInterest { async fn notified_timeout(&self, timeout: Duration) -> Result<(), tokio::time::error::Elapsed> { - actix_rt::time::timeout( - timeout, - self.interest.as_ref().expect("interest exists").notified(), - ) - .await + self.interest + .as_ref() + .expect("interest exists") + .notified() + .with_timeout(timeout) + .await } } @@ -1214,7 +1215,9 @@ impl QueueRepo for PostgresRepo { } drop(conn); - if actix_rt::time::timeout(Duration::from_secs(5), notifier.notified()) + if notifier + .notified() + .with_timeout(Duration::from_secs(5)) .await .is_ok() { From 62ebc1c60a722bcdf22852cd79b57b8e09c0d852 Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 22:02:59 -0500 Subject: [PATCH 24/27] Move timeout after metrics - timeouts count as incomplete --- src/repo/postgres.rs | 102 +++++++++++++++++++++---------------------- 1 file changed, 51 insertions(+), 51 deletions(-) diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 3117ecd..48ee06d 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -451,8 +451,8 @@ impl HashRepo for PostgresRepo { let count = hashes .count() .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.count") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -470,8 +470,8 @@ impl HashRepo for PostgresRepo { .select(created_at) .filter(hash.eq(&input_hash)) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.bound") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map(time::PrimitiveDateTime::assume_utc) @@ -501,8 +501,8 @@ impl HashRepo for PostgresRepo { .filter(created_at.lt(timestamp)) .order(created_at.desc()) .get_result::<(time::PrimitiveDateTime, Hash)>(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.ordered-hash") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -540,8 +540,8 @@ impl HashRepo for PostgresRepo { .then_order_by(hash.desc()) .limit(limit as i64 + 1) .get_results::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.next-hashes") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -554,8 +554,8 @@ impl HashRepo for PostgresRepo { .then_order_by(hash) .limit(limit as i64) .get_results::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.prev-hashes") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)? @@ -569,8 +569,8 @@ impl HashRepo for PostgresRepo { .then_order_by(hash.desc()) .limit(limit as i64 + 1) .get_results::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.first-hashes") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -608,8 +608,8 @@ impl HashRepo for PostgresRepo { created_at.eq(×tamp), )) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.create-hash") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)?; @@ -637,8 +637,8 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .set(identifier.eq(input_identifier.as_ref())) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.update-identifier") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -656,8 +656,8 @@ impl HashRepo for PostgresRepo { .select(identifier) .filter(hash.eq(&input_hash)) .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.identifier") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -684,8 +684,8 @@ impl HashRepo for PostgresRepo { identifier.eq(input_identifier.as_ref()), )) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.variants.relate-variant-identifier") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)?; @@ -714,8 +714,8 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .filter(variant.eq(&input_variant)) .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.variants.identifier") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -735,8 +735,8 @@ impl HashRepo for PostgresRepo { .select((variant, identifier)) .filter(hash.eq(&input_hash)) .get_results::<(String, String)>(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.variants.for-hash") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)? @@ -761,8 +761,8 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .filter(variant.eq(&input_variant)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.variants.remove") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -784,8 +784,8 @@ impl HashRepo for PostgresRepo { .filter(hash.eq(&input_hash)) .set(motion_identifier.eq(input_identifier.as_ref())) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.relate-motion-identifier") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -803,8 +803,8 @@ impl HashRepo for PostgresRepo { .select(motion_identifier) .filter(hash.eq(&input_hash)) .get_result::>(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.hashes.motion-identifier") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -861,8 +861,8 @@ impl AliasRepo for PostgresRepo { token.eq(delete_token), )) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.aliases.create") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)?; @@ -886,8 +886,8 @@ impl AliasRepo for PostgresRepo { .select(token) .filter(alias.eq(input_alias)) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.aliases.delete-token") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -906,8 +906,8 @@ impl AliasRepo for PostgresRepo { .select(hash) .filter(alias.eq(input_alias)) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.aliases.hash") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -926,8 +926,8 @@ impl AliasRepo for PostgresRepo { .select(alias) .filter(hash.eq(&input_hash)) .get_results(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.aliases.for-hash") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -944,8 +944,8 @@ impl AliasRepo for PostgresRepo { diesel::delete(aliases) .filter(alias.eq(input_alias)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.aliases.cleanup") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -970,8 +970,8 @@ impl SettingsRepo for PostgresRepo { .do_update() .set(value.eq(&input_value)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.settings.set") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -989,8 +989,8 @@ impl SettingsRepo for PostgresRepo { .select(value) .filter(key.eq(input_key)) .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.settings.get") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1012,8 +1012,8 @@ impl SettingsRepo for PostgresRepo { diesel::delete(settings) .filter(key.eq(input_key)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.settings.remove") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1040,8 +1040,8 @@ impl DetailsRepo for PostgresRepo { diesel::insert_into(details) .values((identifier.eq(input_identifier.as_ref()), json.eq(&value))) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.details.relate") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1059,8 +1059,8 @@ impl DetailsRepo for PostgresRepo { .select(json) .filter(identifier.eq(input_identifier.as_ref())) .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.details.get") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1082,8 +1082,8 @@ impl DetailsRepo for PostgresRepo { diesel::delete(details) .filter(identifier.eq(input_identifier.as_ref())) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.details.cleanup") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1110,8 +1110,8 @@ impl QueueRepo for PostgresRepo { .values((queue.eq(queue_name), job.eq(job_json))) .returning(id) .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.queue.push") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1143,8 +1143,8 @@ impl QueueRepo for PostgresRepo { diesel::sql_query("LISTEN queue_status_channel;") .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.queue.listen") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1158,8 +1158,8 @@ impl QueueRepo for PostgresRepo { status.eq(JobStatus::New), )) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.queue.requeue") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1176,8 +1176,8 @@ impl QueueRepo for PostgresRepo { .order(queue_time) .limit(1) .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.queue.select") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1197,8 +1197,8 @@ impl QueueRepo for PostgresRepo { )) .returning((id, job)) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.queue.claim") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1249,8 +1249,8 @@ impl QueueRepo for PostgresRepo { ) .set(heartbeat.eq(timestamp)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.queue.heartbeat") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1276,8 +1276,8 @@ impl QueueRepo for PostgresRepo { .and(worker.eq(worker_id)), ) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.queue.complete") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1297,8 +1297,8 @@ impl StoreMigrationRepo for PostgresRepo { let count = store_migrations .count() .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.store-migration.count") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1324,8 +1324,8 @@ impl StoreMigrationRepo for PostgresRepo { .on_conflict((old_identifier, new_identifier)) .do_nothing() .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.store-migration.mark-migrated") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1343,8 +1343,8 @@ impl StoreMigrationRepo for PostgresRepo { store_migrations.filter(old_identifier.eq(input_old_identifier.as_ref())), )) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.store-migration.is-migrated") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1360,8 +1360,8 @@ impl StoreMigrationRepo for PostgresRepo { diesel::delete(store_migrations) .execute(&mut conn) - .with_timeout(Duration::from_secs(20)) .with_metrics("pict-rs.postgres.store-migration.clear") + .with_timeout(Duration::from_secs(20)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1381,8 +1381,8 @@ impl ProxyRepo for PostgresRepo { diesel::insert_into(proxies) .values((url.eq(input_url.as_str()), alias.eq(&input_alias))) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.proxy.relate-url") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1400,8 +1400,8 @@ impl ProxyRepo for PostgresRepo { .select(alias) .filter(url.eq(input_url.as_str())) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.proxy.related") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1419,8 +1419,8 @@ impl ProxyRepo for PostgresRepo { diesel::delete(proxies) .filter(alias.eq(&input_alias)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.proxy.remove-relation") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1447,8 +1447,8 @@ impl AliasAccessRepo for PostgresRepo { .filter(alias.eq(&input_alias)) .set(accessed.eq(timestamp)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.alias-access.set-accessed") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1469,8 +1469,8 @@ impl AliasAccessRepo for PostgresRepo { .select(accessed) .filter(alias.eq(&input_alias)) .get_result::(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.alias-access.accessed-at") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1502,8 +1502,8 @@ impl AliasAccessRepo for PostgresRepo { .order(accessed.desc()) .limit(100) .get_results(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.alias-access.older-aliases") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1539,8 +1539,8 @@ impl VariantAccessRepo for PostgresRepo { .filter(hash.eq(&input_hash).and(variant.eq(&input_variant))) .set(accessed.eq(timestamp)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.variant-access.set-accessed") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1562,8 +1562,8 @@ impl VariantAccessRepo for PostgresRepo { .select(accessed) .filter(hash.eq(&input_hash).and(variant.eq(&input_variant))) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.variant-access.accessed-at") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1595,8 +1595,8 @@ impl VariantAccessRepo for PostgresRepo { .order(accessed.desc()) .limit(100) .get_results(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.variant-access.older-variants") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1661,8 +1661,8 @@ impl UploadRepo for PostgresRepo { .default_values() .returning(id) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.uploads.create") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1684,8 +1684,8 @@ impl UploadRepo for PostgresRepo { diesel::sql_query("LISTEN upload_completion_channel;") .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.uploads.listen") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1694,8 +1694,8 @@ impl UploadRepo for PostgresRepo { .select(result) .filter(id.eq(upload_id.id)) .get_result(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.uploads.wait") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .optional() @@ -1736,8 +1736,8 @@ impl UploadRepo for PostgresRepo { diesel::delete(uploads) .filter(id.eq(upload_id.id)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.uploads.claim") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; @@ -1763,8 +1763,8 @@ impl UploadRepo for PostgresRepo { .filter(id.eq(upload_id.id)) .set(result.eq(upload_result)) .execute(&mut conn) - .with_timeout(Duration::from_secs(5)) .with_metrics("pict-rs.postgres.uploads.complete") + .with_timeout(Duration::from_secs(5)) .await .map_err(|_| PostgresError::DbTimeout)? .map_err(PostgresError::Diesel)?; From 492b99922be1764a8f4a1128e8a786e403479ecc Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 22:04:15 -0500 Subject: [PATCH 25/27] 10 second worker delay after disconnect error --- src/queue.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/queue.rs b/src/queue.rs index a3cbe15..0c4447e 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -204,7 +204,7 @@ async fn process_jobs( tracing::warn!("{}", format!("{e:?}")); if e.is_disconnected() { - actix_rt::time::sleep(Duration::from_secs(3)).await; + actix_rt::time::sleep(Duration::from_secs(10)).await; } continue; From 9c50bbe23a5c6dc648addf142845b5175aa0101b Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 22:06:28 -0500 Subject: [PATCH 26/27] match -> matches! --- src/repo/postgres.rs | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/repo/postgres.rs b/src/repo/postgres.rs index 48ee06d..40bfe57 100644 --- a/src/repo/postgres.rs +++ b/src/repo/postgres.rs @@ -130,17 +130,18 @@ impl PostgresError { } pub(super) const fn is_disconnected(&self) -> bool { - match self { + matches!( + self, Self::Pool( PoolError::Closed - | PoolError::Backend(diesel_async::pooled_connection::PoolError::ConnectionError(_)), - ) - | Self::Diesel(diesel::result::Error::DatabaseError( + | PoolError::Backend( + diesel_async::pooled_connection::PoolError::ConnectionError(_) + ), + ) | Self::Diesel(diesel::result::Error::DatabaseError( diesel::result::DatabaseErrorKind::ClosedConnection, _, - )) => true, - _ => false, - } + )) + ) } } From bc49f8ca3703586c76444a568527dc8625f99855 Mon Sep 17 00:00:00 2001 From: asonix Date: Mon, 4 Sep 2023 22:25:11 -0500 Subject: [PATCH 27/27] Migrate with timestamp --- src/repo/migrate.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/repo/migrate.rs b/src/repo/migrate.rs index 60dbe68..e10177a 100644 --- a/src/repo/migrate.rs +++ b/src/repo/migrate.rs @@ -206,10 +206,14 @@ async fn do_migrate_hash(old_repo: &ArcRepo, new_repo: &ArcRepo, hash: Hash) -> return Ok(()); }; - let _ = new_repo.create_hash(hash.clone(), &identifier).await?; - if let Some(details) = old_repo.details(&identifier).await? { + let _ = new_repo + .create_hash_with_timestamp(hash.clone(), &identifier, details.created_at()) + .await?; + new_repo.relate_details(&identifier, &details).await?; + } else { + let _ = new_repo.create_hash(hash.clone(), &identifier).await?; } if let Some(identifier) = old_repo.motion_identifier(hash.clone()).await? {