diff --git a/Cargo.lock b/Cargo.lock index 5711e02..39d3be6 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]] @@ -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" @@ -494,7 +500,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -670,6 +676,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 +720,69 @@ 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", + "serde_json", + "time", + "uuid", +] + +[[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-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.31", +] + +[[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.31", +] + +[[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.31", +] + [[package]] name = "digest" version = "0.10.7" @@ -740,6 +831,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" @@ -842,7 +939,7 @@ checksum = "89ca545a94061b6365f2c7355b4b32bd20df3ff95f02da9329b34ccc3bd6ee72" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -1280,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" @@ -1329,7 +1426,7 @@ checksum = "ddece26afd34c31585c74a4db0630c376df271c285d682d1e55012197830b6df" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -1440,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", ] @@ -1670,7 +1767,7 @@ dependencies = [ "pest_meta", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -1684,6 +1781,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" @@ -1694,12 +1809,17 @@ dependencies = [ "actix-web", "anyhow", "async-trait", + "barrel", "base64 0.21.3", "clap", "color-eyre", "config", "console-subscriber", "dashmap", + "deadpool", + "diesel", + "diesel-async", + "diesel-derive-enum", "flume", "futures-core", "hex", @@ -1713,6 +1833,7 @@ dependencies = [ "opentelemetry-otlp", "pin-project-lite", "quick-xml 0.30.0", + "refinery", "reqwest", "reqwest-middleware", "reqwest-tracing", @@ -1728,13 +1849,13 @@ dependencies = [ "thiserror", "time", "tokio", + "tokio-postgres", "tokio-uring", "tokio-util", "toml 0.7.6", "tracing", "tracing-actix-web", "tracing-error", - "tracing-futures", "tracing-log", "tracing-opentelemetry", "tracing-subscriber", @@ -1759,7 +1880,7 @@ checksum = "4359fd9c9171ec6e8c62926d6faaf553a8dc3f64e1507e76da7911b4f6a04405" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -1780,6 +1901,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" @@ -1930,14 +2098,60 @@ dependencies = [ ] [[package]] -name = "regex" -version = "1.9.4" +name = "refinery" +version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12de2eff854e5fa4b1295edd650e227e9d8fb0c9e90b12e7f36d6a6811791a29" +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.31", +] + +[[package]] +name = "regex" +version = "1.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "697061221ea1b4a94a624f67d0ae2bfe4e22b8a17b6a192afb11046542cc8c47" dependencies = [ "aho-corasick", "memchr", - "regex-automata 0.3.7", + "regex-automata 0.3.8", "regex-syntax 0.7.5", ] @@ -1952,9 +2166,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", @@ -2046,6 +2260,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 +2388,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" @@ -2232,7 +2471,7 @@ checksum = "4eca7ac642d82aa35b60049a6eccb4be6be75e599bd2e9adb5f875a737654af2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -2316,6 +2555,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 +2640,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" @@ -2420,9 +2675,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", @@ -2446,22 +2701,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]] @@ -2555,7 +2810,33 @@ checksum = "630bdcf245f78637c13ec01ffae6187cca34625e8c63150d424b59e55af2675e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", +] + +[[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]] @@ -2727,8 +3008,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", @@ -2746,7 +3026,7 @@ checksum = "5f4f31f56159e98206da9efd823404b79b6ef3143b4a7ab76e67b1751b25a4ab" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", ] [[package]] @@ -2769,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" @@ -2792,12 +3062,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", @@ -2936,6 +3208,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" @@ -2972,7 +3254,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", "wasm-bindgen-shared", ] @@ -3006,7 +3288,7 @@ checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.29", + "syn 2.0.31", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -3046,6 +3328,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 +3354,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..071a322 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,12 +20,17 @@ 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" config = "0.13.0" console-subscriber = "0.1" 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" @@ -39,6 +44,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 +60,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", @@ -62,9 +69,8 @@ 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.20" +tracing-opentelemetry = "0.21" tracing-subscriber = { version = "0.3.0", features = [ "ansi", "env-filter", @@ -79,4 +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" diff --git a/dev.toml b/dev.toml index 209af82..691b1fb 100644 --- a/dev.toml +++ b/dev.toml @@ -11,8 +11,12 @@ 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' +targets = 'info,pict_rs=debug' + +[metrics] +prometheus_address = "127.0.0.1:8070" [old_repo] path = 'data/sled-repo-local' @@ -59,10 +63,12 @@ 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://pictrs:1234@localhost:5432/pictrs' + +# [repo] +# type = 'sled' +# path = 'data/sled-repo-local' [store] type = 'filesystem' 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/docs/postgres-planning.md b/docs/postgres-planning.md index 158ce69..d39d5cb 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); ``` @@ -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, @@ -165,20 +165,20 @@ 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 ```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/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/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 diff --git a/src/backgrounded.rs b/src/backgrounded.rs index 0210dac..a1aa87a 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(); @@ -87,14 +82,12 @@ where 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 { @@ -102,14 +95,12 @@ where 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/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/details.rs b/src/details.rs index edf921b..f5b9a86 100644 --- a/src/details.rs +++ b/src/details.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{ bytes_stream::BytesStream, discover::Discovery, @@ -101,9 +103,10 @@ impl Details { )) } + #[tracing::instrument(level = "DEBUG")] 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/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/error.rs b/src/error.rs index 38c481e..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, } } } @@ -166,12 +180,20 @@ 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, } } + + 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/error_code.rs b/src/error_code.rs index 67eaef2..0a940ad 100644 --- a/src/error_code.rs +++ b/src/error_code.rs @@ -56,12 +56,19 @@ 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", + }; pub(crate) const CONFLICTED_RECORD: ErrorCode = ErrorCode { code: "conflicted-record", }; 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/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/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/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..75f27b3 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(); @@ -224,14 +217,12 @@ where 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() { @@ -240,14 +231,12 @@ where 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() { @@ -255,14 +244,12 @@ where 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/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 b7faff8..b0e5680 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; @@ -26,6 +27,7 @@ mod repo_04; mod serde_str; mod store; mod stream; +mod sync; mod tmp_file; mod validate; @@ -36,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; @@ -45,14 +48,15 @@ use reqwest_middleware::{ClientBuilder, ClientWithMiddleware}; use reqwest_tracing::TracingMiddleware; use rusty_s3::UrlStyle; use std::{ + marker::PhantomData, path::Path, path::PathBuf, sync::Arc, time::{Duration, SystemTime}, }; use tokio::sync::Semaphore; +use tracing::Instrument; use tracing_actix_web::TracingLogger; -use tracing_futures::Instrument; use self::{ backgrounded::Backgrounded, @@ -69,7 +73,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}, }; @@ -83,8 +87,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( @@ -93,7 +98,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 +122,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 +177,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 +246,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 +264,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 +275,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 +317,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 +376,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 +399,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?; @@ -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?; @@ -560,10 +565,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 +766,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 +856,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 +976,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 +1039,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 +1173,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 +1238,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 +1256,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 +1291,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 +1543,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,10 +1552,11 @@ async fn identifier( Ok(HttpResponse::Ok().json(&serde_json::json!({ "msg": "ok", - "identifier": identifier.string_repr(), + "identifier": identifier.as_ref(), }))) } +#[tracing::instrument(skip(repo, store))] async fn healthz( repo: web::Data, store: web::Data, @@ -1691,44 +1680,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( @@ -1810,7 +1794,7 @@ async fn launch_object_store( - repo: Repo, + repo: ArcRepo, client: ClientWithMiddleware, from: S1, to: config::primitives::Store, @@ -1824,11 +1808,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 +1842,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 +1946,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 +2010,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 +2026,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()) @@ -2066,15 +2042,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(_) => { + launch_file_store(arc_repo, store, client, config, |_| {}).await?; + } } } config::Store::ObjectStorage(config::ObjectStorage { @@ -2089,6 +2064,8 @@ impl PictRsConfiguration { client_timeout, public_endpoint, }) => { + let arc_repo = repo.to_arc(); + let store = ObjectStore::build( endpoint, bucket_name, @@ -2104,13 +2081,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 +2103,9 @@ impl PictRsConfiguration { }) .await?; } + Repo::Postgres(_) => { + launch_object_store(arc_repo, store, client, config, |_| {}).await?; + } } } } 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/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/migrate_store.rs b/src/migrate_store.rs index 5be16b9..d2ce902 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, }; @@ -58,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(()) @@ -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, @@ -365,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; } } } @@ -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/process.rs b/src/process.rs index 9621444..61cd538 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,11 +11,10 @@ use std::{ use tokio::{ io::{AsyncRead, AsyncWriteExt, ReadBuf}, process::{Child, ChildStdin, ChildStdout, Command}, - sync::oneshot::{channel, Receiver}, }; use tracing::{Instrument, Span}; -use crate::error_code::ErrorCode; +use crate::{error_code::ErrorCode, future::WithTimeout}; struct MetricsGuard { start: Instant, @@ -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, @@ -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() => { @@ -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 child_fut.with_timeout(timeout).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.rs b/src/queue.rs index 8a3fc20..0c4447e 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -3,15 +3,13 @@ use crate::{ config::Configuration, error::{Error, UploadError}, formats::InputProcessableFormat, + future::LocalBoxFuture, repo::{Alias, DeleteToken, FullRepo, Hash, JobId, UploadId}, serde_str::Serde, - store::{Identifier, Store}, + store::Store, }; -use base64::{prelude::BASE64_STANDARD, Engine}; use std::{ - future::Future, path::PathBuf, - pin::Pin, sync::Arc, time::{Duration, Instant}, }; @@ -20,32 +18,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"; @@ -55,7 +27,7 @@ enum Cleanup { hash: Hash, }, Identifier { - identifier: Base64Bytes, + identifier: String, }, Alias { alias: Serde, @@ -74,7 +46,7 @@ enum Cleanup { #[derive(Debug, serde::Deserialize, serde::Serialize)] enum Process { Ingest { - identifier: Base64Bytes, + identifier: String, upload_id: Serde, declared_alias: Option>, }, @@ -91,30 +63,30 @@ 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_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_vec(&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(()) } -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_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,42 +96,42 @@ async fn cleanup_variants( variant: Option, ) -> Result<(), Error> { let job = - serde_json::to_vec(&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_vec(&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_vec(&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_vec(&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(()) } 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_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 +142,14 @@ 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_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(()) } @@ -206,8 +178,6 @@ pub(crate) async fn process_images( .await } -type LocalBoxFuture<'a, T> = Pin + 'a>>; - async fn process_jobs( repo: &Arc, store: &S, @@ -220,7 +190,7 @@ async fn process_jobs( &'a Arc, &'a S, &'a Configuration, - &'a [u8], + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { @@ -232,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(10)).await; + } + continue; } @@ -284,13 +259,13 @@ where &'a Arc, &'a S, &'a Configuration, - &'a [u8], + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { let fut = async { - let (job_id, bytes) = 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 +278,7 @@ where queue, worker_id, job_id, - (callback)(repo, store, config, bytes.as_ref()), + (callback)(repo, store, config, job), ) }) .instrument(span) @@ -337,7 +312,7 @@ async fn process_image_jobs( &'a S, &'a ProcessMap, &'a Configuration, - &'a [u8], + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { @@ -350,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; } @@ -373,13 +353,13 @@ where &'a S, &'a ProcessMap, &'a Configuration, - &'a [u8], + serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> + Copy, { loop { let fut = async { - let (job_id, bytes) = 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 +372,7 @@ where queue, worker_id, job_id, - (callback)(repo, store, process_map, config, bytes.as_ref()), + (callback)(repo, store, process_map, config, job), ) }) .instrument(span) diff --git a/src/queue/cleanup.rs b/src/queue/cleanup.rs index 34b2424..ea3ff75 100644 --- a/src/queue/cleanup.rs +++ b/src/queue/cleanup.rs @@ -1,10 +1,13 @@ +use std::sync::Arc; + use crate::{ config::Configuration, error::{Error, UploadError}, - queue::{Base64Bytes, Cleanup, LocalBoxFuture}, + future::LocalBoxFuture, + queue::Cleanup, repo::{Alias, ArcRepo, DeleteToken, Hash}, serde_str::Serde, - store::{Identifier, Store}, + store::Store, stream::IntoStreamer, }; @@ -12,18 +15,18 @@ pub(super) fn perform<'a, S>( repo: &'a ArcRepo, store: &'a S, configuration: &'a Configuration, - job: &'a [u8], + job: serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> where S: Store, { Box::pin(async move { - match serde_json::from_slice(job) { + match serde_json::from_value(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 +53,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 +101,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 +137,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 +194,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 +204,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..14b0fce 100644 --- a/src/queue/process.rs +++ b/src/queue/process.rs @@ -3,37 +3,38 @@ use crate::{ config::Configuration, error::{Error, UploadError}, formats::InputProcessableFormat, + future::LocalBoxFuture, ingest::Session, - queue::{Base64Bytes, LocalBoxFuture, Process}, + queue::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: serde_json::Value, ) -> LocalBoxFuture<'a, Result<(), Error>> where S: Store + 'static, { Box::pin(async move { - match serde_json::from_slice(job) { + match serde_json::from_value(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, @@ -68,11 +69,11 @@ where }) } -#[tracing::instrument(skip_all)] +#[tracing::instrument(skip(repo, store, media))] 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,14 +82,12 @@ 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(); 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? @@ -97,7 +96,7 @@ where let session = crate::ingest::ingest(&repo, &store2, stream, declared_alias, &media).await?; - Ok(session) as Result, Error> + Ok(session) as Result }) .await; @@ -128,7 +127,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/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 2518cb3..e75fb5d 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -1,8 +1,14 @@ +mod alias; +mod delete_token; +mod hash; +mod metrics; +mod migrate; + use crate::{ config, details::Details, error_code::{ErrorCode, OwnedErrorCode}, - store::{Identifier, StoreError}, + future::LocalBoxFuture, stream::LocalBoxStream, }; use base64::Engine; @@ -10,10 +16,11 @@ use std::{fmt::Debug, sync::Arc}; use url::Url; use uuid::Uuid; -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}; @@ -22,6 +29,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)] @@ -30,17 +38,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)] @@ -70,6 +67,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, @@ -81,10 +81,18 @@ 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, } } + + pub(crate) const fn is_disconnected(&self) -> bool { + match self { + Self::PostgresError(e) => e.is_disconnected(), + _ => false, + } + } } #[async_trait::async_trait(?Send)] @@ -106,7 +114,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); }; @@ -127,7 +135,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); }; @@ -367,13 +375,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: serde_json::Value) -> Result; async fn pop( &self, queue: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc<[u8]>), RepoError>; + ) -> Result<(JobId, serde_json::Value), RepoError>; async fn heartbeat( &self, @@ -395,7 +403,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: serde_json::Value) -> Result { T::push(self, queue, job).await } @@ -403,7 +411,7 @@ where &self, queue: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc<[u8]>), RepoError> { + ) -> Result<(JobId, serde_json::Value), RepoError> { T::pop(self, queue, worker_id).await } @@ -455,12 +463,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)] @@ -470,17 +478,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 } } @@ -491,11 +499,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>; } @@ -511,13 +519,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 } @@ -526,7 +534,7 @@ where } } -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] pub(crate) struct OrderedHash { timestamp: time::OffsetDateTime, hash: Hash, @@ -564,12 +572,88 @@ impl HashPage { } } +type PageFuture = LocalBoxFuture<'static, Result>; + +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); @@ -599,8 +683,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 } @@ -608,38 +692,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>; } @@ -653,10 +733,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 } @@ -680,21 +756,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 } @@ -702,8 +774,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 } @@ -711,11 +783,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 } @@ -726,12 +798,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 } @@ -791,7 +863,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 +874,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(postgres_repo) => Arc::new(postgres_repo.clone()), } } } @@ -829,106 +907,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() } @@ -961,253 +939,3 @@ 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}; - - #[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 new file mode 100644 index 0000000..a0d021c --- /dev/null +++ b/src/repo/alias.rs @@ -0,0 +1,274 @@ +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) + } + } +} + +#[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 new file mode 100644 index 0000000..ca1a50f --- /dev/null +++ b/src/repo/delete_token.rs @@ -0,0 +1,160 @@ +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) + } +} + +#[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/hash.rs b/src/repo/hash.rs index 18c6dab..23d9f8b 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(|| "Invalid base64 hash".to_string().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/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/migrate.rs b/src/repo/migrate.rs index d27dd24..e10177a 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 { @@ -204,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? { @@ -266,7 +272,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 +282,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 +330,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 +346,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 new file mode 100644 index 0000000..40bfe57 --- /dev/null +++ b/src/repo/postgres.rs @@ -0,0 +1,1854 @@ +mod embedded; +mod job_status; +mod schema; + +use std::{ + collections::{BTreeSet, VecDeque}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Weak, + }, + time::{Duration, Instant}, +}; + +use dashmap::DashMap; +use deadpool::managed::Hook; +use diesel::prelude::*; +use diesel_async::{ + pooled_connection::{ + deadpool::{BuildError, Object, Pool, PoolError}, + AsyncDieselConnectionManager, ManagerConfig, + }, + AsyncConnection, AsyncPgConnection, RunQueryDsl, +}; +use tokio::sync::Notify; +use tokio_postgres::{tls::NoTlsStream, AsyncMessage, Connection, NoTls, Notification, Socket}; +use tracing::Instrument; +use url::Url; +use uuid::Uuid; + +use crate::{ + details::Details, + error_code::{ErrorCode, OwnedErrorCode}, + future::{LocalBoxFuture, WithMetrics, WithTimeout}, + serde_str::Serde, + stream::LocalBoxStream, +}; + +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, + 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>, + upload_notifications: DashMap>, +} + +struct UploadInterest { + inner: Arc, + interest: Option>, + upload_id: UploadId, +} + +struct JobNotifierState<'a> { + inner: &'a Inner, + capacity: usize, + jobs: BTreeSet, + jobs_ordered: VecDeque, +} + +struct UploadNotifierState<'a> { + inner: &'a Inner, +} + +#[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)] +pub(crate) enum PostgresError { + #[error("Error in db pool")] + Pool(#[source] PoolError), + + #[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), + + #[error("Error serializing upload result")] + SerializeUploadResult(#[source] serde_json::Error), + + #[error("Error deserializing upload result")] + DeserializeUploadResult(#[source] serde_json::Error), + + #[error("Timed out waiting for postgres")] + DbTimeout, +} + +impl PostgresError { + pub(super) const fn error_code(&self) -> ErrorCode { + match self { + Self::Pool(_) + | Self::Diesel(_) + | Self::SerializeDetails(_) + | Self::SerializeUploadResult(_) + | 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 { + matches!( + self, + Self::Pool( + PoolError::Closed + | PoolError::Backend( + diesel_async::pooled_connection::PoolError::ConnectionError(_) + ), + ) | Self::Diesel(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::ClosedConnection, + _, + )) + ) + } +} + +impl PostgresRepo { + pub(crate) async fn connect(postgres_url: Url) -> Result { + let (mut client, conn) = tokio_postgres::connect(postgres_url.as_str(), NoTls) + .await + .map_err(ConnectPostgresError::ConnectForMigration)?; + + let handle = crate::sync::spawn(conn); + + embedded::migrations::runner() + .run_async(&mut client) + .await + .map_err(ConnectPostgresError::Migration)?; + + 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(); + config.custom_setup = build_handler(tx); + + let mgr = AsyncDieselConnectionManager::::new_with_config( + 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)?; + + let inner = Arc::new(Inner { + health_count: AtomicU64::new(0), + pool, + queue_notifications: DashMap::new(), + upload_notifications: DashMap::new(), + }); + + let handle = crate::sync::spawn(delegate_notifications(rx, inner.clone(), parallelism * 8)); + + let notifications = Arc::new(handle); + + Ok(PostgresRepo { + inner, + notifications, + }) + } + + async fn get_connection(&self) -> Result, PostgresError> { + self.inner.get_connection().await + } +} + +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", "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 = "TRACE", skip(self))] + async fn get_connection(&self) -> Result, PostgresError> { + let guard = GetConnectionMetricsGuard::guard(); + + let obj = self.pool.get().await.map_err(PostgresError::Pool)?; + + guard.disarm(); + + Ok(obj) + } + + 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> { + self.interest + .as_ref() + .expect("interest exists") + .notified() + .with_timeout(timeout) + .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"); + } + } +} + +type BoxFuture<'a, T> = std::pin::Pin + Send + 'a>>; +type ConfigFn = + Box BoxFuture<'_, ConnectionResult> + Send + Sync + 'static>; + +async fn delegate_notifications( + receiver: flume::Receiver, + inner: Arc, + capacity: usize, +) { + let mut job_notifier_state = JobNotifierState { + inner: &inner, + capacity, + 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> { + let sender = sender.clone(); + + let connect_span = tracing::trace_span!(parent: None, "connect future"); + + Box::pin( + async move { + let (client, 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 + 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()) +} + +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::*; + + let mut conn = self.get_connection().await?; + + let count = hashes + .count() + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.hashes.count") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + 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::*; + + let mut conn = self.get_connection().await?; + + let timestamp = hashes + .select(created_at) + .filter(hash.eq(&input_hash)) + .get_result(&mut conn) + .with_metrics("pict-rs.postgres.hashes.bound") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map(time::PrimitiveDateTime::assume_utc) + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(timestamp.map(|timestamp| OrderedHash { + timestamp, + hash: input_hash, + })) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn hash_page_by_date( + &self, + date: time::OffsetDateTime, + limit: usize, + ) -> Result { + use schema::hashes::dsl::*; + + let mut conn = self.get_connection().await?; + + let timestamp = to_primitive(date); + + let ordered_hash = hashes + .select((created_at, hash)) + .filter(created_at.lt(timestamp)) + .order(created_at.desc()) + .get_result::<(time::PrimitiveDateTime, Hash)>(&mut conn) + .with_metrics("pict-rs.postgres.hashes.ordered-hash") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)? + .map(|tup| OrderedHash { + timestamp: tup.0.assume_utc(), + hash: tup.1, + }); + + self.hashes_ordered(ordered_hash, limit).await + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn hashes_ordered( + &self, + bound: Option, + limit: usize, + ) -> Result { + use schema::hashes::dsl::*; + + let mut conn = self.get_connection().await?; + + 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) + .get_results::(&mut conn) + .with_metrics("pict-rs.postgres.hashes.next-hashes") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .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) + .limit(limit as i64) + .get_results::(&mut conn) + .with_metrics("pict-rs.postgres.hashes.prev-hashes") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)? + .pop(); + + (page, prev) + } else { + let page = hashes + .select(hash) + .order(created_at.desc()) + .then_order_by(hash.desc()) + .limit(limit as i64 + 1) + .get_results::(&mut conn) + .with_metrics("pict-rs.postgres.hashes.first-hashes") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + (page, None) + }; + + let next = if page.len() > limit { page.pop() } else { None }; + + Ok(HashPage { + limit, + prev, + next, + hashes: page, + }) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + 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.get_connection().await?; + + let timestamp = to_primitive(timestamp); + + let res = diesel::insert_into(hashes) + .values(( + hash.eq(&input_hash), + identifier.eq(input_identifier.as_ref()), + created_at.eq(×tamp), + )) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.hashes.create-hash") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)?; + + match res { + Ok(_) => Ok(Ok(())), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => Ok(Err(HashAlreadyExists)), + Err(e) => Err(PostgresError::Diesel(e).into()), + } + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn update_identifier( + &self, + input_hash: Hash, + input_identifier: &Arc, + ) -> Result<(), RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::update(hashes) + .filter(hash.eq(&input_hash)) + .set(identifier.eq(input_identifier.as_ref())) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.hashes.update-identifier") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn identifier(&self, input_hash: Hash) -> Result>, RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = hashes + .select(identifier) + .filter(hash.eq(&input_hash)) + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.hashes.identifier") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt.map(Arc::from)) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn relate_variant_identifier( + &self, + input_hash: Hash, + input_variant: String, + input_identifier: &Arc, + ) -> Result, RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.get_connection().await?; + + let res = diesel::insert_into(variants) + .values(( + hash.eq(&input_hash), + variant.eq(&input_variant), + identifier.eq(input_identifier.as_ref()), + )) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.variants.relate-variant-identifier") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)?; + + match res { + Ok(_) => Ok(Ok(())), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => Ok(Err(VariantAlreadyExists)), + Err(e) => Err(PostgresError::Diesel(e).into()), + } + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn variant_identifier( + &self, + input_hash: Hash, + input_variant: String, + ) -> Result>, RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = variants + .select(identifier) + .filter(hash.eq(&input_hash)) + .filter(variant.eq(&input_variant)) + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.variants.identifier") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)? + .map(Arc::from); + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn variants(&self, input_hash: Hash) -> Result)>, RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.get_connection().await?; + + let vec = variants + .select((variant, identifier)) + .filter(hash.eq(&input_hash)) + .get_results::<(String, String)>(&mut conn) + .with_metrics("pict-rs.postgres.variants.for-hash") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)? + .into_iter() + .map(|(s, i)| (s, Arc::from(i))) + .collect(); + + Ok(vec) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn remove_variant( + &self, + input_hash: Hash, + input_variant: String, + ) -> Result<(), RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::delete(variants) + .filter(hash.eq(&input_hash)) + .filter(variant.eq(&input_variant)) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.variants.remove") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn relate_motion_identifier( + &self, + input_hash: Hash, + input_identifier: &Arc, + ) -> Result<(), RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::update(hashes) + .filter(hash.eq(&input_hash)) + .set(motion_identifier.eq(input_identifier.as_ref())) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.hashes.relate-motion-identifier") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn motion_identifier(&self, input_hash: Hash) -> Result>, RepoError> { + use schema::hashes::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = hashes + .select(motion_identifier) + .filter(hash.eq(&input_hash)) + .get_result::>(&mut conn) + .with_metrics("pict-rs.postgres.hashes.motion-identifier") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)? + .flatten() + .map(Arc::from); + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn cleanup_hash(&self, input_hash: Hash) -> Result<(), RepoError> { + let mut conn = self.get_connection().await?; + + conn.transaction(|conn| { + Box::pin(async move { + 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 + }) + }) + .await + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl AliasRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn create_alias( + &self, + input_alias: &Alias, + delete_token: &DeleteToken, + input_hash: Hash, + ) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.get_connection().await?; + + let res = diesel::insert_into(aliases) + .values(( + alias.eq(input_alias), + hash.eq(&input_hash), + token.eq(delete_token), + )) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.aliases.create") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)?; + + match res { + Ok(_) => Ok(Ok(())), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => Ok(Err(AliasAlreadyExists)), + Err(e) => Err(PostgresError::Diesel(e).into()), + } + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn delete_token(&self, input_alias: &Alias) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = aliases + .select(token) + .filter(alias.eq(input_alias)) + .get_result(&mut conn) + .with_metrics("pict-rs.postgres.aliases.delete-token") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn hash(&self, input_alias: &Alias) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = aliases + .select(hash) + .filter(alias.eq(input_alias)) + .get_result(&mut conn) + .with_metrics("pict-rs.postgres.aliases.hash") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn aliases_for_hash(&self, input_hash: Hash) -> Result, RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.get_connection().await?; + + let vec = aliases + .select(alias) + .filter(hash.eq(&input_hash)) + .get_results(&mut conn) + .with_metrics("pict-rs.postgres.aliases.for-hash") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(vec) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn cleanup_alias(&self, input_alias: &Alias) -> Result<(), RepoError> { + use schema::aliases::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::delete(aliases) + .filter(alias.eq(input_alias)) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.aliases.cleanup") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[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::*; + + let input_value = hex::encode(input_value); + + let mut conn = self.get_connection().await?; + + diesel::insert_into(settings) + .values((key.eq(input_key), value.eq(&input_value))) + .on_conflict(key) + .do_update() + .set(value.eq(&input_value)) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.settings.set") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn get(&self, input_key: &'static str) -> Result>, RepoError> { + use schema::settings::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = settings + .select(value) + .filter(key.eq(input_key)) + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.settings.get") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)? + .map(hex::decode) + .transpose() + .map_err(PostgresError::Hex)? + .map(Arc::from); + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn remove(&self, input_key: &'static str) -> Result<(), RepoError> { + use schema::settings::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::delete(settings) + .filter(key.eq(input_key)) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.settings.remove") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[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, + input_details: &Details, + ) -> Result<(), RepoError> { + use schema::details::dsl::*; + + let mut conn = self.get_connection().await?; + + 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) + .with_metrics("pict-rs.postgres.details.relate") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn details(&self, input_identifier: &Arc) -> Result, RepoError> { + use schema::details::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = details + .select(json) + .filter(identifier.eq(input_identifier.as_ref())) + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.details.get") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)? + .map(serde_json::from_value) + .transpose() + .map_err(PostgresError::DeserializeDetails)? + .map(|inner| Details { inner }); + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn cleanup_details(&self, input_identifier: &Arc) -> Result<(), RepoError> { + use schema::details::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::delete(details) + .filter(identifier.eq(input_identifier.as_ref())) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.details.cleanup") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[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, + job_json: serde_json::Value, + ) -> Result { + let guard = PushMetricsGuard::guard(queue_name); + + use schema::job_queue::dsl::*; + + let mut conn = self.get_connection().await?; + + let job_id = diesel::insert_into(job_queue) + .values((queue.eq(queue_name), job.eq(job_json))) + .returning(id) + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.queue.push") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + guard.disarm(); + + Ok(JobId(job_id)) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn pop( + &self, + 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 { + let mut conn = self.get_connection().await?; + + let notifier: Arc = self + .inner + .queue_notifications + .entry(String::from(queue_name)) + .or_insert_with(crate::sync::notify) + .clone(); + + diesel::sql_query("LISTEN queue_status_channel;") + .execute(&mut conn) + .with_metrics("pict-rs.postgres.queue.listen") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + let timestamp = to_primitive(time::OffsetDateTime::now_utc()); + + let count = 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) + .with_metrics("pict-rs.postgres.queue.requeue") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + if count > 0 { + tracing::info!("Reset {count} jobs"); + } + + // 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) + .with_metrics("pict-rs.postgres.queue.select") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .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) + .with_metrics("pict-rs.postgres.queue.claim") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)?; + + if let Some(tup) = opt { + break Some(tup); + } + }; + + if let Some((job_id, job_json)) = opt { + guard.disarm(); + return Ok((JobId(job_id), job_json)); + } + + drop(conn); + if notifier + .notified() + .with_timeout(Duration::from_secs(5)) + .await + .is_ok() + { + tracing::debug!("Notified"); + } else { + tracing::debug!("Timed out"); + } + } + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + 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.get_connection().await?; + + 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) + .with_metrics("pict-rs.postgres.queue.heartbeat") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + 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.get_connection().await?; + + diesel::delete(job_queue) + .filter( + id.eq(job_id.0) + .and(queue.eq(queue_name)) + .and(worker.eq(worker_id)), + ) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.queue.complete") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[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::*; + + let mut conn = self.get_connection().await?; + + let count = store_migrations + .count() + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.store-migration.count") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(count > 0) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn mark_migrated( + &self, + input_old_identifier: &Arc, + input_new_identifier: &Arc, + ) -> Result<(), RepoError> { + use schema::store_migrations::dsl::*; + + let mut conn = self.get_connection().await?; + + 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) + .with_metrics("pict-rs.postgres.store-migration.mark-migrated") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn is_migrated(&self, input_old_identifier: &Arc) -> Result { + use schema::store_migrations::dsl::*; + + 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())), + )) + .get_result(&mut conn) + .with_metrics("pict-rs.postgres.store-migration.is-migrated") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(b) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn clear(&self) -> Result<(), RepoError> { + use schema::store_migrations::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::delete(store_migrations) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.store-migration.clear") + .with_timeout(Duration::from_secs(20)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[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::*; + + let mut conn = self.get_connection().await?; + + diesel::insert_into(proxies) + .values((url.eq(input_url.as_str()), alias.eq(&input_alias))) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.proxy.relate-url") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn related(&self, input_url: Url) -> Result, RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = proxies + .select(alias) + .filter(url.eq(input_url.as_str())) + .get_result(&mut conn) + .with_metrics("pict-rs.postgres.proxy.related") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)?; + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn remove_relation(&self, input_alias: Alias) -> Result<(), RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::delete(proxies) + .filter(alias.eq(&input_alias)) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.proxy.remove-relation") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[async_trait::async_trait(?Send)] +impl AliasAccessRepo for PostgresRepo { + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn set_accessed_alias( + &self, + input_alias: Alias, + timestamp: time::OffsetDateTime, + ) -> Result<(), RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.get_connection().await?; + + let timestamp = to_primitive(timestamp); + + diesel::update(proxies) + .filter(alias.eq(&input_alias)) + .set(accessed.eq(timestamp)) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.alias-access.set-accessed") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn alias_accessed_at( + &self, + input_alias: Alias, + ) -> Result, RepoError> { + use schema::proxies::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = proxies + .select(accessed) + .filter(alias.eq(&input_alias)) + .get_result::(&mut conn) + .with_metrics("pict-rs.postgres.alias-access.accessed-at") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)? + .map(time::PrimitiveDateTime::assume_utc); + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + 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.get_connection().await?; + + let vec = proxies + .select((accessed, alias)) + .filter(accessed.lt(older_than)) + .order(accessed.desc()) + .limit(100) + .get_results(&mut conn) + .with_metrics("pict-rs.postgres.alias-access.older-aliases") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .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 { + #[tracing::instrument(level = "DEBUG", skip(self))] + 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.get_connection().await?; + + 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) + .with_metrics("pict-rs.postgres.variant-access.set-accessed") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn variant_accessed_at( + &self, + input_hash: Hash, + input_variant: String, + ) -> Result, RepoError> { + use schema::variants::dsl::*; + + let mut conn = self.get_connection().await?; + + let opt = variants + .select(accessed) + .filter(hash.eq(&input_hash).and(variant.eq(&input_variant))) + .get_result(&mut conn) + .with_metrics("pict-rs.postgres.variant-access.accessed-at") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)? + .map(time::PrimitiveDateTime::assume_utc); + + Ok(opt) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + 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.get_connection().await?; + + let vec = variants + .select((accessed, (hash, variant))) + .filter(accessed.lt(older_than)) + .order(accessed.desc()) + .limit(100) + .get_results(&mut conn) + .with_metrics("pict-rs.postgres.variant-access.older-variants") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .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 { + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn create_upload(&self) -> Result { + use schema::uploads::dsl::*; + + let mut conn = self.get_connection().await?; + + let uuid = diesel::insert_into(uploads) + .default_values() + .returning(id) + .get_result(&mut conn) + .with_metrics("pict-rs.postgres.uploads.create") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(UploadId { id: uuid }) + } + + #[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;") + .execute(&mut conn) + .with_metrics("pict-rs.postgres.uploads.listen") + .with_timeout(Duration::from_secs(5)) + .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_metrics("pict-rs.postgres.uploads.wait") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .optional() + .map_err(PostgresError::Diesel)?; + + 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)?; + + guard.disarm(); + return Ok(upload_result.into()); + } + } + None => { + return Err(RepoError::AlreadyClaimed); + } + } + + drop(conn); + + if interest_future.await.is_ok() { + tracing::debug!("Notified"); + } else { + tracing::debug!("Timed out"); + } + } + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn claim(&self, upload_id: UploadId) -> Result<(), RepoError> { + use schema::uploads::dsl::*; + + let mut conn = self.get_connection().await?; + + diesel::delete(uploads) + .filter(id.eq(upload_id.id)) + .execute(&mut conn) + .with_metrics("pict-rs.postgres.uploads.claim") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } + + #[tracing::instrument(level = "DEBUG", skip(self))] + async fn complete_upload( + &self, + upload_id: UploadId, + upload_result: UploadResult, + ) -> Result<(), RepoError> { + use schema::uploads::dsl::*; + + let mut conn = self.get_connection().await?; + + 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) + .with_metrics("pict-rs.postgres.uploads.complete") + .with_timeout(Duration::from_secs(5)) + .await + .map_err(|_| PostgresError::DbTimeout)? + .map_err(PostgresError::Diesel)?; + + Ok(()) + } +} + +#[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); + + self.set("health-value", Arc::from(next.to_be_bytes())) + .await?; + + Ok(()) + } +} + +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") + .field("pool", &"pool") + .finish() + } +} 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/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/V0000__enable_pgcrypto.rs b/src/repo/postgres/migrations/V0000__enable_pgcrypto.rs new file mode 100644 index 0000000..4c8bd57 --- /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 IF NOT EXISTS 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 new file mode 100644 index 0000000..1c97f08 --- /dev/null +++ b/src/repo/postgres/migrations/V0001__create_hashes.rs @@ -0,0 +1,33 @@ +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::text() + .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), + ); + + 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..4e62f1b --- /dev/null +++ b/src/repo/postgres/migrations/V0002__create_variants.rs @@ -0,0 +1,28 @@ +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::text().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), + ); + }); + + m.make::().to_string() +} 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..007f123 --- /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::text().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..440aa02 --- /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("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 new file mode 100644 index 0000000..cfabb83 --- /dev/null +++ b/src/repo/postgres/migrations/V0006__create_queue.rs @@ -0,0 +1,57 @@ +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("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)); + t.add_column("worker", types::uuid().nullable(true)); + t.add_column( + "status", + types::custom("job_status").nullable(false).default("new"), + ); + t.add_column( + "queue_time", + types::datetime() + .nullable(false) + .default(AutogenFunction::CurrentTimestamp), + ); + 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"])); + }); + + m.inject_custom( + r#" +CREATE OR REPLACE FUNCTION queue_status_notify() + RETURNS trigger AS +$$ +BEGIN + PERFORM pg_notify('queue_status_channel', NEW.id::text || ' ' || NEW.queue::text); + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + "# + .trim(), + ); + + m.inject_custom( + r#" +CREATE TRIGGER queue_status + AFTER INSERT OR UPDATE OF status + ON job_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..db106fc --- /dev/null +++ b/src/repo/postgres/migrations/V0007__create_store_migrations.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("store_migrations", |t| { + t.add_column( + "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/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..4cfbafa --- /dev/null +++ b/src/repo/postgres/migrations/V0009__create_uploads.rs @@ -0,0 +1,44 @@ +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").nullable(true)); + t.add_column( + "created_at", + types::datetime() + .nullable(false) + .default(AutogenFunction::CurrentTimestamp), + ); + }); + + 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..72ba9aa --- /dev/null +++ b/src/repo/postgres/schema.rs @@ -0,0 +1,115 @@ +// @generated automatically by Diesel CLI. + +pub mod sql_types { + #[derive(diesel::query_builder::QueryId, diesel::sql_types::SqlType)] + #[diesel(postgres_type(name = "job_status"))] + pub struct JobStatus; +} + +diesel::table! { + aliases (alias) { + alias -> Text, + hash -> Text, + token -> Text, + } +} + +diesel::table! { + details (identifier) { + identifier -> Text, + json -> Jsonb, + } +} + +diesel::table! { + hashes (hash) { + 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, + worker -> Nullable, + status -> JobStatus, + queue_time -> Timestamp, + heartbeat -> Nullable, + } +} + +diesel::table! { + proxies (url) { + url -> Text, + alias -> Text, + accessed -> 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 (old_identifier) { + old_identifier -> Text, + new_identifier -> Text, + } +} + +diesel::table! { + uploads (id) { + id -> Uuid, + result -> Nullable, + created_at -> Timestamp, + } +} + +diesel::table! { + variants (id) { + id -> Uuid, + hash -> Text, + 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, + job_queue, + proxies, + refinery_schema_history, + settings, + store_migrations, + uploads, + variants, +); diff --git a/src/repo/sled.rs b/src/repo/sled.rs index 57527c2..77eb054 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}; @@ -13,26 +12,25 @@ 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, Identifier, 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 { ($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)? @@ -47,14 +45,20 @@ 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), + #[error("Invalid string data in db")] + Utf8(#[source] std::str::Utf8Error), + + #[error("Invalid job json")] + Job(#[source] serde_json::Error), + #[error("Operation panicked")] Panic, @@ -65,9 +69,10 @@ 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::Job(_) => ErrorCode::EXTRACT_JOB, Self::Panic => ErrorCode::PANIC, Self::Conflict => ErrorCode::CONFLICTED_RECORD, } @@ -169,7 +174,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); }) @@ -252,7 +257,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)? { @@ -318,7 +323,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())? { @@ -358,7 +363,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)? { @@ -428,7 +433,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)? { @@ -486,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))] @@ -547,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(); @@ -556,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 { @@ -571,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()); } } @@ -648,19 +609,31 @@ 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: 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(); - 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(); @@ -687,7 +660,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(); @@ -700,7 +673,7 @@ impl QueueRepo for SledRepo { &self, queue_name: &'static str, worker_id: Uuid, - ) -> Result<(JobId, Arc<[u8]>), RepoError> { + ) -> Result<(JobId, serde_json::Value), RepoError> { let metrics_guard = PopMetricsGuard::guard(queue_name); let now = time::OffsetDateTime::now_utc(); @@ -710,7 +683,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. @@ -755,9 +728,12 @@ impl QueueRepo for SledRepo { let opt = queue .get(&key)? - .map(|job_bytes| (job_id, Arc::from(job_bytes.to_vec()))); + .map(|ivec| serde_json::from_slice(&ivec[..])) + .transpose() + .map_err(SledError::Job)?; - return Ok(opt) as Result)>, SledError>; + return Ok(opt.map(|job| (job_id, job))) + as Result, SledError>; } Ok(None) @@ -781,9 +757,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) }; @@ -802,7 +776,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); @@ -832,7 +806,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[..])?; @@ -949,43 +923,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 +976,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 +1043,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())); @@ -1095,7 +1065,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() @@ -1147,7 +1117,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() @@ -1197,10 +1167,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 +1204,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 || { + crate::sync::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 +1261,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 +1271,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 +1284,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 +1313,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 +1339,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..07ccd9e 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}; @@ -35,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)? @@ -56,6 +53,9 @@ pub(crate) enum SledError { #[error("Operation panicked")] Panic, + + #[error("Error reading string")] + Utf8(#[from] std::str::Utf8Error), } #[derive(Clone)] @@ -179,17 +179,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 +219,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 +252,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 +261,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..b6f4d1e 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; @@ -40,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 { @@ -70,32 +77,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 +93,7 @@ pub(crate) trait Store: Clone + Debug { &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static; @@ -111,28 +101,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 +130,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 +138,7 @@ where &self, reader: Reader, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static, { @@ -162,7 +149,7 @@ where &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static, { @@ -173,26 +160,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 +188,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 +274,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 +282,7 @@ where &self, reader: Reader, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static, { @@ -237,7 +293,7 @@ where &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static, { @@ -248,26 +304,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 +332,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 b2a1cb7..e208261 100644 --- a/src/store/file_store.rs +++ b/src/store/file_store.rs @@ -1,23 +1,17 @@ use crate::{ - error_code::ErrorCode, - file::File, - repo::{Repo, SettingsRepo}, - store::Store, + 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 @@ -33,12 +27,12 @@ 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")] + #[error("Couldn't strip root dir")] PrefixError, + #[error("Couldn't convert Path to String")] + StringError, + #[error("Tried to save over existing file")] FileExists, } @@ -49,7 +43,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 | Self::PrefixError => ErrorCode::FORMAT_FILE_ID_ERROR, } } } @@ -58,14 +52,12 @@ impl FileError { pub(crate) struct FileStore { path_gen: Generator, root_dir: PathBuf, - repo: Repo, + repo: ArcRepo, } #[async_trait::async_trait(?Send)] impl Store for FileStore { - type Identifier = FileId; - type Stream = Pin>>>; - + #[tracing::instrument(level = "DEBUG", skip(self))] async fn health_check(&self) -> Result<(), StoreError> { tokio::fs::metadata(&self.root_dir) .await @@ -74,12 +66,12 @@ impl Store for FileStore { Ok(()) } - #[tracing::instrument(skip(reader))] + #[tracing::instrument(skip(self, reader))] async fn save_async_read( &self, mut reader: Reader, _content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where Reader: AsyncRead + Unpin + 'static, { @@ -97,7 +89,7 @@ impl Store for FileStore { &self, stream: S, content_type: mime::Mime, - ) -> Result + ) -> Result, StoreError> where S: Stream> + Unpin + 'static, { @@ -105,12 +97,12 @@ impl Store for FileStore { .await } - #[tracing::instrument(skip(bytes))] + #[tracing::instrument(skip(self, bytes))] async fn save_bytes( &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 { @@ -121,17 +113,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] + #[tracing::instrument(skip(self))] 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"); @@ -149,10 +141,10 @@ impl Store for FileStore { Ok(Box::pin(stream)) } - #[tracing::instrument(skip(writer))] + #[tracing::instrument(skip(self, writer))] async fn read_into( &self, - identifier: &Self::Identifier, + identifier: &Arc, writer: &mut Writer, ) -> Result<(), std::io::Error> where @@ -165,8 +157,8 @@ impl Store for FileStore { Ok(()) } - #[tracing::instrument] - async fn len(&self, identifier: &Self::Identifier) -> Result { + #[tracing::instrument(skip(self))] + async fn len(&self, identifier: &Arc) -> Result { let path = self.path_from_file_id(identifier); let len = tokio::fs::metadata(path) @@ -177,8 +169,8 @@ impl Store for FileStore { Ok(len) } - #[tracing::instrument] - async fn remove(&self, identifier: &Self::Identifier) -> Result<(), StoreError> { + #[tracing::instrument(skip(self))] + async fn remove(&self, identifier: &Arc) -> Result<(), StoreError> { let path = self.path_from_file_id(identifier); self.safe_remove_file(path).await?; @@ -189,7 +181,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?; @@ -201,16 +193,24 @@ impl FileStore { }) } + fn file_id_from_path(&self, path: PathBuf) -> Result, FileError> { + 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 { + self.root_dir.join(file_id.as_ref()) + } + 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() { @@ -227,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; @@ -308,18 +309,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/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 a37f465..324a0eb 100644 --- a/src/store/object_store.rs +++ b/src/store/object_store.rs @@ -1,9 +1,9 @@ use crate::{ bytes_stream::BytesStream, error_code::ErrorCode, - repo::{Repo, SettingsRepo}, + 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); @@ -107,7 +104,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 +116,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, @@ -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, { @@ -283,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( @@ -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() @@ -493,7 +487,7 @@ impl ObjectStore { signature_expiration: u64, client_timeout: u64, public_endpoint: Option, - repo: Repo, + repo: ArcRepo, ) -> Result { let path_gen = init_generator(&repo).await?; @@ -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) } @@ -714,13 +708,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 +723,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()) } } 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 - } -} 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())); } }