diff --git a/Cargo.lock b/Cargo.lock index f50ebda890..cc002f079e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2890,14 +2890,18 @@ dependencies = [ "chirp-client", "chirp-workflow-macros", "cjson", + "fdb-util", "formatted-error", + "foundationdb", "futures-util", "global-error", + "include_dir", "indoc 2.0.5", "lazy_static", "md5", "prost 0.12.6", "prost-types 0.12.6", + "rand", "rivet-cache", "rivet-config", "rivet-connection", @@ -2915,6 +2919,7 @@ dependencies = [ "tokio", "tokio-util 0.7.12", "tracing", + "tracing-logfmt", "tracing-subscriber", "uuid", ] @@ -5971,6 +5976,17 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "fdb-util" +version = "24.6.2-rc.1" +dependencies = [ + "anyhow", + "foundationdb", + "lazy_static", + "tokio", + "tracing", +] + [[package]] name = "fdeflate" version = "0.3.7" @@ -12032,9 +12048,13 @@ dependencies = [ name = "rivet-pools" version = "24.6.2-rc.1" dependencies = [ + "anyhow", "async-nats", "clickhouse", + "fdb-util", + "foundationdb", "funty 1.1.0", + "futures-util", "global-error", "governor", "hyper 0.14.31", @@ -12045,11 +12065,13 @@ dependencies = [ "rivet-config", "rivet-metrics", "sqlx", + "tempfile", "thiserror 1.0.69", "tokio", "tokio-util 0.7.12", "tracing", "url", + "uuid", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 7775e66a36..26a1a11a01 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [workspace] resolver = "2" -members = ["packages/api/actor","packages/api/auth","packages/api/cf-verification","packages/api/cloud","packages/api/games","packages/api/group","packages/api/identity","packages/api/job","packages/api/matchmaker","packages/api/monolith-edge","packages/api/monolith-public","packages/api/portal","packages/api/provision","packages/api/status","packages/api/traefik-provider","packages/api/ui","packages/common/api-helper/build","packages/common/api-helper/macros","packages/common/cache/build","packages/common/cache/result","packages/common/chirp-workflow/core","packages/common/chirp-workflow/macros","packages/common/chirp/client","packages/common/chirp/metrics","packages/common/chirp/perf","packages/common/chirp/types","packages/common/chirp/worker","packages/common/chirp/worker-attributes","packages/common/claims","packages/common/config","packages/common/connection","packages/common/convert","packages/common/deno-embed","packages/common/env","packages/common/formatted-error","packages/common/global-error","packages/common/health-checks","packages/common/hub-embed","packages/common/kv-str","packages/common/metrics","packages/common/migrate","packages/common/nomad-util","packages/common/operation/core","packages/common/operation/macros","packages/common/pools","packages/common/redis-util","packages/common/runtime","packages/common/s3-util","packages/common/schemac","packages/common/service-manager","packages/common/smithy-output/api-auth/rust","packages/common/smithy-output/api-auth/rust-server","packages/common/smithy-output/api-cf-verification/rust","packages/common/smithy-output/api-cf-verification/rust-server","packages/common/smithy-output/api-cloud/rust","packages/common/smithy-output/api-cloud/rust-server","packages/common/smithy-output/api-group/rust","packages/common/smithy-output/api-group/rust-server","packages/common/smithy-output/api-identity/rust","packages/common/smithy-output/api-identity/rust-server","packages/common/smithy-output/api-job/rust","packages/common/smithy-output/api-job/rust-server","packages/common/smithy-output/api-kv/rust","packages/common/smithy-output/api-kv/rust-server","packages/common/smithy-output/api-matchmaker/rust","packages/common/smithy-output/api-matchmaker/rust-server","packages/common/smithy-output/api-party/rust","packages/common/smithy-output/api-party/rust-server","packages/common/smithy-output/api-portal/rust","packages/common/smithy-output/api-portal/rust-server","packages/common/smithy-output/api-status/rust","packages/common/smithy-output/api-status/rust-server","packages/common/smithy-output/api-traefik-provider/rust","packages/common/smithy-output/api-traefik-provider/rust-server","packages/common/test","packages/common/test-images","packages/common/types-proto/build","packages/common/types-proto/core","packages/common/util/core","packages/common/util/macros","packages/common/util/search","packages/infra/client/actor-kv","packages/infra/client/config","packages/infra/client/container-runner","packages/infra/client/echo","packages/infra/client/isolate-v8-runner","packages/infra/client/logs","packages/infra/client/manager","packages/infra/legacy/job-runner","packages/infra/schema-generator","packages/infra/server","packages/services/build","packages/services/build/ops/create","packages/services/build/ops/get","packages/services/build/ops/list-for-env","packages/services/build/ops/list-for-game","packages/services/build/standalone/default-create","packages/services/build/util","packages/services/captcha/ops/hcaptcha-config-get","packages/services/captcha/ops/hcaptcha-verify","packages/services/captcha/ops/request","packages/services/captcha/ops/turnstile-config-get","packages/services/captcha/ops/turnstile-verify","packages/services/captcha/ops/verify","packages/services/captcha/util","packages/services/cdn/ops/namespace-auth-user-remove","packages/services/cdn/ops/namespace-auth-user-update","packages/services/cdn/ops/namespace-create","packages/services/cdn/ops/namespace-domain-create","packages/services/cdn/ops/namespace-domain-remove","packages/services/cdn/ops/namespace-get","packages/services/cdn/ops/namespace-resolve-domain","packages/services/cdn/ops/ns-auth-type-set","packages/services/cdn/ops/ns-enable-domain-public-auth-set","packages/services/cdn/ops/site-create","packages/services/cdn/ops/site-get","packages/services/cdn/ops/site-list-for-game","packages/services/cdn/ops/version-get","packages/services/cdn/ops/version-prepare","packages/services/cdn/ops/version-publish","packages/services/cdn/util","packages/services/cdn/worker","packages/services/cf-custom-hostname/ops/get","packages/services/cf-custom-hostname/ops/list-for-namespace-id","packages/services/cf-custom-hostname/ops/resolve-hostname","packages/services/cf-custom-hostname/worker","packages/services/cloud/ops/device-link-create","packages/services/cloud/ops/game-config-create","packages/services/cloud/ops/game-config-get","packages/services/cloud/ops/game-token-create","packages/services/cloud/ops/namespace-create","packages/services/cloud/ops/namespace-get","packages/services/cloud/ops/namespace-token-development-create","packages/services/cloud/ops/namespace-token-public-create","packages/services/cloud/ops/version-get","packages/services/cloud/ops/version-publish","packages/services/cloud/standalone/default-create","packages/services/cloud/worker","packages/services/cluster","packages/services/cluster/standalone/datacenter-tls-renew","packages/services/cluster/standalone/default-update","packages/services/cluster/standalone/gc","packages/services/cluster/standalone/metrics-publish","packages/services/custom-user-avatar/ops/list-for-game","packages/services/custom-user-avatar/ops/upload-complete","packages/services/debug/ops/email-res","packages/services/ds","packages/services/ds-log/ops/export","packages/services/ds-log/ops/read","packages/services/dynamic-config","packages/services/email-verification/ops/complete","packages/services/email-verification/ops/create","packages/services/email/ops/send","packages/services/external/ops/request-validate","packages/services/external/worker","packages/services/faker/ops/build","packages/services/faker/ops/cdn-site","packages/services/faker/ops/game","packages/services/faker/ops/game-namespace","packages/services/faker/ops/game-version","packages/services/faker/ops/job-run","packages/services/faker/ops/job-template","packages/services/faker/ops/mm-lobby","packages/services/faker/ops/mm-lobby-row","packages/services/faker/ops/mm-player","packages/services/faker/ops/region","packages/services/faker/ops/team","packages/services/faker/ops/user","packages/services/game/ops/banner-upload-complete","packages/services/game/ops/create","packages/services/game/ops/get","packages/services/game/ops/list-all","packages/services/game/ops/list-for-team","packages/services/game/ops/logo-upload-complete","packages/services/game/ops/namespace-create","packages/services/game/ops/namespace-get","packages/services/game/ops/namespace-list","packages/services/game/ops/namespace-resolve-name-id","packages/services/game/ops/namespace-resolve-url","packages/services/game/ops/namespace-validate","packages/services/game/ops/namespace-version-history-list","packages/services/game/ops/namespace-version-set","packages/services/game/ops/recommend","packages/services/game/ops/resolve-name-id","packages/services/game/ops/resolve-namespace-id","packages/services/game/ops/token-development-validate","packages/services/game/ops/validate","packages/services/game/ops/version-create","packages/services/game/ops/version-get","packages/services/game/ops/version-list","packages/services/game/ops/version-validate","packages/services/ip/ops/info","packages/services/job-log/ops/read","packages/services/job-log/worker","packages/services/job-run","packages/services/job/standalone/gc","packages/services/job/util","packages/services/linode","packages/services/linode/standalone/gc","packages/services/load-test/standalone/api-cloud","packages/services/load-test/standalone/mm","packages/services/load-test/standalone/mm-sustain","packages/services/load-test/standalone/sqlx","packages/services/load-test/standalone/watch-requests","packages/services/mm-config/ops/game-get","packages/services/mm-config/ops/game-upsert","packages/services/mm-config/ops/lobby-group-get","packages/services/mm-config/ops/lobby-group-resolve-name-id","packages/services/mm-config/ops/lobby-group-resolve-version","packages/services/mm-config/ops/namespace-config-set","packages/services/mm-config/ops/namespace-config-validate","packages/services/mm-config/ops/namespace-create","packages/services/mm-config/ops/namespace-get","packages/services/mm-config/ops/version-get","packages/services/mm-config/ops/version-prepare","packages/services/mm-config/ops/version-publish","packages/services/mm/ops/dev-player-token-create","packages/services/mm/ops/lobby-find-fail","packages/services/mm/ops/lobby-find-lobby-query-list","packages/services/mm/ops/lobby-find-try-complete","packages/services/mm/ops/lobby-for-run-id","packages/services/mm/ops/lobby-get","packages/services/mm/ops/lobby-history","packages/services/mm/ops/lobby-idle-update","packages/services/mm/ops/lobby-list-for-namespace","packages/services/mm/ops/lobby-list-for-user-id","packages/services/mm/ops/lobby-player-count","packages/services/mm/ops/lobby-runtime-aggregate","packages/services/mm/ops/lobby-state-get","packages/services/mm/ops/player-count-for-namespace","packages/services/mm/ops/player-get","packages/services/mm/standalone/gc","packages/services/mm/util","packages/services/mm/worker","packages/services/monolith/standalone/worker","packages/services/monolith/standalone/workflow-worker","packages/services/nomad/standalone/monitor","packages/services/pegboard","packages/services/pegboard/standalone/dc-init","packages/services/pegboard/standalone/gc","packages/services/pegboard/standalone/metrics-publish","packages/services/pegboard/standalone/ws","packages/services/region/ops/get","packages/services/region/ops/list","packages/services/region/ops/list-for-game","packages/services/region/ops/recommend","packages/services/region/ops/resolve","packages/services/region/ops/resolve-for-game","packages/services/server-spec","packages/services/team-invite/ops/get","packages/services/team-invite/worker","packages/services/team/ops/avatar-upload-complete","packages/services/team/ops/get","packages/services/team/ops/join-request-list","packages/services/team/ops/member-count","packages/services/team/ops/member-get","packages/services/team/ops/member-list","packages/services/team/ops/member-relationship-get","packages/services/team/ops/profile-validate","packages/services/team/ops/recommend","packages/services/team/ops/resolve-display-name","packages/services/team/ops/user-ban-get","packages/services/team/ops/user-ban-list","packages/services/team/ops/validate","packages/services/team/util","packages/services/team/worker","packages/services/telemetry/standalone/beacon","packages/services/tier","packages/services/token/ops/create","packages/services/token/ops/exchange","packages/services/token/ops/get","packages/services/token/ops/revoke","packages/services/upload/ops/complete","packages/services/upload/ops/file-list","packages/services/upload/ops/get","packages/services/upload/ops/list-for-user","packages/services/upload/ops/prepare","packages/services/upload/worker","packages/services/user","packages/services/user-identity/ops/create","packages/services/user-identity/ops/delete","packages/services/user-identity/ops/get","packages/services/user/ops/avatar-upload-complete","packages/services/user/ops/get","packages/services/user/ops/pending-delete-toggle","packages/services/user/ops/profile-validate","packages/services/user/ops/resolve-email","packages/services/user/ops/team-list","packages/services/user/ops/token-create","packages/services/user/standalone/delete-pending","packages/services/user/worker","packages/services/workflow/standalone/gc","packages/services/workflow/standalone/metrics-publish","packages/toolchain/actors-sdk-embed","packages/toolchain/cli","packages/toolchain/js-utils-embed","packages/toolchain/toolchain","sdks/api/full/rust"] +members = ["packages/api/actor","packages/api/auth","packages/api/cf-verification","packages/api/cloud","packages/api/games","packages/api/group","packages/api/identity","packages/api/job","packages/api/matchmaker","packages/api/monolith-edge","packages/api/monolith-public","packages/api/portal","packages/api/provision","packages/api/status","packages/api/traefik-provider","packages/api/ui","packages/common/api-helper/build","packages/common/api-helper/macros","packages/common/cache/build","packages/common/cache/result","packages/common/chirp-workflow/core","packages/common/chirp-workflow/macros","packages/common/chirp/client","packages/common/chirp/metrics","packages/common/chirp/perf","packages/common/chirp/types","packages/common/chirp/worker","packages/common/chirp/worker-attributes","packages/common/claims","packages/common/config","packages/common/connection","packages/common/convert","packages/common/deno-embed","packages/common/env","packages/common/fdb-util","packages/common/formatted-error","packages/common/global-error","packages/common/health-checks","packages/common/hub-embed","packages/common/kv-str","packages/common/metrics","packages/common/migrate","packages/common/nomad-util","packages/common/operation/core","packages/common/operation/macros","packages/common/pools","packages/common/redis-util","packages/common/runtime","packages/common/s3-util","packages/common/schemac","packages/common/service-manager","packages/common/smithy-output/api-auth/rust","packages/common/smithy-output/api-auth/rust-server","packages/common/smithy-output/api-cf-verification/rust","packages/common/smithy-output/api-cf-verification/rust-server","packages/common/smithy-output/api-cloud/rust","packages/common/smithy-output/api-cloud/rust-server","packages/common/smithy-output/api-group/rust","packages/common/smithy-output/api-group/rust-server","packages/common/smithy-output/api-identity/rust","packages/common/smithy-output/api-identity/rust-server","packages/common/smithy-output/api-job/rust","packages/common/smithy-output/api-job/rust-server","packages/common/smithy-output/api-kv/rust","packages/common/smithy-output/api-kv/rust-server","packages/common/smithy-output/api-matchmaker/rust","packages/common/smithy-output/api-matchmaker/rust-server","packages/common/smithy-output/api-party/rust","packages/common/smithy-output/api-party/rust-server","packages/common/smithy-output/api-portal/rust","packages/common/smithy-output/api-portal/rust-server","packages/common/smithy-output/api-status/rust","packages/common/smithy-output/api-status/rust-server","packages/common/smithy-output/api-traefik-provider/rust","packages/common/smithy-output/api-traefik-provider/rust-server","packages/common/test","packages/common/test-images","packages/common/types-proto/build","packages/common/types-proto/core","packages/common/util/core","packages/common/util/macros","packages/common/util/search","packages/infra/client/actor-kv","packages/infra/client/config","packages/infra/client/container-runner","packages/infra/client/echo","packages/infra/client/isolate-v8-runner","packages/infra/client/logs","packages/infra/client/manager","packages/infra/legacy/job-runner","packages/infra/schema-generator","packages/infra/server","packages/services/build","packages/services/build/ops/create","packages/services/build/ops/get","packages/services/build/ops/list-for-env","packages/services/build/ops/list-for-game","packages/services/build/standalone/default-create","packages/services/build/util","packages/services/captcha/ops/hcaptcha-config-get","packages/services/captcha/ops/hcaptcha-verify","packages/services/captcha/ops/request","packages/services/captcha/ops/turnstile-config-get","packages/services/captcha/ops/turnstile-verify","packages/services/captcha/ops/verify","packages/services/captcha/util","packages/services/cdn/ops/namespace-auth-user-remove","packages/services/cdn/ops/namespace-auth-user-update","packages/services/cdn/ops/namespace-create","packages/services/cdn/ops/namespace-domain-create","packages/services/cdn/ops/namespace-domain-remove","packages/services/cdn/ops/namespace-get","packages/services/cdn/ops/namespace-resolve-domain","packages/services/cdn/ops/ns-auth-type-set","packages/services/cdn/ops/ns-enable-domain-public-auth-set","packages/services/cdn/ops/site-create","packages/services/cdn/ops/site-get","packages/services/cdn/ops/site-list-for-game","packages/services/cdn/ops/version-get","packages/services/cdn/ops/version-prepare","packages/services/cdn/ops/version-publish","packages/services/cdn/util","packages/services/cdn/worker","packages/services/cf-custom-hostname/ops/get","packages/services/cf-custom-hostname/ops/list-for-namespace-id","packages/services/cf-custom-hostname/ops/resolve-hostname","packages/services/cf-custom-hostname/worker","packages/services/cloud/ops/device-link-create","packages/services/cloud/ops/game-config-create","packages/services/cloud/ops/game-config-get","packages/services/cloud/ops/game-token-create","packages/services/cloud/ops/namespace-create","packages/services/cloud/ops/namespace-get","packages/services/cloud/ops/namespace-token-development-create","packages/services/cloud/ops/namespace-token-public-create","packages/services/cloud/ops/version-get","packages/services/cloud/ops/version-publish","packages/services/cloud/standalone/default-create","packages/services/cloud/worker","packages/services/cluster","packages/services/cluster/standalone/datacenter-tls-renew","packages/services/cluster/standalone/default-update","packages/services/cluster/standalone/gc","packages/services/cluster/standalone/metrics-publish","packages/services/custom-user-avatar/ops/list-for-game","packages/services/custom-user-avatar/ops/upload-complete","packages/services/debug/ops/email-res","packages/services/ds","packages/services/ds-log/ops/export","packages/services/ds-log/ops/read","packages/services/dynamic-config","packages/services/email-verification/ops/complete","packages/services/email-verification/ops/create","packages/services/email/ops/send","packages/services/external/ops/request-validate","packages/services/external/worker","packages/services/faker/ops/build","packages/services/faker/ops/cdn-site","packages/services/faker/ops/game","packages/services/faker/ops/game-namespace","packages/services/faker/ops/game-version","packages/services/faker/ops/job-run","packages/services/faker/ops/job-template","packages/services/faker/ops/mm-lobby","packages/services/faker/ops/mm-lobby-row","packages/services/faker/ops/mm-player","packages/services/faker/ops/region","packages/services/faker/ops/team","packages/services/faker/ops/user","packages/services/game/ops/banner-upload-complete","packages/services/game/ops/create","packages/services/game/ops/get","packages/services/game/ops/list-all","packages/services/game/ops/list-for-team","packages/services/game/ops/logo-upload-complete","packages/services/game/ops/namespace-create","packages/services/game/ops/namespace-get","packages/services/game/ops/namespace-list","packages/services/game/ops/namespace-resolve-name-id","packages/services/game/ops/namespace-resolve-url","packages/services/game/ops/namespace-validate","packages/services/game/ops/namespace-version-history-list","packages/services/game/ops/namespace-version-set","packages/services/game/ops/recommend","packages/services/game/ops/resolve-name-id","packages/services/game/ops/resolve-namespace-id","packages/services/game/ops/token-development-validate","packages/services/game/ops/validate","packages/services/game/ops/version-create","packages/services/game/ops/version-get","packages/services/game/ops/version-list","packages/services/game/ops/version-validate","packages/services/ip/ops/info","packages/services/job-log/ops/read","packages/services/job-log/worker","packages/services/job-run","packages/services/job/standalone/gc","packages/services/job/util","packages/services/linode","packages/services/linode/standalone/gc","packages/services/load-test/standalone/api-cloud","packages/services/load-test/standalone/mm","packages/services/load-test/standalone/mm-sustain","packages/services/load-test/standalone/sqlx","packages/services/load-test/standalone/watch-requests","packages/services/mm-config/ops/game-get","packages/services/mm-config/ops/game-upsert","packages/services/mm-config/ops/lobby-group-get","packages/services/mm-config/ops/lobby-group-resolve-name-id","packages/services/mm-config/ops/lobby-group-resolve-version","packages/services/mm-config/ops/namespace-config-set","packages/services/mm-config/ops/namespace-config-validate","packages/services/mm-config/ops/namespace-create","packages/services/mm-config/ops/namespace-get","packages/services/mm-config/ops/version-get","packages/services/mm-config/ops/version-prepare","packages/services/mm-config/ops/version-publish","packages/services/mm/ops/dev-player-token-create","packages/services/mm/ops/lobby-find-fail","packages/services/mm/ops/lobby-find-lobby-query-list","packages/services/mm/ops/lobby-find-try-complete","packages/services/mm/ops/lobby-for-run-id","packages/services/mm/ops/lobby-get","packages/services/mm/ops/lobby-history","packages/services/mm/ops/lobby-idle-update","packages/services/mm/ops/lobby-list-for-namespace","packages/services/mm/ops/lobby-list-for-user-id","packages/services/mm/ops/lobby-player-count","packages/services/mm/ops/lobby-runtime-aggregate","packages/services/mm/ops/lobby-state-get","packages/services/mm/ops/player-count-for-namespace","packages/services/mm/ops/player-get","packages/services/mm/standalone/gc","packages/services/mm/util","packages/services/mm/worker","packages/services/monolith/standalone/worker","packages/services/monolith/standalone/workflow-worker","packages/services/nomad/standalone/monitor","packages/services/pegboard","packages/services/pegboard/standalone/dc-init","packages/services/pegboard/standalone/gc","packages/services/pegboard/standalone/metrics-publish","packages/services/pegboard/standalone/ws","packages/services/region/ops/get","packages/services/region/ops/list","packages/services/region/ops/list-for-game","packages/services/region/ops/recommend","packages/services/region/ops/resolve","packages/services/region/ops/resolve-for-game","packages/services/server-spec","packages/services/team-invite/ops/get","packages/services/team-invite/worker","packages/services/team/ops/avatar-upload-complete","packages/services/team/ops/get","packages/services/team/ops/join-request-list","packages/services/team/ops/member-count","packages/services/team/ops/member-get","packages/services/team/ops/member-list","packages/services/team/ops/member-relationship-get","packages/services/team/ops/profile-validate","packages/services/team/ops/recommend","packages/services/team/ops/resolve-display-name","packages/services/team/ops/user-ban-get","packages/services/team/ops/user-ban-list","packages/services/team/ops/validate","packages/services/team/util","packages/services/team/worker","packages/services/telemetry/standalone/beacon","packages/services/tier","packages/services/token/ops/create","packages/services/token/ops/exchange","packages/services/token/ops/get","packages/services/token/ops/revoke","packages/services/upload/ops/complete","packages/services/upload/ops/file-list","packages/services/upload/ops/get","packages/services/upload/ops/list-for-user","packages/services/upload/ops/prepare","packages/services/upload/worker","packages/services/user","packages/services/user-identity/ops/create","packages/services/user-identity/ops/delete","packages/services/user-identity/ops/get","packages/services/user/ops/avatar-upload-complete","packages/services/user/ops/get","packages/services/user/ops/pending-delete-toggle","packages/services/user/ops/profile-validate","packages/services/user/ops/resolve-email","packages/services/user/ops/team-list","packages/services/user/ops/token-create","packages/services/user/standalone/delete-pending","packages/services/user/worker","packages/services/workflow/standalone/gc","packages/services/workflow/standalone/metrics-publish","packages/toolchain/actors-sdk-embed","packages/toolchain/cli","packages/toolchain/js-utils-embed","packages/toolchain/toolchain","sdks/api/full/rust"] [workspace.package] version = "25.1.0-rc.1" @@ -37,6 +37,10 @@ git = "https://github.com/rivet-gg/sqlx" rev = "e7120f59" default-features = false +[workspace.dependencies.foundationdb] +version = "0.9.1" +features = [ "fdb-7_1", "embedded-fdb-include" ] + [workspace.dependencies.nomad_client] git = "https://github.com/rivet-gg/nomad-client" rev = "abb66bf" @@ -177,6 +181,9 @@ path = "packages/common/deno-embed" [workspace.dependencies.rivet-env] path = "packages/common/env" +[workspace.dependencies.fdb-util] +path = "packages/common/fdb-util" + [workspace.dependencies.formatted-error] path = "packages/common/formatted-error" diff --git a/packages/common/chirp-workflow/core/Cargo.toml b/packages/common/chirp-workflow/core/Cargo.toml index 84d04eb615..2c65243a70 100644 --- a/packages/common/chirp-workflow/core/Cargo.toml +++ b/packages/common/chirp-workflow/core/Cargo.toml @@ -6,13 +6,17 @@ license.workspace = true edition.workspace = true [dependencies] +anyhow = "1.0" async-trait = "0.1.80" chirp-client.workspace = true chirp-workflow-macros.workspace = true cjson = "0.1" +fdb-util.workspace = true formatted-error.workspace = true +foundationdb.workspace = true futures-util = "0.3" global-error.workspace = true +include_dir = "0.7.4" indoc = "2.0.5" lazy_static = "1.4" md5 = "0.7.0" @@ -29,11 +33,13 @@ rivet-runtime.workspace = true rivet-util.workspace = true serde = { version = "1.0.198", features = ["derive"] } serde_json = "1.0.116" +sqlite-util.workspace = true strum = { version = "0.26", features = ["derive"] } thiserror = "1.0.59" tokio = { version = "1.40.0", features = ["full"] } tokio-util = "0.7" tracing = "0.1.40" +tracing-logfmt.workspace = true tracing-subscriber = { version = "0.3.18", features = ["env-filter"] } uuid = { version = "1.8.0", features = ["v4", "serde"] } @@ -44,8 +50,10 @@ features = [ "postgres", "uuid", "json", - "ipnetwork" + "ipnetwork", + "sqlite" ] [dev-dependencies] anyhow = "1.0.82" +rand = "0.8" \ No newline at end of file diff --git a/packages/common/chirp-workflow/core/src/compat.rs b/packages/common/chirp-workflow/core/src/compat.rs index b09555d5bd..1804982fda 100644 --- a/packages/common/chirp-workflow/core/src/compat.rs +++ b/packages/common/chirp-workflow/core/src/compat.rs @@ -12,7 +12,7 @@ use crate::{ common, message::{MessageCtx, SubscriptionHandle}, }, - db::{DatabaseCrdbNats, DatabaseHandle}, + db::{Database, DatabaseCrdbNats, DatabaseHandle}, message::{AsTags, Message}, operation::{Operation, OperationInput}, signal::Signal, @@ -105,20 +105,16 @@ where msg_ctx.subscribe::(tags).await.map_err(GlobalError::raw) } -// Get pool as a trait object async fn db_from_ctx( ctx: &rivet_operation::OperationContext, ) -> GlobalResult { - let crdb = ctx.crdb().await?; - let nats = ctx.conn().nats().await?; - - Ok(DatabaseCrdbNats::from_pools(crdb, nats)) + DatabaseCrdbNats::from_pools(ctx.pools().clone()) + .map(|db| db as DatabaseHandle) + .map_err(Into::into) } -// Get crdb pool as a trait object pub async fn db_from_pools(pools: &rivet_pools::Pools) -> GlobalResult { - let crdb = pools.crdb()?; - let nats = pools.nats()?; - - Ok(DatabaseCrdbNats::from_pools(crdb, nats)) + DatabaseCrdbNats::from_pools(pools.clone()) + .map(|db| db as DatabaseHandle) + .map_err(Into::into) } diff --git a/packages/common/chirp-workflow/core/src/ctx/activity.rs b/packages/common/chirp-workflow/core/src/ctx/activity.rs index 9123859fd3..ce4b9d4c30 100644 --- a/packages/common/chirp-workflow/core/src/ctx/activity.rs +++ b/packages/common/chirp-workflow/core/src/ctx/activity.rs @@ -16,6 +16,7 @@ use crate::{ #[derive(Clone)] pub struct ActivityCtx { workflow_id: Uuid, + workflow_name: String, ray_id: Uuid, name: &'static str, ts: i64, @@ -33,6 +34,7 @@ pub struct ActivityCtx { impl ActivityCtx { pub async fn new( workflow_id: Uuid, + workflow_name: String, db: DatabaseHandle, config: &rivet_config::Config, conn: &rivet_connection::Connection, @@ -60,6 +62,7 @@ impl ActivityCtx { Ok(ActivityCtx { workflow_id, + workflow_name, ray_id, name, ts, @@ -94,9 +97,11 @@ impl ActivityCtx { .await } + // TODO: Theres nothing preventing this from being able to be called from the workflow ctx also, but for + // now its only in the activity ctx so it isn't called again during workflow retries pub async fn update_workflow_tags(&self, tags: &serde_json::Value) -> GlobalResult<()> { self.db - .update_workflow_tags(self.workflow_id, tags) + .update_workflow_tags(self.workflow_id, &self.workflow_name, tags) .await .map_err(GlobalError::raw) } diff --git a/packages/common/chirp-workflow/core/src/ctx/common.rs b/packages/common/chirp-workflow/core/src/ctx/common.rs index a3c6ede8d7..f56ee25594 100644 --- a/packages/common/chirp-workflow/core/src/ctx/common.rs +++ b/packages/common/chirp-workflow/core/src/ctx/common.rs @@ -4,7 +4,7 @@ use global_error::{GlobalError, GlobalResult}; use uuid::Uuid; /// Poll interval when polling for a sub workflow in-process -pub const SUB_WORKFLOW_RETRY: Duration = Duration::from_millis(150); +pub const SUB_WORKFLOW_RETRY: Duration = Duration::from_millis(500); /// Time to delay a workflow from retrying after an error pub const RETRY_TIMEOUT_MS: usize = 2000; pub const WORKFLOW_TIMEOUT: Duration = Duration::from_secs(60); diff --git a/packages/common/chirp-workflow/core/src/ctx/operation.rs b/packages/common/chirp-workflow/core/src/ctx/operation.rs index c4bf7a6fe1..1f16ca81c5 100644 --- a/packages/common/chirp-workflow/core/src/ctx/operation.rs +++ b/packages/common/chirp-workflow/core/src/ctx/operation.rs @@ -88,6 +88,7 @@ impl OperationCtx { /// Creates a signal builder. pub fn signal(&self, body: T) -> builder::signal::SignalBuilder { + // TODO: Add check for from_workflow so you cant dispatch a signal builder::signal::SignalBuilder::new(self.db.clone(), self.ray_id, body) } } diff --git a/packages/common/chirp-workflow/core/src/ctx/test.rs b/packages/common/chirp-workflow/core/src/ctx/test.rs index 0e4d6f8d5b..caa6cbf704 100644 --- a/packages/common/chirp-workflow/core/src/ctx/test.rs +++ b/packages/common/chirp-workflow/core/src/ctx/test.rs @@ -10,7 +10,7 @@ use crate::{ message::{SubscriptionHandle, TailAnchor, TailAnchorResponse}, MessageCtx, }, - db::{DatabaseCrdbNats, DatabaseHandle}, + db::{Database, DatabaseHandle}, message::{AsTags, Message, NatsMessage}, operation::{Operation, OperationInput}, signal::Signal, @@ -34,14 +34,28 @@ pub struct TestCtx { } impl TestCtx { - pub async fn from_env(test_name: &str) -> TestCtx { + pub async fn from_env( + test_name: &str, + no_config: bool, + ) -> TestCtx { let service_name = format!("{}-test--{}", rivet_env::service_name(), test_name); let ray_id = Uuid::new_v4(); - let config = rivet_config::Config::load::(&[]).await.unwrap(); - let pools = rivet_pools::Pools::new(config.clone()) - .await - .expect("failed to create pools"); + let (config, pools) = if no_config { + let config = rivet_config::Config::from_root(rivet_config::config::Root::default()); + let pools = rivet_pools::Pools::test(config.clone()) + .await + .expect("failed to create pools"); + + (config, pools) + } else { + let config = rivet_config::Config::load::(&[]).await.unwrap(); + let pools = rivet_pools::Pools::new(config.clone()) + .await + .expect("failed to create pools"); + + (config, pools) + }; let shared_client = chirp_client::SharedClient::from_env(pools.clone()) .expect("failed to create chirp client"); let cache = @@ -68,10 +82,7 @@ impl TestCtx { (), ); - let db = DatabaseCrdbNats::from_pools( - pools.crdb().unwrap(), - pools.nats_option().clone().unwrap(), - ); + let db = DB::from_pools(pools).unwrap(); let msg_ctx = MessageCtx::new(&conn, ray_id).await.unwrap(); TestCtx { @@ -224,6 +235,10 @@ impl TestCtx { self.conn.cache_handle() } + pub fn pools(&self) -> &rivet_pools::Pools { + self.conn.pools() + } + pub async fn crdb(&self) -> Result { self.conn.crdb().await } diff --git a/packages/common/chirp-workflow/core/src/ctx/workflow.rs b/packages/common/chirp-workflow/core/src/ctx/workflow.rs index 42198c9c85..85cafb2e22 100644 --- a/packages/common/chirp-workflow/core/src/ctx/workflow.rs +++ b/packages/common/chirp-workflow/core/src/ctx/workflow.rs @@ -162,7 +162,11 @@ impl WorkflowCtx { interval.tick().await; // Write output - if let Err(err) = self.db.commit_workflow(self.workflow_id, &output).await { + if let Err(err) = self + .db + .complete_workflow(self.workflow_id, &self.name, &output) + .await + { if retries > MAX_DB_ACTION_RETRIES { return Err(err); } @@ -213,8 +217,9 @@ impl WorkflowCtx { // Write output let res = self .db - .fail_workflow( + .commit_workflow( self.workflow_id, + &self.name, false, deadline_ts, wake_signals, @@ -250,6 +255,7 @@ impl WorkflowCtx { let ctx = ActivityCtx::new( self.workflow_id, + self.name.clone(), self.db.clone(), &self.config, &self.conn, @@ -446,10 +452,10 @@ impl WorkflowCtx { loop { interval.tick().await; - // Check if state finished + // Check if workflow completed let workflow = self .db - .get_workflow(sub_workflow_id) + .get_sub_workflow(self.workflow_id, &self.name, sub_workflow_id) .await .map_err(GlobalError::raw)? .ok_or(WorkflowError::WorkflowNotFound) @@ -788,6 +794,8 @@ impl WorkflowCtx { &loop_location, self.version, 0, + // TODO: + &serde_json::value::RawValue::from_string("null".to_string())?, None, self.loop_location(), ) @@ -851,6 +859,8 @@ impl WorkflowCtx { &loop_location, self.version, iteration, + // TODO: + &serde_json::value::RawValue::from_string("null".to_string())?, None, self.loop_location(), ) @@ -875,6 +885,8 @@ impl WorkflowCtx { &loop_location, self.version, iteration, + // TODO: + &serde_json::value::RawValue::from_string("null".to_string())?, Some(&output_val), self.loop_location(), ) diff --git a/packages/common/chirp-workflow/core/src/db/crdb_nats.rs b/packages/common/chirp-workflow/core/src/db/crdb_nats.rs index 9c22ccd000..4a4476de16 100644 --- a/packages/common/chirp-workflow/core/src/db/crdb_nats.rs +++ b/packages/common/chirp-workflow/core/src/db/crdb_nats.rs @@ -20,7 +20,7 @@ use crate::{ event::{EventId, EventType, SleepState}, location::Location, }, - message, metrics, worker, + metrics, worker, }; // HACK: We alias global error here because its hardcoded into the sql macros @@ -34,6 +34,10 @@ const QUERY_RETRY_MS: usize = 500; const TXN_RETRY: Duration = Duration::from_millis(100); /// Maximum times a query ran by this database adapter is retried. const MAX_QUERY_RETRIES: usize = 16; +/// For SQL macros. +const CONTEXT_NAME: &str = "chirp_workflow_crdb_nats_engine"; +/// For NATS wake mechanism. +const WORKER_WAKE_SUBJECT: &str = "chirp.workflow.crdb_nats.worker.wake"; pub struct DatabaseCrdbNats { pool: PgPool, @@ -42,15 +46,6 @@ pub struct DatabaseCrdbNats { } impl DatabaseCrdbNats { - pub fn from_pools(pool: PgPool, nats: NatsPool) -> Arc { - Arc::new(DatabaseCrdbNats { - pool, - // Lazy load the nats sub - sub: Mutex::new(None), - nats, - }) - } - async fn conn(&self) -> WorkflowResult> { // Attempt to use an existing connection if let Some(conn) = self.pool.try_acquire() { @@ -66,9 +61,9 @@ impl DatabaseCrdbNats { Ok(self.pool.clone()) } - // For sql macro - pub fn name(&self) -> &str { - "chirp_workflow_crdb_nats_engine" + // For SQL macros + fn name(&self) -> &str { + CONTEXT_NAME } /// Spawns a new thread and publishes a worker wake message to nats. @@ -78,10 +73,7 @@ impl DatabaseCrdbNats { let spawn_res = tokio::task::Builder::new().name("wake").spawn( async move { // Fail gracefully - if let Err(err) = nats - .publish(message::WORKER_WAKE_SUBJECT, Vec::new().into()) - .await - { + if let Err(err) = nats.publish(WORKER_WAKE_SUBJECT, Vec::new().into()).await { tracing::warn!(?err, "failed to publish wake message"); } } @@ -139,6 +131,15 @@ impl DatabaseCrdbNats { #[async_trait::async_trait] impl Database for DatabaseCrdbNats { + fn from_pools(pools: rivet_pools::Pools) -> Result, rivet_pools::Error> { + Ok(Arc::new(DatabaseCrdbNats { + pool: pools.crdb()?, + nats: pools.nats()?, + // Lazy load the nats sub + sub: Mutex::new(None), + })) + } + async fn wake(&self) -> WorkflowResult<()> { let mut sub = self.sub.try_lock().map_err(WorkflowError::WakeLock)?; @@ -146,7 +147,7 @@ impl Database for DatabaseCrdbNats { if sub.is_none() { *sub = Some( self.nats - .subscribe(message::WORKER_WAKE_SUBJECT) + .subscribe(WORKER_WAKE_SUBJECT) .await .map_err(|x| WorkflowError::CreateSubscription(x.into()))?, ); @@ -244,6 +245,14 @@ impl Database for DatabaseCrdbNats { .map(|row| row.map(Into::into)) } + async fn find_workflow( + &self, + _workflow_name: &str, + _tags: &serde_json::Value, + ) -> WorkflowResult> { + todo!(); + } + async fn pull_workflows( &self, worker_instance_id: Uuid, @@ -361,7 +370,7 @@ impl Database for DatabaseCrdbNats { last_pull_ts = $3 FROM select_pending_workflows AS pw WHERE w.workflow_id = pw.workflow_id - RETURNING w.workflow_id, workflow_name, create_ts, ray_id, input, wake_deadline_ts + RETURNING w.workflow_id, workflow_name, create_ts, ray_id, input ), -- Update last ping worker_instance_update AS ( @@ -641,9 +650,10 @@ impl Database for DatabaseCrdbNats { Ok(workflows) } - async fn commit_workflow( + async fn complete_workflow( &self, workflow_id: Uuid, + _workflow_name: &str, output: &serde_json::value::RawValue, ) -> WorkflowResult<()> { self.query(|| async { @@ -667,11 +677,12 @@ impl Database for DatabaseCrdbNats { Ok(()) } - async fn fail_workflow( + async fn commit_workflow( &self, workflow_id: Uuid, + _workflow_name: &str, immediate: bool, - deadline_ts: Option, + wake_deadline_ts: Option, wake_signals: &[&str], wake_sub_workflow_id: Option, error: &str, @@ -692,7 +703,7 @@ impl Database for DatabaseCrdbNats { )) .bind(workflow_id) .bind(immediate) - .bind(deadline_ts) + .bind(wake_deadline_ts) .bind(wake_signals) .bind(wake_sub_workflow_id) .bind(error) @@ -702,139 +713,6 @@ impl Database for DatabaseCrdbNats { }) .await?; - // Wake worker again if the deadline is before the next tick - if let Some(deadline_ts) = deadline_ts { - if deadline_ts - < rivet_util::timestamp::now() + worker::TICK_INTERVAL.as_millis() as i64 + 1 - { - self.wake_worker(); - } - } - - Ok(()) - } - - // TODO: Theres nothing preventing this from being able to be called from the workflow ctx also, but for - // now its only in the activity ctx so it isn't called again during workflow retries - async fn update_workflow_tags( - &self, - workflow_id: Uuid, - tags: &serde_json::Value, - ) -> WorkflowResult<()> { - self.query(|| async { - sqlx::query(indoc!( - " - UPDATE db_workflow.workflows - SET tags = $2 - WHERE workflow_id = $1 - ", - )) - .bind(workflow_id) - .bind(tags) - .execute(&mut *self.conn().await?) - .await - .map_err(WorkflowError::Sqlx) - }) - .await?; - - Ok(()) - } - - async fn commit_workflow_activity_event( - &self, - workflow_id: Uuid, - location: &Location, - version: usize, - event_id: &EventId, - create_ts: i64, - input: &serde_json::value::RawValue, - res: Result<&serde_json::value::RawValue, &str>, - loop_location: Option<&Location>, - ) -> WorkflowResult<()> { - match res { - Ok(output) => { - self.query(|| async { - sqlx::query(indoc!( - " - INSERT INTO db_workflow.workflow_activity_events ( - workflow_id, - location2, - version, - activity_name, - input_hash, - input, - output, - create_ts, - loop_location2 - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) - ON CONFLICT (workflow_id, location2_hash) DO UPDATE - SET output = EXCLUDED.output - ", - )) - .bind(workflow_id) - .bind(location) - .bind(version as i64) - .bind(&event_id.name) - .bind(event_id.input_hash.to_le_bytes()) - .bind(sqlx::types::Json(input)) - .bind(sqlx::types::Json(output)) - .bind(create_ts) - .bind(loop_location) - .execute(&mut *self.conn().await?) - .await - .map_err(WorkflowError::Sqlx) - }) - .await?; - } - Err(err) => { - self.query(|| async { - sqlx::query(indoc!( - " - WITH - event AS ( - INSERT INTO db_workflow.workflow_activity_events ( - workflow_id, - location2, - version, - activity_name, - input_hash, - input, - create_ts, - loop_location2 - ) - VALUES ($1, $2, $3, $4, $5, $6, $8, $9) - ON CONFLICT (workflow_id, location2_hash) DO NOTHING - RETURNING 1 - ), - err AS ( - INSERT INTO db_workflow.workflow_activity_errors ( - workflow_id, location2, activity_name, error, ts - ) - VALUES ($1, $2, $4, $7, $10) - RETURNING 1 - ) - SELECT 1 - ", - )) - .bind(workflow_id) - .bind(location) - .bind(version as i64) - .bind(&event_id.name) - .bind(event_id.input_hash.to_le_bytes()) - .bind(sqlx::types::Json(input)) - .bind(err) - .bind(create_ts) - .bind(loop_location) - .bind(rivet_util::timestamp::now()) - .execute(&mut *self.conn().await?) - .await - .map_err(WorkflowError::Sqlx) - }) - .await?; - } - } - Ok(()) } @@ -938,6 +816,15 @@ impl Database for DatabaseCrdbNats { Ok(signal) } + async fn get_sub_workflow( + &self, + _workflow_id: Uuid, + _workflow_name: &str, + sub_workflow_id: Uuid, + ) -> WorkflowResult> { + self.get_workflow(sub_workflow_id).await + } + async fn publish_signal( &self, ray_id: Uuid, @@ -1082,7 +969,7 @@ impl Database for DatabaseCrdbNats { RETURNING 1 ), send_event AS ( - INSERT INTO db_workflow.workflow_signal_send_events( + INSERT INTO db_workflow.workflow_signal_send_events ( workflow_id, location2, version, signal_id, signal_name, body, loop_location2 ) VALUES($7, $8, $9, $1, $3, $4, $10) @@ -1146,7 +1033,7 @@ impl Database for DatabaseCrdbNats { RETURNING workflow_id ), insert_sub_workflow_event AS ( - INSERT INTO db_workflow.workflow_sub_workflow_events( + INSERT INTO db_workflow.workflow_sub_workflow_events ( workflow_id, location2, version, sub_workflow_id, create_ts, loop_location2 ) SELECT $1, $7, $8, $9, $3, $10 @@ -1171,7 +1058,7 @@ impl Database for DatabaseCrdbNats { RETURNING workflow_id ), insert_sub_workflow_event AS ( - INSERT INTO db_workflow.workflow_sub_workflow_events( + INSERT INTO db_workflow.workflow_sub_workflow_events ( workflow_id, location2, version, sub_workflow_id, create_ts, loop_location2 ) VALUES($1, $7, $8, $9, $3, $10) @@ -1208,6 +1095,129 @@ impl Database for DatabaseCrdbNats { Ok(actual_sub_workflow_id) } + async fn update_workflow_tags( + &self, + workflow_id: Uuid, + _workflow_name: &str, + tags: &serde_json::Value, + ) -> WorkflowResult<()> { + self.query(|| async { + sqlx::query(indoc!( + " + UPDATE db_workflow.workflows + SET tags = $2 + WHERE workflow_id = $1 + ", + )) + .bind(workflow_id) + .bind(tags) + .execute(&mut *self.conn().await?) + .await + .map_err(WorkflowError::Sqlx) + }) + .await?; + + Ok(()) + } + + async fn commit_workflow_activity_event( + &self, + workflow_id: Uuid, + location: &Location, + version: usize, + event_id: &EventId, + create_ts: i64, + input: &serde_json::value::RawValue, + res: Result<&serde_json::value::RawValue, &str>, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + match res { + Ok(output) => { + self.query(|| async { + sqlx::query(indoc!( + " + INSERT INTO db_workflow.workflow_activity_events ( + workflow_id, + location2, + version, + activity_name, + input_hash, + input, + output, + create_ts, + loop_location2 + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + ON CONFLICT (workflow_id, location2_hash) DO UPDATE + SET output = EXCLUDED.output + ", + )) + .bind(workflow_id) + .bind(location) + .bind(version as i64) + .bind(&event_id.name) + .bind(event_id.input_hash.to_le_bytes()) + .bind(sqlx::types::Json(input)) + .bind(sqlx::types::Json(output)) + .bind(create_ts) + .bind(loop_location) + .execute(&mut *self.conn().await?) + .await + .map_err(WorkflowError::Sqlx) + }) + .await?; + } + Err(err) => { + self.query(|| async { + sqlx::query(indoc!( + " + WITH + event AS ( + INSERT INTO db_workflow.workflow_activity_events ( + workflow_id, + location2, + version, + activity_name, + input_hash, + input, + create_ts, + loop_location2 + ) + VALUES ($1, $2, $3, $4, $5, $6, $8, $9) + ON CONFLICT (workflow_id, location2_hash) DO NOTHING + RETURNING 1 + ), + err AS ( + INSERT INTO db_workflow.workflow_activity_errors ( + workflow_id, location2, activity_name, error, ts + ) + VALUES ($1, $2, $4, $7, $10) + RETURNING 1 + ) + SELECT 1 + ", + )) + .bind(workflow_id) + .bind(location) + .bind(version as i64) + .bind(&event_id.name) + .bind(event_id.input_hash.to_le_bytes()) + .bind(sqlx::types::Json(input)) + .bind(err) + .bind(create_ts) + .bind(loop_location) + .bind(rivet_util::timestamp::now()) + .execute(&mut *self.conn().await?) + .await + .map_err(WorkflowError::Sqlx) + }) + .await?; + } + } + + Ok(()) + } + async fn commit_workflow_message_send_event( &self, from_workflow_id: Uuid, @@ -1221,7 +1231,7 @@ impl Database for DatabaseCrdbNats { self.query(|| async { sqlx::query(indoc!( " - INSERT INTO db_workflow.workflow_message_send_events( + INSERT INTO db_workflow.workflow_message_send_events ( workflow_id, location2, version, tags, message_name, body, loop_location2 ) VALUES($1, $2, $3, $4, $5, $6, $7) @@ -1250,6 +1260,8 @@ impl Database for DatabaseCrdbNats { location: &Location, version: usize, iteration: usize, + // TODO: + _state: &serde_json::value::RawValue, output: Option<&serde_json::value::RawValue>, loop_location: Option<&Location>, ) -> WorkflowResult<()> { @@ -1408,7 +1420,7 @@ impl Database for DatabaseCrdbNats { self.query(|| async { sqlx::query(indoc!( " - INSERT INTO db_workflow.workflow_sleep_events( + INSERT INTO db_workflow.workflow_sleep_events ( workflow_id, location2, version, deadline_ts, loop_location2, state ) VALUES($1, $2, $3, $4, $5, $6) @@ -1466,7 +1478,7 @@ impl Database for DatabaseCrdbNats { self.query(|| async { sqlx::query(indoc!( " - INSERT INTO db_workflow.workflow_branch_events( + INSERT INTO db_workflow.workflow_branch_events ( workflow_id, location, version, loop_location ) VALUES($1, $2, $3, $4) @@ -1497,7 +1509,7 @@ impl Database for DatabaseCrdbNats { self.query(|| async { sqlx::query(indoc!( " - INSERT INTO db_workflow.workflow_removed_events( + INSERT INTO db_workflow.workflow_removed_events ( workflow_id, location, event_type, event_name, loop_location ) VALUES($1, $2, $3, $4, $5) @@ -1528,7 +1540,7 @@ impl Database for DatabaseCrdbNats { self.query(|| async { sqlx::query(indoc!( " - INSERT INTO db_workflow.workflow_version_check_events( + INSERT INTO db_workflow.workflow_version_check_events ( workflow_id, location, version, loop_location ) VALUES($1, $2, $3, $4) @@ -1593,7 +1605,6 @@ mod types { create_ts: i64, ray_id: Uuid, input: RawJson, - wake_deadline_ts: Option, } #[derive(sqlx::FromRow)] @@ -1685,7 +1696,7 @@ mod types { fn try_from(value: AmalgamEventRow) -> WorkflowResult { Ok(ActivityEvent { - event_id: EventId::from_bytes( + event_id: EventId::from_le_bytes( value.name.ok_or(WorkflowError::MissingEventData)?, value.hash.ok_or(WorkflowError::MissingEventData)?, )?, @@ -1751,6 +1762,9 @@ mod types { fn try_from(value: AmalgamEventRow) -> WorkflowResult { Ok(LoopEvent { + // TODO: + state: serde_json::value::RawValue::from_string("null".to_string()) + .map_err(WorkflowError::SerializeLoopState)?, output: value.output.map(|x| x.0), iteration: value .iteration @@ -1791,63 +1805,6 @@ mod types { } } - // Implements sqlx postgres types for `Location` - impl sqlx::Type for Location - where - DB: sqlx::Database, - serde_json::Value: sqlx::Type, - { - fn type_info() -> DB::TypeInfo { - >::type_info() - } - - fn compatible(ty: &DB::TypeInfo) -> bool { - >::compatible(ty) - } - } - - impl<'q> sqlx::Encode<'q, sqlx::Postgres> for Location { - fn encode_by_ref( - &self, - buf: &mut sqlx::postgres::PgArgumentBuffer, - ) -> Result { - >::encode( - serialize_location(self), - buf, - ) - } - } - - impl sqlx::Decode<'_, sqlx::Postgres> for Location { - fn decode(value: sqlx::postgres::PgValueRef) -> Result { - let value = - ]>> as sqlx::Decode>::decode( - value, - )?; - - Ok(IntoIterator::into_iter(value.0).collect()) - } - } - - // TODO: Implement serde serialize and deserialize for `Location` - /// Convert location to json as `number[][]`. - pub fn serialize_location(location: &Location) -> serde_json::Value { - serde_json::Value::Array( - location - .as_ref() - .iter() - .map(|coord| { - serde_json::Value::Array( - coord - .iter() - .map(|x| serde_json::Value::Number((*x).into())) - .collect(), - ) - }) - .collect(), - ) - } - // IMPORTANT: Must match the hashing algorithm used in the `db-workflow` `loop_location2_hash` generated /// column expression. pub fn hash_location(location: &Location) -> Vec { @@ -2009,7 +1966,6 @@ mod types { create_ts: row.create_ts, ray_id: row.ray_id, input: row.input.0, - wake_deadline_ts: row.wake_deadline_ts, events: events_by_location, } }) diff --git a/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/mod.rs b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/mod.rs new file mode 100644 index 0000000000..cdcd2d7db2 --- /dev/null +++ b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/mod.rs @@ -0,0 +1,27 @@ +// TODO: Use concrete error types +use anyhow::*; +use foundationdb::future::FdbValue; + +pub mod signal; +pub mod wake; +pub mod workflow; + +pub trait FormalKey { + type Value; + + fn deserialize(&self, raw: &[u8]) -> Result; + + fn serialize(&self, value: Self::Value) -> Result>; +} + +pub trait FormalChunkedKey { + type Value; + type ChunkKey; + + fn chunk(&self, chunk: usize) -> Self::ChunkKey; + + /// Assumes chunks are in order. + fn combine(&self, chunks: Vec) -> Result; + + // fn split(&self, value: Self::Value) -> Result>>; +} diff --git a/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/signal.rs b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/signal.rs new file mode 100644 index 0000000000..76078c654c --- /dev/null +++ b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/signal.rs @@ -0,0 +1,378 @@ +use std::{borrow::Cow, result::Result::Ok}; + +// TODO: Use concrete error types +use anyhow::*; +use foundationdb::{ + future::FdbValue, + tuple::{PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset}, +}; +use uuid::Uuid; + +use super::{FormalChunkedKey, FormalKey}; + +pub struct BodyKey { + signal_id: Uuid, +} + +impl BodyKey { + pub fn new(signal_id: Uuid) -> Self { + BodyKey { signal_id } + } + + pub fn split_ref(&self, value: &serde_json::value::RawValue) -> Result>> { + // TODO: Chunk + Ok(vec![value.get().as_bytes().to_vec()]) + } +} + +impl FormalChunkedKey for BodyKey { + type ChunkKey = BodyChunkKey; + type Value = Box; + + fn chunk(&self, chunk: usize) -> Self::ChunkKey { + BodyChunkKey { + signal_id: self.signal_id, + chunk, + } + } + + fn combine(&self, chunks: Vec) -> Result { + serde_json::value::RawValue::from_string(String::from_utf8( + chunks + .iter() + .map(|x| x.value().iter().map(|x| *x)) + .flatten() + .collect(), + )?) + .map_err(Into::into) + } + + // fn split(&self, value: Self::Value) -> Result>> { + // self.split_ref(value.as_ref()) + // } +} + +impl TuplePack for BodyKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("signal", "data", self.signal_id, "body"); + t.pack(w, tuple_depth) + } +} + +pub struct BodyChunkKey { + signal_id: Uuid, + chunk: usize, +} + +impl TuplePack for BodyChunkKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("signal", "data", self.signal_id, "body", self.chunk); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for BodyChunkKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, signal_id, _, chunk)) = + <(Cow, Cow, Uuid, Cow, usize)>::unpack(input, tuple_depth)?; + let v = BodyChunkKey { signal_id, chunk }; + + Ok((input, v)) + } +} + +pub struct AckKey { + signal_id: Uuid, +} + +impl AckKey { + pub fn new(signal_id: Uuid) -> Self { + AckKey { signal_id } + } +} + +impl FormalKey for AckKey { + type Value = (); + + fn deserialize(&self, _raw: &[u8]) -> Result { + Ok(()) + } + + fn serialize(&self, _value: Self::Value) -> Result> { + Ok(Vec::new()) + } +} + +impl TuplePack for AckKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("signal", "data", self.signal_id, "ack"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for AckKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, signal_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = AckKey { signal_id }; + + Ok((input, v)) + } +} + +#[derive(Debug)] +pub struct CreateTsKey { + signal_id: Uuid, +} + +impl CreateTsKey { + pub fn new(signal_id: Uuid) -> Self { + CreateTsKey { signal_id } + } +} + +impl FormalKey for CreateTsKey { + // Timestamp. + type Value = i64; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(i64::from_be_bytes(raw.try_into()?)) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.to_be_bytes().to_vec()) + } +} + +impl TuplePack for CreateTsKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("signal", "data", self.signal_id, "create_ts"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for CreateTsKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, signal_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = CreateTsKey { signal_id }; + + Ok((input, v)) + } +} + +pub struct TaggedPendingKey { + pub signal_name: String, + /// For ordering. + pub ts: i64, + pub signal_id: Uuid, +} + +impl TaggedPendingKey { + pub fn new(signal_name: String, signal_id: Uuid) -> Self { + TaggedPendingKey { + signal_name, + ts: rivet_util::timestamp::now(), + signal_id, + } + } + + pub fn subspace(signal_name: String) -> TaggedPendingSubspaceKey { + TaggedPendingSubspaceKey::new(signal_name) + } +} + +impl FormalKey for TaggedPendingKey { + /// Signal tags. + type Value = Vec<(String, String)>; + + fn deserialize(&self, raw: &[u8]) -> Result { + serde_json::from_slice(raw).map_err(Into::into) + } + + fn serialize(&self, value: Self::Value) -> Result> { + serde_json::to_vec(&value).map_err(Into::into) + } +} + +impl TuplePack for TaggedPendingKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "tagged_signal", + "pending", + &self.signal_name, + self.ts, + self.signal_id, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for TaggedPendingKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, signal_name, ts, signal_id)) = + <(Cow, Cow, String, i64, Uuid)>::unpack(input, tuple_depth)?; + let v = TaggedPendingKey { + signal_name, + ts, + signal_id, + }; + + Ok((input, v)) + } +} + +pub struct TaggedPendingSubspaceKey { + signal_name: String, +} + +impl TaggedPendingSubspaceKey { + fn new(signal_name: String) -> Self { + TaggedPendingSubspaceKey { signal_name } + } +} + +impl TuplePack for TaggedPendingSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("tagged_signal", "pending", &self.signal_name); + t.pack(w, tuple_depth) + } +} + +pub struct TagKey { + signal_id: Uuid, + pub k: String, + pub v: String, +} + +impl TagKey { + pub fn new(signal_id: Uuid, k: String, v: String) -> Self { + TagKey { signal_id, k, v } + } + + // pub fn subspace(signal_id: Uuid) -> TagSubspaceKey { + // TagSubspaceKey::new(signal_id) + // } +} + +impl FormalKey for TagKey { + type Value = (); + + fn deserialize(&self, _raw: &[u8]) -> Result { + Ok(()) + } + + fn serialize(&self, _value: Self::Value) -> Result> { + Ok(Vec::new()) + } +} + +impl TuplePack for TagKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("signal", "data", self.signal_id, "tag", &self.k, &self.v); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for TagKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, signal_id, _, k, v)) = + <(Cow, Cow, Uuid, Cow, String, String)>::unpack(input, tuple_depth)?; + let v = TagKey { signal_id, k, v }; + + Ok((input, v)) + } +} + +// pub struct TagSubspaceKey { +// signal_id: Uuid, +// } + +// impl TagSubspaceKey { +// pub fn new(signal_id: Uuid) -> Self { +// TagSubspaceKey { signal_id } +// } +// } + +// impl TuplePack for TagSubspaceKey { +// fn pack( +// &self, +// w: &mut W, +// tuple_depth: TupleDepth, +// ) -> std::io::Result { +// let t = ("signal", "data", self.signal_id, "tag"); +// t.pack(w, tuple_depth) +// } +// } + +#[derive(Debug)] +pub struct RayIdKey { + signal_id: Uuid, +} + +impl RayIdKey { + pub fn new(signal_id: Uuid) -> Self { + RayIdKey { signal_id } + } +} + +impl FormalKey for RayIdKey { + type Value = Uuid; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(Uuid::from_slice(raw)?) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.as_bytes().to_vec()) + } +} + +impl TuplePack for RayIdKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("signal", "data", self.signal_id, "ray_id"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for RayIdKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, signal_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = RayIdKey { signal_id }; + + Ok((input, v)) + } +} diff --git a/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/wake.rs b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/wake.rs new file mode 100644 index 0000000000..49b32cbb2a --- /dev/null +++ b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/wake.rs @@ -0,0 +1,434 @@ +use std::{borrow::Cow, result::Result::Ok}; + +// TODO: Use concrete error types +use anyhow::*; +use foundationdb::tuple::{ + PackError, PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset, +}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use super::FormalKey; + +#[derive(Debug)] +pub enum WakeCondition { + Immediate, + Deadline { deadline_ts: i64 }, + SubWorkflow { sub_workflow_id: Uuid }, + Signal { signal_id: Uuid }, + TaggedSignal { signal_id: Uuid }, +} + +#[derive(Debug)] +pub struct WorkflowWakeConditionKey { + pub workflow_name: String, + pub ts: i64, + pub workflow_id: Uuid, + pub condition: WakeCondition, +} + +impl WorkflowWakeConditionKey { + pub fn new(workflow_name: String, workflow_id: Uuid, condition: WakeCondition) -> Self { + WorkflowWakeConditionKey { + workflow_name, + // NOTE: Override current ts with deadline + ts: if let WakeCondition::Deadline { deadline_ts } = &condition { + *deadline_ts + } else { + rivet_util::timestamp::now() + }, + workflow_id, + condition, + } + } + + pub fn subspace(workflow_name: String, ts: i64) -> WorkflowWakeConditionSubspaceKey { + WorkflowWakeConditionSubspaceKey::new(workflow_name, ts) + } + + pub fn subspace_without_ts(workflow_name: String) -> WorkflowWakeConditionSubspaceKey { + WorkflowWakeConditionSubspaceKey::new_without_ts(workflow_name) + } +} + +impl FormalKey for WorkflowWakeConditionKey { + type Value = (); + + fn deserialize(&self, _raw: &[u8]) -> Result { + Ok(()) + } + + fn serialize(&self, _value: Self::Value) -> Result> { + Ok(Vec::new()) + } +} + +impl TuplePack for WorkflowWakeConditionKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + match &self.condition { + WakeCondition::Immediate => { + let t = ( + "wake", + "workflow", + &self.workflow_name, + self.ts, + self.workflow_id, + "immediate", + ); + t.pack(w, tuple_depth) + } + WakeCondition::Deadline { .. } => { + let t = ( + "wake", + "workflow", + &self.workflow_name, + // Already matches deadline ts, see `WorkflowWakeConditionKey::new` + self.ts, + self.workflow_id, + "deadline", + ); + t.pack(w, tuple_depth) + } + WakeCondition::SubWorkflow { sub_workflow_id } => { + let t = ( + "wake", + "workflow", + &self.workflow_name, + self.ts, + self.workflow_id, + "sub_workflow", + sub_workflow_id, + ); + t.pack(w, tuple_depth) + } + WakeCondition::Signal { signal_id } => { + let t = ( + "wake", + "workflow", + &self.workflow_name, + self.ts, + self.workflow_id, + "signal", + signal_id, + ); + t.pack(w, tuple_depth) + } + WakeCondition::TaggedSignal { signal_id } => { + let t = ( + "wake", + "workflow", + &self.workflow_name, + self.ts, + self.workflow_id, + "tagged_signal", + signal_id, + ); + t.pack(w, tuple_depth) + } + } + } +} + +impl<'de> TupleUnpack<'de> for WorkflowWakeConditionKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_name, ts, workflow_id, condition_name)) = + <(Cow, Cow, String, i64, Uuid, String)>::unpack(input, tuple_depth)?; + + let (input, v) = match &*condition_name { + "immediate" => ( + input, + WorkflowWakeConditionKey { + workflow_name, + ts, + workflow_id, + condition: WakeCondition::Immediate, + }, + ), + "deadline" => ( + input, + WorkflowWakeConditionKey { + workflow_name, + ts, + workflow_id, + condition: WakeCondition::Deadline { + // See `WorkflowWakeConditionKey::new` + deadline_ts: ts, + }, + }, + ), + "sub_workflow" => { + let (input, sub_workflow_id) = Uuid::unpack(input, tuple_depth)?; + + ( + input, + WorkflowWakeConditionKey { + workflow_name, + ts, + workflow_id, + condition: WakeCondition::SubWorkflow { sub_workflow_id }, + }, + ) + } + "signal" => { + let (input, signal_id) = Uuid::unpack(input, tuple_depth)?; + + ( + input, + WorkflowWakeConditionKey { + workflow_name, + ts, + workflow_id, + condition: WakeCondition::Signal { signal_id }, + }, + ) + } + "tagged_signal" => { + let (input, signal_id) = Uuid::unpack(input, tuple_depth)?; + + ( + input, + WorkflowWakeConditionKey { + workflow_name, + ts, + workflow_id, + condition: WakeCondition::TaggedSignal { signal_id }, + }, + ) + } + _ => { + return Err(PackError::Message( + format!("invalid wake condition type {condition_name:?} in key").into(), + )) + } + }; + + Ok((input, v)) + } +} + +// Structure should match `WorkflowWakeConditionKey` +pub struct WorkflowWakeConditionSubspaceKey { + workflow_name: String, + ts: Option, +} + +impl WorkflowWakeConditionSubspaceKey { + pub fn new(workflow_name: String, ts: i64) -> Self { + WorkflowWakeConditionSubspaceKey { + workflow_name, + ts: Some(ts), + } + } + + pub fn new_without_ts(workflow_name: String) -> Self { + WorkflowWakeConditionSubspaceKey { + workflow_name, + ts: None, + } + } +} + +impl TuplePack for WorkflowWakeConditionSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let mut offset = VersionstampOffset::None { size: 0 }; + + let t = ("wake", "workflow", &self.workflow_name); + offset += t.pack(w, tuple_depth)?; + + if let Some(ts) = &self.ts { + offset += ts.pack(w, tuple_depth)?; + } + + Ok(offset) + } +} + +#[derive(Debug)] +pub struct TaggedSignalWakeKey { + pub signal_name: String, + /// For ordering. + pub ts: i64, + pub workflow_id: Uuid, +} + +impl TaggedSignalWakeKey { + pub fn new(signal_name: String, workflow_id: Uuid) -> Self { + TaggedSignalWakeKey { + signal_name, + ts: rivet_util::timestamp::now(), + workflow_id, + } + } + + pub fn subspace(signal_name: String) -> TaggedSignalWakeSubspaceKey { + TaggedSignalWakeSubspaceKey::new(signal_name) + } +} + +impl FormalKey for TaggedSignalWakeKey { + /// Workflow name and tags. + type Value = TaggedSignalWakeData; + + fn deserialize(&self, raw: &[u8]) -> Result { + serde_json::from_slice(raw).map_err(Into::into) + } + + fn serialize(&self, value: Self::Value) -> Result> { + serde_json::to_vec(&value).map_err(Into::into) + } +} + +impl TuplePack for TaggedSignalWakeKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "wake", + "tagged_signal", + &self.signal_name, + self.ts, + self.workflow_id, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for TaggedSignalWakeKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, signal_name, ts, workflow_id)) = + <(Cow, Cow, String, i64, Uuid)>::unpack(input, tuple_depth)?; + let v = TaggedSignalWakeKey { + signal_name, + ts, + workflow_id, + }; + + Ok((input, v)) + } +} + +#[derive(Serialize, Deserialize)] +pub struct TaggedSignalWakeData { + pub workflow_name: String, + pub tags: Vec<(String, String)>, +} + +#[derive(Debug)] +pub struct TaggedSignalWakeSubspaceKey { + signal_name: String, +} + +impl TaggedSignalWakeSubspaceKey { + pub fn new(signal_name: String) -> Self { + TaggedSignalWakeSubspaceKey { signal_name } + } +} + +impl TuplePack for TaggedSignalWakeSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("wake", "tagged_signal", &self.signal_name); + t.pack(w, tuple_depth) + } +} + +#[derive(Debug)] +pub struct SubWorkflowWakeKey { + pub sub_workflow_id: Uuid, + /// For ordering. + pub ts: i64, + pub workflow_id: Uuid, +} + +impl SubWorkflowWakeKey { + pub fn new(sub_workflow_id: Uuid, workflow_id: Uuid) -> Self { + SubWorkflowWakeKey { + sub_workflow_id, + ts: rivet_util::timestamp::now(), + workflow_id, + } + } + + pub fn subspace(sub_workflow_id: Uuid) -> SubWorkflowWakeSubspaceKey { + SubWorkflowWakeSubspaceKey::new(sub_workflow_id) + } +} + +impl FormalKey for SubWorkflowWakeKey { + /// Workflow name (not sub workflow name). + type Value = String; + + fn deserialize(&self, raw: &[u8]) -> Result { + String::from_utf8(raw.to_vec()).map_err(Into::into) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.into_bytes()) + } +} + +impl TuplePack for SubWorkflowWakeKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "wake", + "sub_workflow", + &self.sub_workflow_id, + self.ts, + self.workflow_id, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for SubWorkflowWakeKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, sub_workflow_id, ts, workflow_id)) = + <(Cow, Cow, Uuid, i64, Uuid)>::unpack(input, tuple_depth)?; + let v = SubWorkflowWakeKey { + sub_workflow_id, + ts, + workflow_id, + }; + + Ok((input, v)) + } +} + +#[derive(Debug)] +pub struct SubWorkflowWakeSubspaceKey { + sub_workflow_id: Uuid, +} + +impl SubWorkflowWakeSubspaceKey { + pub fn new(sub_workflow_id: Uuid) -> Self { + SubWorkflowWakeSubspaceKey { sub_workflow_id } + } +} + +impl TuplePack for SubWorkflowWakeSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("wake", "sub_workflow", self.sub_workflow_id); + t.pack(w, tuple_depth) + } +} diff --git a/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/workflow.rs b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/workflow.rs new file mode 100644 index 0000000000..c270559717 --- /dev/null +++ b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/keys/workflow.rs @@ -0,0 +1,813 @@ +use std::{borrow::Cow, result::Result::Ok}; + +// TODO: Use concrete error types +use anyhow::*; +use foundationdb::{ + future::FdbValue, + tuple::{PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset}, +}; +use uuid::Uuid; + +use super::{FormalChunkedKey, FormalKey}; + +pub struct LeaseKey { + workflow_id: Uuid, +} + +impl LeaseKey { + pub fn new(workflow_id: Uuid) -> Self { + LeaseKey { workflow_id } + } +} + +impl FormalKey for LeaseKey { + type Value = Uuid; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(Uuid::from_slice(raw)?) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.as_bytes().to_vec()) + } +} + +impl TuplePack for LeaseKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "lease", self.workflow_id); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for LeaseKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id)) = + <(Cow, Cow, Uuid)>::unpack(input, tuple_depth)?; + let v = LeaseKey { workflow_id }; + + Ok((input, v)) + } +} + +pub struct TagKey { + workflow_id: Uuid, + pub k: String, + pub v: String, +} + +impl TagKey { + pub fn new(workflow_id: Uuid, k: String, v: String) -> Self { + TagKey { workflow_id, k, v } + } + + pub fn subspace(workflow_id: Uuid) -> TagSubspaceKey { + TagSubspaceKey::new(workflow_id) + } +} + +impl FormalKey for TagKey { + type Value = (); + + fn deserialize(&self, _raw: &[u8]) -> Result { + Ok(()) + } + + fn serialize(&self, _value: Self::Value) -> Result> { + Ok(Vec::new()) + } +} + +impl TuplePack for TagKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "workflow", + "data", + self.workflow_id, + "tag", + &self.k, + &self.v, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for TagKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _, k, v)) = + <(Cow, Cow, Uuid, Cow, String, String)>::unpack(input, tuple_depth)?; + let v = TagKey { workflow_id, k, v }; + + Ok((input, v)) + } +} + +pub struct TagSubspaceKey { + workflow_id: Uuid, +} + +impl TagSubspaceKey { + pub fn new(workflow_id: Uuid) -> Self { + TagSubspaceKey { workflow_id } + } +} + +impl TuplePack for TagSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "tag"); + t.pack(w, tuple_depth) + } +} + +pub struct InputKey { + workflow_id: Uuid, +} + +impl InputKey { + pub fn new(workflow_id: Uuid) -> Self { + InputKey { workflow_id } + } + + pub fn split_ref(&self, value: &serde_json::value::RawValue) -> Result>> { + // TODO: Chunk + Ok(vec![value.get().as_bytes().to_vec()]) + } +} + +impl FormalChunkedKey for InputKey { + type ChunkKey = InputChunkKey; + type Value = Box; + + fn chunk(&self, chunk: usize) -> Self::ChunkKey { + InputChunkKey { + workflow_id: self.workflow_id, + chunk, + } + } + + fn combine(&self, chunks: Vec) -> Result { + serde_json::value::RawValue::from_string(String::from_utf8( + chunks + .iter() + .map(|x| x.value().iter().map(|x| *x)) + .flatten() + .collect(), + )?) + .map_err(Into::into) + } + + // fn split(&self, value: Self::Value) -> Result>> { + // self.split_ref(value.as_ref()) + // } +} + +impl TuplePack for InputKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "input"); + t.pack(w, tuple_depth) + } +} + +pub struct InputChunkKey { + workflow_id: Uuid, + chunk: usize, +} + +impl TuplePack for InputChunkKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "input", self.chunk); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for InputChunkKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, workflow_id, _, chunk)) = + <(Cow, Uuid, Cow, usize)>::unpack(input, tuple_depth)?; + let v = InputChunkKey { workflow_id, chunk }; + + Ok((input, v)) + } +} + +pub struct OutputKey { + workflow_id: Uuid, +} + +impl OutputKey { + pub fn new(workflow_id: Uuid) -> Self { + OutputKey { workflow_id } + } + + pub fn split_ref(&self, value: &serde_json::value::RawValue) -> Result>> { + // TODO: Chunk + Ok(vec![value.get().as_bytes().to_vec()]) + } +} + +impl FormalChunkedKey for OutputKey { + type ChunkKey = OutputChunkKey; + type Value = Box; + + fn chunk(&self, chunk: usize) -> Self::ChunkKey { + OutputChunkKey { + workflow_id: self.workflow_id, + chunk, + } + } + + fn combine(&self, chunks: Vec) -> Result { + serde_json::value::RawValue::from_string(String::from_utf8( + chunks + .iter() + .map(|x| x.value().iter().map(|x| *x)) + .flatten() + .collect(), + )?) + .map_err(Into::into) + } + + // fn split(&self, value: Self::Value) -> Result>> { + // self.split_ref(value.as_ref()) + // } +} + +impl TuplePack for OutputKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "output"); + t.pack(w, tuple_depth) + } +} + +pub struct OutputChunkKey { + workflow_id: Uuid, + chunk: usize, +} + +impl TuplePack for OutputChunkKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "output", self.chunk); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for OutputChunkKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, workflow_id, _, chunk)) = + <(Cow, Uuid, Cow, usize)>::unpack(input, tuple_depth)?; + let v = OutputChunkKey { workflow_id, chunk }; + + Ok((input, v)) + } +} + +pub struct WakeSignalKey { + workflow_id: Uuid, + pub signal_name: String, +} + +impl WakeSignalKey { + pub fn new(workflow_id: Uuid, signal_name: String) -> Self { + WakeSignalKey { + workflow_id, + signal_name, + } + } + + pub fn subspace(workflow_id: Uuid) -> WakeSignalSubspaceKey { + WakeSignalSubspaceKey::new(workflow_id) + } +} + +impl FormalKey for WakeSignalKey { + type Value = (); + + fn deserialize(&self, _raw: &[u8]) -> Result { + Ok(()) + } + + fn serialize(&self, _value: Self::Value) -> Result> { + Ok(Vec::new()) + } +} + +impl TuplePack for WakeSignalKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "workflow", + "data", + self.workflow_id, + "wake_signal", + &self.signal_name, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for WakeSignalKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _, signal_name)) = + <(Cow, Cow, Uuid, Cow, String)>::unpack(input, tuple_depth)?; + let v = WakeSignalKey { + workflow_id, + signal_name, + }; + + Ok((input, v)) + } +} + +pub struct WakeSignalSubspaceKey { + workflow_id: Uuid, +} + +impl WakeSignalSubspaceKey { + pub fn new(workflow_id: Uuid) -> Self { + WakeSignalSubspaceKey { workflow_id } + } +} + +impl TuplePack for WakeSignalSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "wake_signal"); + t.pack(w, tuple_depth) + } +} + +pub struct WakeDeadlineKey { + workflow_id: Uuid, +} + +impl WakeDeadlineKey { + pub fn new(workflow_id: Uuid) -> Self { + WakeDeadlineKey { workflow_id } + } +} + +impl FormalKey for WakeDeadlineKey { + // Timestamp. + type Value = i64; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(i64::from_be_bytes(raw.try_into()?)) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.to_be_bytes().to_vec()) + } +} + +impl TuplePack for WakeDeadlineKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "wake_deadline"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for WakeDeadlineKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = WakeDeadlineKey { workflow_id }; + + Ok((input, v)) + } +} + +#[derive(Debug)] +pub struct NameKey { + workflow_id: Uuid, +} + +impl NameKey { + pub fn new(workflow_id: Uuid) -> Self { + NameKey { workflow_id } + } +} + +impl FormalKey for NameKey { + type Value = String; + + fn deserialize(&self, raw: &[u8]) -> Result { + String::from_utf8(raw.to_vec()).map_err(Into::into) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.into_bytes()) + } +} + +impl TuplePack for NameKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "name"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for NameKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = NameKey { workflow_id }; + + Ok((input, v)) + } +} + +#[derive(Debug)] +pub struct CreateTsKey { + workflow_id: Uuid, +} + +impl CreateTsKey { + pub fn new(workflow_id: Uuid) -> Self { + CreateTsKey { workflow_id } + } +} + +impl FormalKey for CreateTsKey { + // Timestamp. + type Value = i64; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(i64::from_be_bytes(raw.try_into()?)) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.to_be_bytes().to_vec()) + } +} + +impl TuplePack for CreateTsKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "create_ts"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for CreateTsKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = CreateTsKey { workflow_id }; + + Ok((input, v)) + } +} + +#[derive(Debug)] +pub struct RayIdKey { + workflow_id: Uuid, +} + +impl RayIdKey { + pub fn new(workflow_id: Uuid) -> Self { + RayIdKey { workflow_id } + } +} + +impl FormalKey for RayIdKey { + type Value = Uuid; + + fn deserialize(&self, raw: &[u8]) -> Result { + Ok(Uuid::from_slice(raw)?) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.as_bytes().to_vec()) + } +} + +impl TuplePack for RayIdKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "ray_id"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for RayIdKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = RayIdKey { workflow_id }; + + Ok((input, v)) + } +} + +#[derive(Debug)] +pub struct ErrorKey { + workflow_id: Uuid, +} + +impl ErrorKey { + pub fn new(workflow_id: Uuid) -> Self { + ErrorKey { workflow_id } + } +} + +impl FormalKey for ErrorKey { + type Value = String; + + fn deserialize(&self, raw: &[u8]) -> Result { + String::from_utf8(raw.to_vec()).map_err(Into::into) + } + + fn serialize(&self, value: Self::Value) -> Result> { + Ok(value.into_bytes()) + } +} + +impl TuplePack for ErrorKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ("workflow", "data", self.workflow_id, "error"); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for ErrorKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _)) = + <(Cow, Cow, Uuid, Cow)>::unpack(input, tuple_depth)?; + let v = ErrorKey { workflow_id }; + + Ok((input, v)) + } +} + +pub struct PendingSignalKey { + pub workflow_id: Uuid, + pub signal_name: String, + /// For ordering. + pub ts: i64, + pub signal_id: Uuid, +} + +impl PendingSignalKey { + pub fn new(workflow_id: Uuid, signal_name: String, signal_id: Uuid) -> Self { + PendingSignalKey { + workflow_id, + signal_name, + ts: rivet_util::timestamp::now(), + signal_id, + } + } + + pub fn subspace(workflow_id: Uuid, signal_name: String) -> PendingSignalSubspaceKey { + PendingSignalSubspaceKey::new(workflow_id, signal_name) + } +} + +impl FormalKey for PendingSignalKey { + type Value = (); + + fn deserialize(&self, _raw: &[u8]) -> Result { + Ok(()) + } + + fn serialize(&self, _value: Self::Value) -> Result> { + Ok(Vec::new()) + } +} + +impl TuplePack for PendingSignalKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "workflow", + "signal", + self.workflow_id, + "pending", + &self.signal_name, + self.ts, + self.signal_id, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for PendingSignalKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_id, _, signal_name, ts, signal_id)) = + <(Cow, Cow, Uuid, Cow, String, i64, Uuid)>::unpack(input, tuple_depth)?; + let v = PendingSignalKey { + workflow_id, + signal_name, + ts, + signal_id, + }; + + Ok((input, v)) + } +} + +pub struct PendingSignalSubspaceKey { + workflow_id: Uuid, + signal_name: String, +} + +impl PendingSignalSubspaceKey { + fn new(workflow_id: Uuid, signal_name: String) -> Self { + PendingSignalSubspaceKey { + workflow_id, + signal_name, + } + } +} + +impl TuplePack for PendingSignalSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "workflow", + "signal", + self.workflow_id, + "pending", + &self.signal_name, + ); + t.pack(w, tuple_depth) + } +} + +pub struct ByNameAndTagKey { + workflow_name: String, + k: String, + v: String, + pub workflow_id: Uuid, +} + +impl ByNameAndTagKey { + pub fn new(workflow_name: String, k: String, v: String, workflow_id: Uuid) -> Self { + ByNameAndTagKey { + workflow_name, + k, + v, + workflow_id, + } + } + + pub fn subspace(workflow_name: String, k: String, v: String) -> ByNameAndTagSubspaceKey { + ByNameAndTagSubspaceKey::new(workflow_name, k, v) + } + + pub fn null(workflow_name: String, workflow_id: Uuid) -> Self { + ByNameAndTagKey { + workflow_name, + k: String::new(), + v: String::new(), + workflow_id, + } + } + + pub fn null_subspace(workflow_name: String) -> ByNameAndTagSubspaceKey { + ByNameAndTagSubspaceKey::null(workflow_name) + } +} + +impl FormalKey for ByNameAndTagKey { + // Rest of the tags. + type Value = Vec<(String, String)>; + + fn deserialize(&self, raw: &[u8]) -> Result { + serde_json::from_slice(raw).map_err(Into::into) + } + + fn serialize(&self, value: Self::Value) -> Result> { + serde_json::to_vec(&value).map_err(Into::into) + } +} + +impl TuplePack for ByNameAndTagKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "workflow", + "by_name_and_tag", + &self.workflow_name, + &self.k, + &self.v, + self.workflow_id, + ); + t.pack(w, tuple_depth) + } +} + +impl<'de> TupleUnpack<'de> for ByNameAndTagKey { + fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> { + let (input, (_, _, workflow_name, k, v, workflow_id)) = + <(Cow, Cow, String, String, String, Uuid)>::unpack(input, tuple_depth)?; + let v = ByNameAndTagKey { + workflow_name, + k, + v, + workflow_id, + }; + + Ok((input, v)) + } +} + +pub struct ByNameAndTagSubspaceKey { + workflow_name: String, + k: String, + v: String, +} + +impl ByNameAndTagSubspaceKey { + pub fn new(workflow_name: String, k: String, v: String) -> Self { + ByNameAndTagSubspaceKey { + workflow_name, + k, + v, + } + } + + pub fn null(workflow_name: String) -> Self { + ByNameAndTagSubspaceKey { + workflow_name, + k: String::new(), + v: String::new(), + } + } +} + +impl TuplePack for ByNameAndTagSubspaceKey { + fn pack( + &self, + w: &mut W, + tuple_depth: TupleDepth, + ) -> std::io::Result { + let t = ( + "workflow", + "by_name_and_tag", + &self.workflow_name, + &self.k, + &self.v, + ); + t.pack(w, tuple_depth) + } +} diff --git a/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/mod.rs b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/mod.rs new file mode 100644 index 0000000000..942e4195a3 --- /dev/null +++ b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/mod.rs @@ -0,0 +1,2593 @@ +//! Implementation of a workflow database driver with SQLite and FoundationDB. +// TODO: Move code to smaller functions for readability + +use std::{ + collections::HashSet, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + time::Instant, +}; + +use foundationdb::{ + self as fdb, + options::{ConflictRangeType, StreamingMode}, + tuple::Subspace, +}; +use futures_util::{StreamExt, TryStreamExt}; +use indoc::indoc; +use rivet_pools::prelude::*; +use sqlite_util::SqliteConnectionExt; +use tokio::sync::Mutex; +use tracing::Instrument; +use uuid::Uuid; + +use super::{Database, PulledWorkflow, SignalData, WorkflowData}; +use crate::{ + error::{WorkflowError, WorkflowResult}, + history::{ + event::{EventId, EventType, SleepState}, + location::Location, + }, + metrics, worker, +}; +use keys::{FormalChunkedKey, FormalKey}; + +mod keys; +mod sqlite; + +// HACK: We alias global error here because its hardcoded into the sql macros +type GlobalError = WorkflowError; + +/// Base retry for query retry backoff. +const QUERY_RETRY_MS: usize = 500; +/// Maximum times a query ran by this database adapter is retried. +const MAX_QUERY_RETRIES: usize = 4; +/// For SQL macros. +const CONTEXT_NAME: &str = "chirp_workflow_fdb_sqlite_nats_engine"; +/// For NATS wake mechanism. +const WORKER_WAKE_SUBJECT: &str = "chirp.workflow.fdb_sqlite_nats.worker.wake"; +/// Makes the code blatantly obvious if its using a snapshot read. +const SNAPSHOT: bool = true; +const SERIALIZABLE: bool = false; + +/// Once tagged signals get removed, this and all of their code can be as well. This allows us to gradually +/// move off of tagged signals and makes it clear which code relates to tagged signals +const TAGGED_SIGNALS_ENABLED: bool = true; + +pub struct DatabaseFdbSqliteNats { + pools: rivet_pools::Pools, + sub: Mutex>, + subspace: Subspace, +} + +impl DatabaseFdbSqliteNats { + // For SQL macros + fn name(&self) -> &str { + CONTEXT_NAME + } + + /// Spawns a new thread and publishes a worker wake message to nats. + fn wake_worker(&self) { + let Ok(nats) = self.pools.nats() else { + tracing::debug!("failed to acquire nats pool"); + return; + }; + + let spawn_res = tokio::task::Builder::new().name("wake").spawn( + async move { + // Fail gracefully + if let Err(err) = nats.publish(WORKER_WAKE_SUBJECT, Vec::new().into()).await { + tracing::warn!(?err, "failed to publish wake message"); + } + } + .in_current_span(), + ); + if let Err(err) = spawn_res { + tracing::error!(?err, "failed to spawn wake task"); + } + } + + /// Executes SQL queries and explicitly handles retry errors. + async fn query<'a, F, Fut, T>(&self, mut cb: F) -> WorkflowResult + where + F: FnMut() -> Fut, + Fut: std::future::Future> + 'a, + T: 'a, + { + let mut backoff = rivet_util::Backoff::new(4, None, QUERY_RETRY_MS, 50); + let mut i = 0; + + loop { + match cb().await { + Err(WorkflowError::Sqlx(err)) => { + i += 1; + if i > MAX_QUERY_RETRIES { + return Err(WorkflowError::MaxSqlRetries(err)); + } + + use sqlx::Error::*; + match &err { + Database(_) | Io(_) | Tls(_) | Protocol(_) | PoolTimedOut | PoolClosed + | WorkerCrashed => { + tracing::warn!(?err, "query retry"); + backoff.tick().await; + } + // Throw error + _ => return Err(WorkflowError::Sqlx(err)), + } + } + x => return x, + } + } + } +} + +#[async_trait::async_trait] +impl Database for DatabaseFdbSqliteNats { + fn from_pools(pools: rivet_pools::Pools) -> Result, rivet_pools::Error> { + Ok(Arc::new(DatabaseFdbSqliteNats { + pools, + // Lazy load the nats sub + sub: Mutex::new(None), + subspace: Subspace::all().subspace(&("chirp_workflow", "fdb_sqlite_nats")), + })) + } + + async fn wake(&self) -> WorkflowResult<()> { + let mut sub = self.sub.try_lock().map_err(WorkflowError::WakeLock)?; + + // Initialize sub + if sub.is_none() { + *sub = Some( + self.pools + .nats()? + .subscribe(WORKER_WAKE_SUBJECT) + .await + .map_err(|x| WorkflowError::CreateSubscription(x.into()))?, + ); + } + + match sub.as_mut().expect("unreachable").next().await { + Some(_) => Ok(()), + None => Err(WorkflowError::SubscriptionUnsubscribed), + } + } + + async fn dispatch_workflow( + &self, + ray_id: Uuid, + workflow_id: Uuid, + workflow_name: &str, + tags: Option<&serde_json::Value>, + input: &serde_json::value::RawValue, + unique: bool, + ) -> WorkflowResult { + assert!( + !unique, + "unique dispatch unimplemented for fdb_sqlite_nats driver" + ); + + self.pools + .fdb()? + .run(|tx, _mc| async move { + // Write create ts + let create_ts_key = keys::workflow::CreateTsKey::new(workflow_id); + tx.set( + &self.subspace.pack(&create_ts_key), + &create_ts_key + .serialize(rivet_util::timestamp::now()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write name + let name_key = keys::workflow::NameKey::new(workflow_id); + tx.set( + &self.subspace.pack(&name_key), + &name_key + .serialize(workflow_name.to_string()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write ray id + let ray_id_key = keys::workflow::RayIdKey::new(workflow_id); + tx.set( + &self.subspace.pack(&ray_id_key), + &ray_id_key + .serialize(ray_id) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write tags + let tags = tags + .map(|x| { + x.as_object().ok_or_else(|| { + WorkflowError::InvalidTags("must be an object".to_string()) + }) + }) + .transpose() + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))? + .into_iter() + .flatten() + .map(|(k, v)| { + Ok(( + k.clone(), + cjson::to_string(v).map_err(WorkflowError::CjsonSerializeTags)?, + )) + }) + .collect::>>() + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + for (k, v) in &tags { + // Write tag key + let tag_key = keys::workflow::TagKey::new(workflow_id, k.clone(), v.clone()); + tx.set( + &self.subspace.pack(&tag_key), + &tag_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write "by name and first tag" secondary index + let by_name_and_tag_key = keys::workflow::ByNameAndTagKey::new( + workflow_name.to_string(), + k.clone(), + v.clone(), + workflow_id, + ); + let rest_of_tags = tags + .iter() + .filter(|(k2, _)| k2 != k) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + tx.set( + &self.subspace.pack(&by_name_and_tag_key), + &by_name_and_tag_key + .serialize(rest_of_tags) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + // Write null key for the "by name and first tag" secondary index (all workflows have this) + { + // Write secondary index by name and first tag + let by_name_and_tag_key = keys::workflow::ByNameAndTagKey::null( + workflow_name.to_string(), + workflow_id, + ); + tx.set( + &self.subspace.pack(&by_name_and_tag_key), + &by_name_and_tag_key + .serialize(tags) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + // Write input + let input_key = keys::workflow::InputKey::new(workflow_id); + + for (i, chunk) in input_key + .split_ref(input) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))? + .into_iter() + .enumerate() + { + let chunk_key = self.subspace.pack(&input_key.chunk(i)); + + tx.set(&chunk_key, &chunk); + } + + // Write immediate wake condition + let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( + workflow_name.to_string(), + workflow_id, + keys::wake::WakeCondition::Immediate, + ); + + tx.set( + &self.subspace.pack(&wake_condition_key), + &wake_condition_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + Ok(()) + }) + .await?; + + self.wake_worker(); + + Ok(workflow_id) + } + + async fn get_workflow(&self, workflow_id: Uuid) -> WorkflowResult> { + let data = self + .pools + .fdb()? + .run(|tx, _mc| async move { + let input_key = keys::workflow::InputKey::new(workflow_id); + let input_subspace = self.subspace.subspace(&input_key); + let output_key = keys::workflow::OutputKey::new(workflow_id); + let output_subspace = self.subspace.subspace(&output_key); + + // Read input and output + let (input_chunks, output_chunks) = tokio::try_join!( + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&input_subspace).into() + }, + SERIALIZABLE, + ) + .try_collect::>(), + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&output_subspace).into() + }, + SERIALIZABLE, + ) + .try_collect::>() + )?; + + if input_chunks.is_empty() { + Ok(None) + } else { + let input = input_key + .combine(input_chunks) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + let output = if output_chunks.is_empty() { + None + } else { + Some( + output_key + .combine(output_chunks) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ) + }; + + Ok(Some((input, output))) + } + }) + .await?; + + if let Some((input, output)) = data { + Ok(Some(WorkflowData { + workflow_id, + input, + output, + })) + } else { + Ok(None) + } + } + + async fn find_workflow( + &self, + workflow_name: &str, + tags: &serde_json::Value, + ) -> WorkflowResult> { + // Convert to flat vec of strings + let mut tag_iter = tags + .as_object() + .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string()))? + .iter() + .map(|(k, v)| { + Ok(( + k.clone(), + serde_json::to_string(&v).map_err(WorkflowError::DeserializeTags)?, + )) + }); + let first_tag = tag_iter.next().transpose()?; + let rest_of_tags = tag_iter.collect::>>()?; + + let workflow_id = self + .pools + .fdb()? + .run(|tx, _mc| { + let first_tag = first_tag.clone(); + let rest_of_tags = rest_of_tags.clone(); + async move { + let workflow_by_name_and_tag_subspace = + if let Some((first_tag_key, first_tag_value)) = first_tag { + self.subspace + .subspace(&keys::workflow::ByNameAndTagKey::subspace( + workflow_name.to_string(), + first_tag_key, + first_tag_value, + )) + } else { + // No tags provided, use null subspace. Every workflow has a null key for its tags + // under the `ByNameAndTagKey` + self.subspace + .subspace(&keys::workflow::ByNameAndTagKey::null_subspace( + workflow_name.to_string(), + )) + }; + + let mut stream = tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::Iterator, + ..(&workflow_by_name_and_tag_subspace).into() + }, + // NOTE: This is not SERIALIZABLE because we don't want to conflict with + // `update_workflow_tags` + SNAPSHOT, + ); + + loop { + let Some(entry) = stream.try_next().await? else { + return Ok(None); + }; + + // Unpack key + let workflow_by_name_and_tag_key = self + .subspace + .unpack::(&entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + // Deserialize value + let wf_rest_of_tags = workflow_by_name_and_tag_key + .deserialize(entry.value()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + // Compute intersection between wf tags and signal tags + let tags_match = rest_of_tags.iter().all(|(k, v)| { + wf_rest_of_tags + .iter() + .any(|(wf_k, wf_v)| k == wf_k && v == wf_v) + }); + + // Return first signal that matches the tags + if tags_match { + break Ok(Some(workflow_by_name_and_tag_key.workflow_id)); + } + } + } + }) + .await?; + + Ok(workflow_id) + } + + async fn pull_workflows( + &self, + worker_instance_id: Uuid, + filter: &[&str], + ) -> WorkflowResult> { + let start_instant = Instant::now(); + let owned_filter = filter + .into_iter() + .map(|x| x.to_string()) + .collect::>(); + + let partial_workflows = self + .pools + .fdb()? + .run(|tx, _mc| { + let owned_filter = owned_filter.clone(); + + async move { + // All wake conditions with a timestamp after this timestamp will be pulled + let pull_before = rivet_util::timestamp::now() + + i64::try_from(worker::TICK_INTERVAL.as_millis()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + // Pull all available wake conditions from all registered wf names + let entries = futures_util::stream::iter(owned_filter) + .map(|wf_name| { + let wake_subspace_start = self + .subspace + .subspace( + &keys::wake::WorkflowWakeConditionKey::subspace_without_ts( + wf_name.clone(), + ), + ) + .bytes() + .iter() + .map(|x| *x) + // https://github.com/apple/foundationdb/blob/main/design/tuple.md + .chain(std::iter::once(0x00)) + .collect::>(); + let wake_subspace_end = self + .subspace + .subspace(&keys::wake::WorkflowWakeConditionKey::subspace( + wf_name, + pull_before, + )) + .bytes() + .to_vec(); + + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(wake_subspace_start, wake_subspace_end).into() + }, + // Must be a snapshot to not conflict with any new wake conditions being + // inserted + SNAPSHOT, + ) + }) + .flatten() + .map(|res| match res { + Ok(entry) => Ok(( + entry.key().to_vec(), + self.subspace + .unpack::(entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + )), + Err(err) => Err(Into::::into(err)), + }) + .try_collect::>() + .await?; + + // Check leases + let wf_ids = entries + .iter() + .map(|(_, key)| key.workflow_id) + .collect::>(); + let leased_wf_ids = futures_util::stream::iter(wf_ids) + .map(|wf_id| { + let tx = tx.clone(); + async move { + let lease_key = keys::workflow::LeaseKey::new(wf_id); + let lease_key_buf = self.subspace.pack(&lease_key); + + // Check lease + if tx.get(&lease_key_buf, SERIALIZABLE).await?.is_some() { + Result::<_, fdb::FdbBindingError>::Ok(None) + } else { + tx.set( + &lease_key_buf, + &lease_key.serialize(worker_instance_id).map_err(|x| { + fdb::FdbBindingError::CustomError(x.into()) + })?, + ); + Ok(Some(wf_id)) + } + } + }) + // TODO: How to get rid of this buffer? + .buffer_unordered(1024) + .try_filter_map(|x| async move { Ok(x) }) + .try_collect::>() + .await?; + + // TODO: Split this txn into two after checking leases here? + + for (raw_key, key) in &entries { + // Check if leased + if !leased_wf_ids.iter().any(|wf_id| wf_id == &key.workflow_id) { + continue; + } + + // Clear fetched wake conditions + tx.clear(raw_key); + + if TAGGED_SIGNALS_ENABLED { + // Clear secondary indexes of tagged signals + if let keys::wake::WakeCondition::TaggedSignal { .. } = key.condition { + let wake_signals_subspace = self.subspace.subspace( + &keys::workflow::WakeSignalKey::subspace(key.workflow_id), + ); + + // Read current signal names + let mut stream = tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&wake_signals_subspace).into() + }, + SERIALIZABLE, + ); + + // Clear each secondary index + while let Some(entry) = stream.try_next().await? { + let wake_signal_key = self + .subspace + .unpack::(entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + let tagged_signal_wake_key = + self.subspace.pack(&keys::wake::TaggedSignalWakeKey::new( + wake_signal_key.signal_name, + key.workflow_id, + )); + + tx.clear(&tagged_signal_wake_key); + } + } + } + } + + // Read required data for each leased wf + futures_util::stream::iter(entries) + .map(|(_, key)| { + let tx = tx.clone(); + async move { + let create_ts_key = + keys::workflow::CreateTsKey::new(key.workflow_id); + let ray_id_key = keys::workflow::RayIdKey::new(key.workflow_id); + let input_key = keys::workflow::InputKey::new(key.workflow_id); + let input_subspace = self.subspace.subspace(&input_key); + + let (create_ts_entry, ray_id_entry, input_chunks) = tokio::try_join!( + tx.get(&self.subspace.pack(&create_ts_key), SERIALIZABLE), + tx.get(&self.subspace.pack(&ray_id_key), SERIALIZABLE), + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&input_subspace).into() + }, + SERIALIZABLE, + ) + .try_collect::>(), + )?; + + let create_ts = create_ts_key + .deserialize(&create_ts_entry.ok_or( + fdb::FdbBindingError::CustomError( + format!("key should exist: {create_ts_key:?}").into(), + ), + )?) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + let ray_id = ray_id_key + .deserialize(&ray_id_entry.ok_or( + fdb::FdbBindingError::CustomError( + format!("key should exist: {ray_id_key:?}").into(), + ), + )?) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + let input = input_key + .combine(input_chunks) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + Ok(PartialWorkflow { + workflow_id: key.workflow_id, + workflow_name: key.workflow_name, + create_ts, + ray_id, + input, + }) + } + }) + // TODO: How to get rid of this buffer? + .buffer_unordered(512) + .try_collect::>() + .await + } + }) + .await?; + + let dt = start_instant.elapsed().as_secs_f64(); + metrics::PULL_WORKFLOWS_DURATION + .with_label_values(&[&worker_instance_id.to_string()]) + .set(dt); + + if partial_workflows.is_empty() { + return Ok(Vec::new()); + } + + let start_instant2 = Instant::now(); + + // Set up sqlite + let pulled_workflows = futures_util::stream::iter(partial_workflows) + .map(|partial| async move { + let pool = self.pools.sqlite(partial.workflow_id).await?; + sqlite::init(partial.workflow_id, &pool).await?; + + // Fetch all events + let events = sql_fetch_all!( + [self, sqlite::AmalgamEventRow, &pool] + " + -- Activity events + SELECT + json(location) AS location, + version, + 0 AS event_type, -- EventType + activity_name AS name, + NULL AS auxiliary_id, + input_hash AS hash, + NULL AS input, + json(output) AS output, + create_ts AS create_ts, + ( + SELECT COUNT(*) + FROM workflow_activity_errors AS err + WHERE ev.location = err.location + ) AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_activity_events AS ev + WHERE NOT forgotten + GROUP BY ev.location + UNION ALL + -- Signal listen events + SELECT + json(location) AS location, + version, + 1 AS event_type, -- EventType + signal_name AS name, + NULL AS auxiliary_id, + NULL AS hash, + NULL AS input, + json(body) AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_signal_events + WHERE NOT forgotten + UNION ALL + -- Signal send events + SELECT + json(location) AS location, + version, + 2 AS event_type, -- EventType + signal_name AS name, + signal_id AS auxiliary_id, + NULL AS hash, + NULL AS input, + NULL AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_signal_send_events + WHERE NOT forgotten + UNION ALL + -- Message send events + SELECT + json(location) AS location, + version, + 3 AS event_type, -- EventType + message_name AS name, + NULL AS auxiliary_id, + NULL AS hash, + NULL AS input, + NULL AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_message_send_events + WHERE NOT forgotten + UNION ALL + -- Sub workflow events + SELECT + json(location) AS location, + version, + 4 AS event_type, -- crdb_nats::types::EventType + sub_workflow_name AS name, + sub_workflow_id AS auxiliary_id, + NULL AS hash, + NULL AS input, + NULL AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_sub_workflow_events + WHERE NOT forgotten + UNION ALL + -- Loop events + SELECT + json(location) AS location, + version, + 5 AS event_type, -- crdb_nats::types::EventType + NULL AS name, + NULL AS auxiliary_id, + NULL AS hash, + json(state) AS input, + output, + NULL AS create_ts, + NULL AS error_count, + iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_loop_events + WHERE NOT forgotten + UNION ALL + -- Sleep events + SELECT + json(location) AS location, + version, + 6 AS event_type, -- crdb_nats::types::EventType + NULL AS name, + NULL AS auxiliary_id, + NULL AS hash, + NULL AS input, + NULL AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + deadline_ts, + state, + NULL AS inner_event_type + FROM workflow_sleep_events + WHERE NOT forgotten + UNION ALL + -- Branch events + SELECT + json(location) AS location, + version, + 7 AS event_type, -- crdb_nats::types::EventType + NULL AS name, + NULL AS auxiliary_id, + NULL AS hash, + NULL AS input, + NULL AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_branch_events + WHERE NOT forgotten + UNION ALL + -- Removed events + SELECT + json(location) AS location, + 1 AS version, -- Default + 8 AS event_type, -- crdb_nats::types::EventType + event_name AS name, + NULL AS auxiliary_id, + NULL AS hash, + NULL AS input, + NULL AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + event_type AS inner_event_type + FROM workflow_removed_events + WHERE NOT forgotten + UNION ALL + -- Version check events + SELECT + json(location) AS location, + version, + 9 AS event_type, -- crdb_nats::types::EventType + NULL AS name, + NULL AS auxiliary_id, + NULL AS hash, + NULL AS input, + NULL AS output, + NULL AS create_ts, + NULL AS error_count, + NULL AS iteration, + NULL AS deadline_ts, + NULL AS state, + NULL AS inner_event_type + FROM workflow_version_check_events + WHERE NOT forgotten + ", + ) + .await?; + + WorkflowResult::Ok(PulledWorkflow { + workflow_id: partial.workflow_id, + workflow_name: partial.workflow_name, + create_ts: partial.create_ts, + ray_id: partial.ray_id, + input: partial.input, + events: sqlite::build_history(events)?, + }) + }) + .buffer_unordered(512) + .try_collect() + .await?; + + let dt = start_instant2.elapsed().as_secs_f64(); + metrics::PULL_WORKFLOWS_HISTORY_DURATION + .with_label_values(&[&worker_instance_id.to_string()]) + .set(dt); + let dt = start_instant.elapsed().as_secs_f64(); + metrics::PULL_WORKFLOWS_FULL_DURATION + .with_label_values(&[&worker_instance_id.to_string()]) + .set(dt); + + Ok(pulled_workflows) + } + + async fn complete_workflow( + &self, + workflow_id: Uuid, + workflow_name: &str, + output: &serde_json::value::RawValue, + ) -> WorkflowResult<()> { + self.pools + .fdb()? + .run(|tx, _mc| async move { + // Check for other workflows waiting on this one, wake all + let sub_workflow_wake_subspace = self + .subspace + .subspace(&keys::wake::SubWorkflowWakeKey::subspace(workflow_id)); + + let mut stream = tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&sub_workflow_wake_subspace).into() + }, + // Must be serializable to conflict with `get_sub_workflow` + SERIALIZABLE, + ); + + while let Some(entry) = stream.try_next().await.unwrap() { + let sub_workflow_wake_key = self + .subspace + .unpack::(&entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + let workflow_name = sub_workflow_wake_key + .deserialize(entry.value()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( + workflow_name, + sub_workflow_wake_key.workflow_id, + keys::wake::WakeCondition::SubWorkflow { + sub_workflow_id: workflow_id, + }, + ); + + // Add wake condition for workflow + tx.set( + &self.subspace.pack(&wake_condition_key), + &wake_condition_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Clear secondary index + tx.clear(entry.key()); + } + + // Get and clear the pending deadline wake condition, if any. We have to clear here because + // the `pull_workflows` function might not have pulled the condition if its in the future. + let wake_deadline_key = keys::workflow::WakeDeadlineKey::new(workflow_id); + if let Some(raw) = tx + .get(&self.subspace.pack(&wake_deadline_key), SERIALIZABLE) + .await? + { + let deadline_ts = wake_deadline_key + .deserialize(&raw) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( + workflow_name.to_string(), + workflow_id, + keys::wake::WakeCondition::Deadline { deadline_ts }, + ); + + tx.clear(&self.subspace.pack(&wake_condition_key)); + } + + // Write output + let output_key = keys::workflow::OutputKey::new(workflow_id); + + for (i, chunk) in output_key + .split_ref(output) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))? + .into_iter() + .enumerate() + { + let chunk_key = self.subspace.pack(&output_key.chunk(i)); + + tx.set(&chunk_key, &chunk); + } + + Ok(()) + }) + .await?; + + self.wake_worker(); + + Ok(()) + } + + async fn commit_workflow( + &self, + workflow_id: Uuid, + workflow_name: &str, + _immediate: bool, + wake_deadline_ts: Option, + wake_signals: &[&str], + wake_sub_workflow_id: Option, + error: &str, + ) -> WorkflowResult<()> { + self.pools + .fdb()? + .run(|tx, _mc| async move { + let wf_tags_subspace = self + .subspace + .subspace(&keys::workflow::TagKey::subspace(workflow_id)); + let wake_deadline_key = keys::workflow::WakeDeadlineKey::new(workflow_id); + + let (wf_tags, wake_deadline) = tokio::try_join!( + // Collect all wf tags + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&wf_tags_subspace).into() + }, + SERIALIZABLE, + ) + .map(|res| match res { + Ok(entry) => { + let key = self + .subspace + .unpack::(entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + Ok((key.k, key.v)) + } + Err(err) => Err(Into::::into(err)), + }) + .try_collect::>(), + // Get previous deadline wake condition + async { + tx.get(&self.subspace.pack(&wake_deadline_key), SERIALIZABLE) + .await + .map_err(Into::into) + }, + )?; + + // Clear all previous wake signals in list + let wake_signals_subspace = self + .subspace + .subspace(&keys::workflow::WakeSignalKey::subspace(workflow_id)); + tx.clear_subspace_range(&wake_signals_subspace); + + // Write signals wake index (only for tagged signals) + for signal_name in wake_signals { + if TAGGED_SIGNALS_ENABLED { + let signal_wake_key = keys::wake::TaggedSignalWakeKey::new( + signal_name.to_string(), + workflow_id, + ); + + tx.set( + &self.subspace.pack(&signal_wake_key), + &signal_wake_key + .serialize(keys::wake::TaggedSignalWakeData { + workflow_name: workflow_name.to_string(), + tags: wf_tags.clone(), + }) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + // Write to wake signals list + let wake_signal_key = + keys::workflow::WakeSignalKey::new(workflow_id, signal_name.to_string()); + tx.set( + &self.subspace.pack(&wake_signal_key), + &wake_signal_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + // Clear previous deadline wake condition. It has to be cleared here because the + // pull_workflows function might not have pulled and cleared the condition if its in the + // future. + if let Some(raw) = wake_deadline { + let deadline_ts = wake_deadline_key + .deserialize(&raw) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( + workflow_name.to_string(), + workflow_id, + keys::wake::WakeCondition::Deadline { deadline_ts }, + ); + + tx.clear(&self.subspace.pack(&wake_condition_key)); + } + + // Write deadline wake index + if let Some(deadline_ts) = wake_deadline_ts { + let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( + workflow_name.to_string(), + workflow_id, + keys::wake::WakeCondition::Deadline { deadline_ts }, + ); + + // Add wake condition for workflow + tx.set( + &self.subspace.pack(&wake_condition_key), + &wake_condition_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write to wake deadline + tx.set( + &self.subspace.pack(&wake_deadline_key), + &wake_deadline_key + .serialize(deadline_ts) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + // Write sub workflow wake index + if let Some(sub_workflow_id) = wake_sub_workflow_id { + let sub_workflow_wake_key = + keys::wake::SubWorkflowWakeKey::new(sub_workflow_id, workflow_id); + + tx.set( + &self.subspace.pack(&sub_workflow_wake_key), + &sub_workflow_wake_key + .serialize(workflow_name.to_string()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + // Write error + let error_key = keys::workflow::ErrorKey::new(workflow_id); + tx.set( + &self.subspace.pack(&error_key), + &error_key + .serialize(error.to_string()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Clear lease + let lease_key = self + .subspace + .pack(&keys::workflow::LeaseKey::new(workflow_id)); + tx.clear(&lease_key); + + Ok(()) + }) + .await?; + + Ok(()) + } + + async fn pull_next_signal( + &self, + workflow_id: Uuid, + filter: &[&str], + location: &Location, + version: usize, + loop_location: Option<&Location>, + ) -> WorkflowResult> { + let pool = self.pools.sqlite(workflow_id).await?; + + let owned_filter = filter + .into_iter() + .map(|x| x.to_string()) + .collect::>(); + let is_retrying = Arc::new(AtomicBool::new(false)); + + // Fetch signal from FDB + let signal = self + .pools + .fdb()? + .run(|tx, _mc| { + let pool = pool.clone(); + let owned_filter = owned_filter.clone(); + let is_retrying = is_retrying.clone(); + + async move { + let signal = + { + // Create a stream for each signal name subspace + let streams = owned_filter + .iter() + .map(|signal_name| { + let pending_signal_subspace = self.subspace.subspace( + &keys::workflow::PendingSignalKey::subspace( + workflow_id, + signal_name.to_string(), + ), + ); + + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + limit: Some(1), + ..(&pending_signal_subspace).into() + }, + // NOTE: This does not have to be SERIALIZABLE because the conflict occurs + // with acking which is a separate row. See below + SNAPSHOT, + ) + }) + .collect::>(); + + // Fetch the next entry from all streams at the same time + let mut results = futures_util::future::try_join_all( + streams.into_iter().map(|mut stream| async move { + if let Some(entry) = stream.try_next().await? { + Result::<_, fdb::FdbBindingError>::Ok(Some(( + entry.key().to_vec(), + self.subspace + .unpack::( + &entry.key(), + ) + .map_err(|x| { + fdb::FdbBindingError::CustomError(x.into()) + })?, + ))) + } else { + Ok(None) + } + }), + ) + .await?; + + // Sort by ts + results.sort_by_key(|res| res.as_ref().map(|(_, key)| key.ts)); + + results.into_iter().flatten().next().map( + |(raw_key, pending_signal_key)| { + ( + raw_key, + pending_signal_key.signal_name, + pending_signal_key.ts, + pending_signal_key.signal_id, + ) + }, + ) + }; + + let tagged_signal = if TAGGED_SIGNALS_ENABLED { + // Collect all wf tags + let wf_tags_subspace = self + .subspace + .subspace(&keys::workflow::TagKey::subspace(workflow_id)); + + let wf_tags = tx + .get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&wf_tags_subspace).into() + }, + SERIALIZABLE, + ) + .map(|res| match res { + Ok(entry) => Ok(self + .subspace + .unpack::(entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?), + Err(err) => Err(Into::::into(err)), + }) + .try_collect::>() + .await?; + + // Create a stream for each signal name subspace + let mut streams = owned_filter + .iter() + .map(|signal_name| { + let pending_signal_subspace = self.subspace.subspace( + &keys::signal::TaggedPendingKey::subspace( + signal_name.to_string(), + ), + ); + + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::Iterator, + ..(&pending_signal_subspace).into() + }, + // NOTE: This does not have to be SERIALIZABLE because the conflict occurs + // with acking which is a separate row. See below + SNAPSHOT, + ) + }) + .collect::>(); + + 'streams: loop { + let taken_streams = std::mem::take(&mut streams); + + // Fetch the next entry from all streams at the same time + let (mut results, new_streams): (Vec<_>, Vec<_>) = + futures_util::future::try_join_all(taken_streams.into_iter().map( + |mut stream| async move { + if let Some(entry) = stream.try_next().await? { + // Unpack key + let pending_signal_key = self + .subspace + .unpack::( + &entry.key(), + ) + .map_err(|x| { + fdb::FdbBindingError::CustomError(x.into()) + })?; + + // Deserialize value + let signal_tags = pending_signal_key + .deserialize(entry.value()) + .map_err(|x| { + fdb::FdbBindingError::CustomError(x.into()) + })?; + + Result::<_, fdb::FdbBindingError>::Ok(Some(( + ( + entry.key().to_vec(), + pending_signal_key, + signal_tags, + ), + stream, + ))) + } else { + Ok(None) + } + }, + )) + .await? + .into_iter() + .flatten() + .unzip(); + + if results.is_empty() { + break None; + } + + // Sort by ts + results.sort_by_key(|(_, key, _)| key.ts); + + for (raw_key, pending_signal_key, signal_tags) in results { + // Compute intersection between wf tags and signal tags + let tags_match = signal_tags.iter().all(|(k, v)| { + wf_tags.iter().any(|key| k == &key.k && v == &key.v) + }); + + // Return first signal that matches the tags + if tags_match { + break 'streams Some(( + raw_key, + pending_signal_key.signal_name, + pending_signal_key.ts, + pending_signal_key.signal_id, + )); + } + } + + // Write non-empty streams for the next loop iteration + streams = new_streams; + } + } else { + None + }; + + // Choose between the two. This functionality is not combined with the above two blocks + // because tagged signals may be removed in the future. + let earliest_signal = match (signal, tagged_signal) { + (Some(signal), Some(tagged_signal)) => { + if signal.2 < tagged_signal.2 { + Some(signal) + } else { + Some(tagged_signal) + } + } + (signal, None) => signal, + (None, tagged_signal) => tagged_signal, + }; + + // Signal found + if let Some((raw_key, signal_name, ts, signal_id)) = earliest_signal { + let ack_key = keys::signal::AckKey::new(signal_id); + + // Ack signal + tx.add_conflict_range(&raw_key, &raw_key, ConflictRangeType::Read)?; + tx.set( + &self.subspace.pack(&ack_key), + &ack_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // TODO: Split txn into two after acking here? + + // Clear pending signal key + tx.clear(&raw_key); + + // Read signal body + let body_key = keys::signal::BodyKey::new(signal_id); + let body_subspace = self.subspace.subspace(&body_key); + + let chunks = tx + .get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&body_subspace).into() + }, + SERIALIZABLE, + ) + .try_collect::>() + .await?; + + let body = body_key + .combine(chunks) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + // In the event of an FDB txn retry, we have to delete the previously inserted row + if is_retrying.load(Ordering::Relaxed) { + self.query(|| async { + sql_execute!( + [self, &pool] + " + DELETE FROM workflow_signal_events + WHERE location = ? + ", + ) + .await + }) + .await + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + } + + // Insert history event + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_signal_events ( + location, + version, + signal_id, + signal_name, + body, + ack_ts, + loop_location + ) + VALUES (jsonb(?), ?, ?, ?, jsonb(?), ?, jsonb(?)) + ", + location, + version as i64, + signal_id, + &signal_name, + sqlx::types::Json(&body), + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + is_retrying.store(true, Ordering::Relaxed); + + Ok(Some(SignalData { + signal_id, + signal_name, + create_ts: ts, + body, + })) + } else { + Ok(None) + } + } + }) + .await?; + + Ok(signal) + } + + async fn get_sub_workflow( + &self, + workflow_id: Uuid, + workflow_name: &str, + sub_workflow_id: Uuid, + ) -> WorkflowResult> { + let data = self + .pools + .fdb()? + .run(|tx, _mc| async move { + let input_key = keys::workflow::InputKey::new(sub_workflow_id); + let input_subspace = self.subspace.subspace(&input_key); + let output_key = keys::workflow::OutputKey::new(sub_workflow_id); + let output_subspace = self.subspace.subspace(&output_key); + + // Read input and output + let (input_chunks, output_chunks) = tokio::try_join!( + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&input_subspace).into() + }, + SERIALIZABLE, + ) + .try_collect::>(), + tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&output_subspace).into() + }, + SERIALIZABLE, + ) + .try_collect::>() + )?; + + if input_chunks.is_empty() { + Ok(None) + } else { + let input = input_key + .combine(input_chunks) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + let output = if output_chunks.is_empty() { + // Write sub workflow wake index + let sub_workflow_wake_key = + keys::wake::SubWorkflowWakeKey::new(sub_workflow_id, workflow_id); + + tx.set( + &self.subspace.pack(&sub_workflow_wake_key), + &sub_workflow_wake_key + .serialize(workflow_name.to_string()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + None + } else { + Some( + output_key + .combine(output_chunks) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ) + }; + + Ok(Some((input, output))) + } + }) + .await?; + + if let Some((input, output)) = data { + Ok(Some(WorkflowData { + workflow_id: sub_workflow_id, + input, + output, + })) + } else { + Ok(None) + } + } + + async fn publish_signal( + &self, + ray_id: Uuid, + workflow_id: Uuid, + signal_id: Uuid, + signal_name: &str, + body: &serde_json::value::RawValue, + ) -> WorkflowResult<()> { + self.pools + .fdb()? + .run(|tx, _mc| async move { + let workflow_name_key = keys::workflow::NameKey::new(workflow_id); + + // NOTE: This does not have to be serializable because create ts doesn't change, and + // should not conflict with dispatching a new workflow + // Check if the workflow exists + let Some(workflow_name_raw) = tx + .get(&self.subspace.pack(&workflow_name_key), SNAPSHOT) + .await? + else { + return Err(fdb::FdbBindingError::CustomError( + WorkflowError::WorkflowNotFound.into(), + )); + }; + + let workflow_name = workflow_name_key + .deserialize(&workflow_name_raw) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + let signal_body_key = keys::signal::BodyKey::new(signal_id); + + // Write signal body + for (i, chunk) in signal_body_key + .split_ref(body) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))? + .into_iter() + .enumerate() + { + let chunk_key = self.subspace.pack(&signal_body_key.chunk(i)); + + tx.set(&chunk_key, &chunk); + } + + // Write pending key + let pending_signal_key = keys::workflow::PendingSignalKey::new( + workflow_id, + signal_name.to_string(), + signal_id, + ); + + tx.set( + &self.subspace.pack(&pending_signal_key), + &pending_signal_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write create ts + let create_ts_key = keys::signal::CreateTsKey::new(signal_id); + tx.set( + &self.subspace.pack(&create_ts_key), + &create_ts_key + .serialize(pending_signal_key.ts) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write ray id ts + let ray_id_key = keys::signal::RayIdKey::new(signal_id); + tx.set( + &self.subspace.pack(&ray_id_key), + &ray_id_key + .serialize(ray_id) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + let wake_signal_key = + keys::workflow::WakeSignalKey::new(workflow_id, signal_name.to_string()); + + // If the workflow is currently listening for this signal, wake it + if tx + .get(&self.subspace.pack(&wake_signal_key), SERIALIZABLE) + .await? + .is_some() + { + let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( + workflow_name, + workflow_id, + keys::wake::WakeCondition::Signal { signal_id }, + ); + + // Add wake condition for workflow + tx.set( + &self.subspace.pack(&wake_condition_key), + &wake_condition_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + Ok(()) + }) + .await?; + + self.wake_worker(); + + Ok(()) + } + + async fn publish_tagged_signal( + &self, + ray_id: Uuid, + tags: &serde_json::Value, + signal_id: Uuid, + signal_name: &str, + body: &serde_json::value::RawValue, + ) -> WorkflowResult<()> { + if TAGGED_SIGNALS_ENABLED { + // Convert to flat vec of strings + let signal_tags = tags + .as_object() + .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string()))? + .iter() + .map(|(k, v)| { + Ok(( + k.clone(), + serde_json::to_string(&v).map_err(WorkflowError::DeserializeTags)?, + )) + }) + .collect::>>()?; + + self.pools + .fdb()? + .run(|tx, _mc| { + let signal_tags = signal_tags.clone(); + async move { + let signal_body_key = keys::signal::BodyKey::new(signal_id); + + // Write signal body + for (i, chunk) in signal_body_key + .split_ref(body) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))? + .into_iter() + .enumerate() + { + let chunk_key = self.subspace.pack(&signal_body_key.chunk(i)); + + tx.set(&chunk_key, &chunk); + } + + // Write pending key + let pending_signal_key = + keys::signal::TaggedPendingKey::new(signal_name.to_string(), signal_id); + + tx.set( + &self.subspace.pack(&pending_signal_key), + &pending_signal_key + .serialize(signal_tags.clone()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write create ts + let create_ts_key = keys::signal::CreateTsKey::new(signal_id); + tx.set( + &self.subspace.pack(&create_ts_key), + &create_ts_key + .serialize(pending_signal_key.ts) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write ray id ts + let ray_id_key = keys::signal::RayIdKey::new(signal_id); + tx.set( + &self.subspace.pack(&ray_id_key), + &ray_id_key + .serialize(ray_id) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write tags + let tags = tags + .as_object() + .ok_or_else(|| { + WorkflowError::InvalidTags("must be an object".to_string()) + }) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + for (k, v) in tags { + let tag_key = keys::signal::TagKey::new( + signal_id, + k.clone(), + cjson::to_string(v).map_err(|x| { + fdb::FdbBindingError::CustomError( + WorkflowError::CjsonSerializeTags(x).into(), + ) + })?, + ); + tx.set( + &self.subspace.pack(&tag_key), + &tag_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + // Read workflows waiting for the same signal name and tags + let signal_wake_subspace = + self.subspace + .subspace(&keys::wake::TaggedSignalWakeKey::subspace( + signal_name.to_string(), + )); + + let mut stream = tx.get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::Iterator, + ..(&signal_wake_subspace).into() + }, + // Must be a snapshot to not conflict with the pull workflows query + SNAPSHOT, + ); + + // TODO: This is almost a full scan + while let Some(entry) = stream.try_next().await.unwrap() { + let signal_wake_key = self + .subspace + .unpack::(&entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + let wake_data = signal_wake_key + .deserialize(entry.value()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + // Compute intersection between wf tags and signal tags + let tags_match = signal_tags.iter().all(|(k, v)| { + wake_data.tags.iter().any(|(k2, v2)| k == k2 && v == v2) + }); + + if tags_match { + let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new( + wake_data.workflow_name, + signal_wake_key.workflow_id, + keys::wake::WakeCondition::TaggedSignal { signal_id }, + ); + + // Add wake condition for workflow + tx.set( + &self.subspace.pack(&wake_condition_key), + &wake_condition_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + break; + } + } + + Ok(()) + } + }) + .await?; + + self.wake_worker(); + + Ok(()) + } else { + Err(WorkflowError::TaggedSignalsDisabled) + } + } + + async fn publish_signal_from_workflow( + &self, + from_workflow_id: Uuid, + location: &Location, + version: usize, + ray_id: Uuid, + to_workflow_id: Uuid, + signal_id: Uuid, + signal_name: &str, + body: &serde_json::value::RawValue, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(from_workflow_id).await?; + + // Insert history event + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_signal_send_events ( + location, version, signal_id, signal_name, body, create_ts, loop_location + ) + VALUES (jsonb(?), ?, ?, ?, jsonb(?), ?, jsonb(?)) + ", + location, + version as i64, + signal_id, + signal_name, + sqlx::types::Json(body), + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await?; + + if let Err(err) = self + .publish_signal(ray_id, to_workflow_id, signal_id, signal_name, body) + .await + { + // Undo history if FDB failed + self.query(|| async { + sql_execute!( + [self, &pool] + " + DELETE FROM workflow_signal_send_events + WHERE location = ? + ", + location, + ) + .await + }) + .await?; + + Err(err) + } else { + Ok(()) + } + } + + async fn publish_tagged_signal_from_workflow( + &self, + from_workflow_id: Uuid, + location: &Location, + version: usize, + ray_id: Uuid, + tags: &serde_json::Value, + signal_id: Uuid, + signal_name: &str, + body: &serde_json::value::RawValue, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + if TAGGED_SIGNALS_ENABLED { + let pool = self.pools.sqlite(from_workflow_id).await?; + + // Insert history event + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_signal_send_events ( + location, version, signal_id, signal_name, body, create_ts, loop_location + ) + VALUES (jsonb(?), ?, ?, ?, jsonb(?), ?, jsonb(?)) + ", + location, + version as i64, + signal_id, + signal_name, + sqlx::types::Json(body), + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await?; + + if let Err(err) = self + .publish_tagged_signal(ray_id, tags, signal_id, signal_name, body) + .await + { + // Undo history if FDB failed + self.query(|| async { + sql_execute!( + [self, &pool] + " + DELETE FROM workflow_signal_send_events + WHERE location = ? + ", + location, + ) + .await + }) + .await?; + + Err(err) + } else { + Ok(()) + } + } else { + Err(WorkflowError::TaggedSignalsDisabled) + } + } + + async fn dispatch_sub_workflow( + &self, + ray_id: Uuid, + workflow_id: Uuid, + location: &Location, + version: usize, + sub_workflow_id: Uuid, + sub_workflow_name: &str, + tags: Option<&serde_json::Value>, + input: &serde_json::value::RawValue, + loop_location: Option<&Location>, + unique: bool, + ) -> WorkflowResult { + let pool = self.pools.sqlite(workflow_id).await?; + + // Insert history event + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_sub_workflow_events ( + location, version, sub_workflow_id, sub_workflow_name, create_ts, loop_location + ) + VALUES (jsonb(?), ?, ?, ?, ?, jsonb(?)) + ", + location, + version as i64, + sub_workflow_id, + sub_workflow_name, + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await?; + + match self + .dispatch_workflow( + ray_id, + sub_workflow_id, + sub_workflow_name, + tags, + input, + unique, + ) + .await + { + Ok(workflow_id) => Ok(workflow_id), + Err(err) => { + // Undo history if FDB failed + self.query(|| async { + sql_execute!( + [self, &pool] + " + DELETE FROM workflow_sub_workflow_events + WHERE location = ? + ", + location, + ) + .await + }) + .await?; + + Err(err) + } + } + } + + async fn update_workflow_tags( + &self, + workflow_id: Uuid, + workflow_name: &str, + tags: &serde_json::Value, + ) -> WorkflowResult<()> { + self.pools + .fdb()? + .run(|tx, _mc| async move { + let tags_subspace = self + .subspace + .subspace(&keys::workflow::TagKey::subspace(workflow_id)); + + // Read old tags + let tag_keys = tx + .get_ranges_keyvalues( + fdb::RangeOption { + mode: StreamingMode::WantAll, + ..(&tags_subspace).into() + }, + SERIALIZABLE, + ) + .map(|res| match res { + Ok(entry) => self + .subspace + .unpack::(entry.key()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into())), + Err(err) => Err(Into::::into(err)), + }) + .try_collect::>() + .await?; + + // Clear old tags + tx.clear_subspace_range(&tags_subspace); + + // Clear old "by name and first tag" secondary index + for key in tag_keys { + keys::workflow::ByNameAndTagKey::new( + workflow_name.to_string(), + key.k, + key.v, + workflow_id, + ); + } + + // Write new tags + let tags = tags + .as_object() + .ok_or_else(|| WorkflowError::InvalidTags("must be an object".to_string())) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))? + .into_iter() + .map(|(k, v)| { + Ok(( + k.clone(), + cjson::to_string(v).map_err(WorkflowError::CjsonSerializeTags)?, + )) + }) + .collect::>>() + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?; + + for (k, v) in &tags { + let tag_key = keys::workflow::TagKey::new(workflow_id, k.clone(), v.clone()); + tx.set( + &self.subspace.pack(&tag_key), + &tag_key + .serialize(()) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + + // Write new "by name and first tag" secondary index + let by_name_and_tag_key = keys::workflow::ByNameAndTagKey::new( + workflow_name.to_string(), + k.clone(), + v.clone(), + workflow_id, + ); + let rest_of_tags = tags + .iter() + .filter(|(k2, _)| k2 != k) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + tx.set( + &self.subspace.pack(&by_name_and_tag_key), + &by_name_and_tag_key + .serialize(rest_of_tags) + .map_err(|x| fdb::FdbBindingError::CustomError(x.into()))?, + ); + } + + Ok(()) + }) + .await?; + + Ok(()) + } + + async fn commit_workflow_activity_event( + &self, + workflow_id: Uuid, + location: &Location, + version: usize, + event_id: &EventId, + create_ts: i64, + input: &serde_json::value::RawValue, + res: Result<&serde_json::value::RawValue, &str>, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(workflow_id).await?; + let input_hash = event_id.input_hash.to_be_bytes(); + + match res { + Ok(output) => { + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_activity_events ( + location, + version, + activity_name, + input_hash, + input, + output, + create_ts, + loop_location + ) + VALUES (jsonb(?), ?, ?, ?, jsonb(?), jsonb(?), ?, jsonb(?)) + ON CONFLICT (location) DO UPDATE + SET output = EXCLUDED.output + ", + location, + version as i64, + &event_id.name, + input_hash.as_slice(), + sqlx::types::Json(input), + sqlx::types::Json(output), + create_ts, + loop_location, + ) + .await + }) + .await?; + } + Err(err) => { + self.query(|| async { + // Attempt to use an existing connection + let mut conn = if let Some(conn) = pool.try_acquire() { + conn + } else { + // Create a new connection + pool.acquire().await? + }; + let mut tx = conn.begin_immediate().await?; + + sql_execute!( + [self, @tx &mut tx] + " + INSERT INTO workflow_activity_events ( + location, + version, + activity_name, + input_hash, + input, + create_ts, + loop_location + ) + VALUES (jsonb(?), ?, ?, ?, jsonb(?), ?, jsonb(?)) + ON CONFLICT (location) DO NOTHING + ", + location, + version as i64, + &event_id.name, + input_hash.as_slice(), + sqlx::types::Json(input), + create_ts, + loop_location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + INSERT INTO workflow_activity_errors ( + location, activity_name, error, ts + ) + VALUES (jsonb(?), ?, ?, ?) + ", + location, + &event_id.name, + err, + rivet_util::timestamp::now(), + ) + .await?; + + tx.commit().await.map_err(Into::into) + }) + .await?; + } + } + + Ok(()) + } + + async fn commit_workflow_message_send_event( + &self, + from_workflow_id: Uuid, + location: &Location, + version: usize, + tags: &serde_json::Value, + message_name: &str, + body: &serde_json::value::RawValue, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(from_workflow_id).await?; + + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_message_send_events ( + location, version, tags, message_name, body, create_ts, loop_location + ) + VALUES (jsonb(?), ?, jsonb(?), ?, jsonb(?), ?, jsonb(?)) + ", + location, + version as i64, + tags, + message_name, + sqlx::types::Json(body), + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await?; + + Ok(()) + } + + async fn upsert_workflow_loop_event( + &self, + workflow_id: Uuid, + location: &Location, + version: usize, + iteration: usize, + state: &serde_json::value::RawValue, + output: Option<&serde_json::value::RawValue>, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(workflow_id).await?; + + self.query(|| async { + // Attempt to use an existing connection + let mut conn = if let Some(conn) = pool.try_acquire() { + conn + } else { + // Create a new connection + pool.acquire().await? + }; + let mut tx = conn.begin_immediate().await?; + + sql_execute!( + [self, @tx &mut tx] + " + INSERT INTO workflow_loop_events ( + location, + version, + iteration, + state, + output, + create_ts, + loop_location + ) + VALUES (jsonb(?), ?, ?, jsonb(?), jsonb(?), ?, jsonb(?)) + ON CONFLICT (location) DO UPDATE + SET + iteration = ?, + state = jsonb(?), + output = jsonb(?) + ", + location, + version as i64, + iteration as i64, + sqlx::types::Json(state), + output.map(sqlx::types::Json), + rivet_util::timestamp::now(), + loop_location, + iteration as i64, + sqlx::types::Json(state), + output.map(sqlx::types::Json), + ) + .await?; + + // 0-th iteration is the initial insertion + if iteration != 0 { + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_activity_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_signal_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_sub_workflow_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_signal_send_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_message_send_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_loop_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_sleep_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_branch_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_removed_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + + sql_execute!( + [self, @tx &mut tx] + " + UPDATE workflow_version_check_events + SET forgotten = TRUE + WHERE loop_location = jsonb(?) AND NOT forgotten + ", + location, + ) + .await?; + } + + tx.commit().await.map_err(WorkflowError::Sqlx)?; + + Ok(()) + }) + .await?; + + Ok(()) + } + + async fn commit_workflow_sleep_event( + &self, + from_workflow_id: Uuid, + location: &Location, + version: usize, + deadline_ts: i64, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(from_workflow_id).await?; + + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_sleep_events ( + location, version, deadline_ts, create_ts, state, loop_location + ) + VALUES (jsonb(?), ?, ?, ?, ?, jsonb(?)) + ", + location, + version as i64, + deadline_ts, + rivet_util::timestamp::now(), + SleepState::Normal as i64, + loop_location, + ) + .await + }) + .await?; + + Ok(()) + } + + async fn update_workflow_sleep_event_state( + &self, + from_workflow_id: Uuid, + location: &Location, + state: SleepState, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(from_workflow_id).await?; + + self.query(|| async { + sql_execute!( + [self, &pool] + " + UPDATE workflow_sleep_events + SET state = ? + WHERE location = ? + ", + state as i64, + location, + ) + .await + }) + .await?; + + Ok(()) + } + + async fn commit_workflow_branch_event( + &self, + from_workflow_id: Uuid, + location: &Location, + version: usize, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(from_workflow_id).await?; + + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_branch_events ( + location, version, create_ts, loop_location + ) + VALUES (jsonb(?), ?, ?, jsonb(?)) + ", + location, + version as i64, + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await?; + + Ok(()) + } + + async fn commit_workflow_removed_event( + &self, + from_workflow_id: Uuid, + location: &Location, + event_type: EventType, + event_name: Option<&str>, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(from_workflow_id).await?; + + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_removed_events ( + location, event_type, event_name, create_ts, loop_location + ) + VALUES (jsonb(?), ?, ?, ?, jsonb(?)) + ", + location, + event_type as i64, + event_name, + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await?; + + Ok(()) + } + + async fn commit_workflow_version_check_event( + &self, + from_workflow_id: Uuid, + location: &Location, + version: usize, + loop_location: Option<&Location>, + ) -> WorkflowResult<()> { + let pool = self.pools.sqlite(from_workflow_id).await?; + + self.query(|| async { + sql_execute!( + [self, &pool] + " + INSERT INTO workflow_version_check_events ( + location, version, create_ts, loop_location + ) + VALUES (jsonb(?), ?, ?, jsonb(?)) + ", + location, + version as i64, + rivet_util::timestamp::now(), + loop_location, + ) + .await + }) + .await?; + + Ok(()) + } +} + +struct PartialWorkflow { + pub workflow_id: Uuid, + pub workflow_name: String, + pub create_ts: i64, + pub ray_id: Uuid, + pub input: Box, +} diff --git a/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/sqlite/migrations/20250122212060_init.sql b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/sqlite/migrations/20250122212060_init.sql new file mode 100644 index 0000000000..d030795614 --- /dev/null +++ b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/sqlite/migrations/20250122212060_init.sql @@ -0,0 +1,198 @@ +-- Activity events +CREATE TABLE workflow_activity_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + activity_name TEXT NOT NULL, + input_hash BLOB NOT NULL, -- u64 + input BLOB NOT NULL, -- JSONB + -- Null if incomplete + output BLOB, -- JSONB + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_activity_events_location_idx +ON workflow_activity_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_activity_events_loop_location_idx +ON workflow_activity_events (loop_location) +WHERE NOT forgotten; + +CREATE TABLE workflow_activity_errors ( + location BLOB PRIMARY KEY, -- JSONB + activity_name TEXT NOT NULL, + error TEXT NOT NULL, + ts INT NOT NULL +) STRICT; + +-- Signal events +CREATE TABLE workflow_signal_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + signal_id BLOB NOT NULL, -- UUID + signal_name TEXT NOT NULL, + body BLOB NOT NULL, -- JSONB + ack_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_signal_events_location_idx +ON workflow_signal_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_signal_events_loop_location_idx +ON workflow_signal_events (loop_location) +WHERE NOT forgotten; + +-- Sub workflow events +CREATE TABLE workflow_sub_workflow_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + sub_workflow_id BLOB NOT NULL, -- UUID + sub_workflow_name TEXT NOT NULL, + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_sub_workflow_events_location_idx +ON workflow_sub_workflow_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_sub_workflow_events_loop_location_idx +ON workflow_sub_workflow_events (loop_location) +WHERE NOT forgotten; + +-- Signal send events +CREATE TABLE workflow_signal_send_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + signal_id BLOB NOT NULL, -- UUID + signal_name TEXT NOT NULL, + body BLOB NOT NULL, -- JSONB + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_signal_send_events_location_idx +ON workflow_signal_send_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_signal_send_events_loop_location_idx +ON workflow_signal_send_events (loop_location) +WHERE NOT forgotten; + +-- Message send events +CREATE TABLE workflow_message_send_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + tags BLOB NOT NULL, -- JSONB + message_name TEXT NOT NULL, + body BLOB NOT NULL, -- JSONB + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_message_send_events_location_idx +ON workflow_message_send_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_message_send_events_loop_location_idx +ON workflow_message_send_events (loop_location) +WHERE NOT forgotten; + +-- Loop events +CREATE TABLE workflow_loop_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + iteration INT NOT NULL, + state BLOB NOT NULL, -- JSONB + output BLOB, -- JSONB + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_loop_events_location_idx +ON workflow_loop_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_loop_events_loop_location_idx +ON workflow_loop_events (loop_location) +WHERE NOT forgotten; + +-- Sleep events +CREATE TABLE workflow_sleep_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + deadline_ts INT NOT NULL, + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + state INT NOT NULL DEFAULT 0, -- event::SleepState + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_sleep_events_location_idx +ON workflow_sleep_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_sleep_events_loop_location_idx +ON workflow_sleep_events (loop_location) +WHERE NOT forgotten; + +-- Branch events +CREATE TABLE workflow_branch_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_branch_events_location_idx +ON workflow_branch_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_branch_events_loop_location_idx +ON workflow_branch_events (loop_location) +WHERE NOT forgotten; + +-- Removed events +CREATE TABLE workflow_removed_events ( + location BLOB PRIMARY KEY, -- JSONB + event_type INT NOT NULL, -- event::EventType + event_name TEXT NOT NULL, + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_removed_events_location_idx +ON workflow_removed_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_removed_events_loop_location_idx +ON workflow_removed_events (loop_location) +WHERE NOT forgotten; + +-- Version check events +CREATE TABLE workflow_version_check_events ( + location BLOB PRIMARY KEY, -- JSONB + version INT NOT NULL, + create_ts INT NOT NULL, + loop_location BLOB, -- JSONB + forgotten INT NOT NULL DEFAULT false -- BOOLEAN +) STRICT; + +CREATE INDEX workflow_version_check_events_location_idx +ON workflow_version_check_events (location) +WHERE NOT forgotten; + +CREATE INDEX workflow_version_check_events_loop_location_idx +ON workflow_version_check_events (loop_location) +WHERE NOT forgotten; diff --git a/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/sqlite/mod.rs b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/sqlite/mod.rs new file mode 100644 index 0000000000..20e3e8316f --- /dev/null +++ b/packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/sqlite/mod.rs @@ -0,0 +1,398 @@ +use std::collections::HashMap; + +use include_dir::{include_dir, Dir, File}; +use indoc::indoc; +use rivet_pools::prelude::*; +use uuid::Uuid; + +use crate::{ + error::{WorkflowError, WorkflowResult}, + history::{ + event::{ + ActivityEvent, Event, EventData, EventId, EventType, LoopEvent, MessageSendEvent, + RemovedEvent, SignalEvent, SignalSendEvent, SleepEvent, SleepState, SubWorkflowEvent, + }, + location::Location, + }, +}; + +type RawJson = sqlx::types::Json>; + +// HACK: We alias global error here because its hardcoded into the sql macros +type GlobalError = WorkflowError; + +const MIGRATIONS_DIR: Dir = + include_dir!("$CARGO_MANIFEST_DIR/src/db/fdb_sqlite_nats/sqlite/migrations"); + +lazy_static::lazy_static! { + // We use a lazy static because this nly needs to be processed once + static ref MIGRATIONS: Migrations = Migrations::new(); +} + +struct Migrations { + last_migration_index: i64, + files: Vec<(i64, &'static File<'static>)>, +} + +impl Migrations { + fn new() -> Self { + let files = MIGRATIONS_DIR + .files() + .map(|file| (parse_migration_index(file), file)) + .collect::>(); + + Migrations { + last_migration_index: files + .iter() + .fold(0, |max_index, (index, _)| max_index.max(*index)), + files, + } + } +} + +// TODO: Used to stub the sql macros, find better solution +struct SqlStub {} + +impl SqlStub { + // For sql macro + pub fn name(&self) -> &str { + super::CONTEXT_NAME + } +} + +/// Runs migrations that have not been run yet. +pub async fn init(workflow_id: Uuid, pool: &SqlitePool) -> WorkflowResult<()> { + // Create migrations table + sql_execute!( + [SqlStub {}, pool] + " + CREATE TABLE IF NOT EXISTS _migrations( + last_index INTEGER NOT NULL, + locked INTEGER NOT NULL, + tainted INTEGER NOT NULL + ) + ", + ) + .await + .map_err(WorkflowError::into_migration_err)?; + + // Initialize state row if not exists + sql_execute!( + [SqlStub {}, pool] + " + INSERT INTO _migrations (last_index, locked, tainted) + SELECT 0, 0, 0 + WHERE NOT EXISTS (SELECT 1 FROM _migrations) + ", + ) + .await + .map_err(WorkflowError::into_migration_err)?; + + // Attempt to get lock on migrations table if migrations should be ran + let migrations_row = sql_fetch_optional!( + [SqlStub {}, (i64,), pool] + " + UPDATE _migrations + SET locked = TRUE + WHERE + last_index < ? AND + NOT locked AND + NOT tainted + RETURNING last_index + ", + MIGRATIONS.last_migration_index, + ) + .await + .map_err(WorkflowError::into_migration_err)?; + + // Could not get lock + let Some((last_index,)) = migrations_row else { + let (locked, tainted) = sql_fetch_one!( + [SqlStub {}, (bool, bool), pool] + " + SELECT locked, tainted FROM _migrations + ", + ) + .await + .map_err(WorkflowError::into_migration_err)?; + + let msg = if tainted { + "tainted" + } else if locked { + "already locked" + } else { + // Already on latest migration + return Ok(()); + }; + + return Err(WorkflowError::MigrationLock(workflow_id, msg.to_string())); + }; + + tracing::debug!(?workflow_id, "running sqlite migrations"); + + // Run migrations + if let Err(err) = run_migrations(&pool, last_index).await { + tracing::debug!(?workflow_id, "sqlite migrations failed"); + + // Mark as tainted + sql_execute!( + [SqlStub {}, pool] + " + UPDATE _migrations + SET + locked = FALSE, + tainted = TRUE + ", + ) + .await + .map_err(WorkflowError::into_migration_err)?; + + return Err(err); + } else { + tracing::debug!(?workflow_id, "sqlite migrations succeeded"); + + // Ack latest migration + sql_execute!( + [SqlStub {}, pool] + " + UPDATE _migrations + SET + locked = FALSE, + last_index = ? + ", + MIGRATIONS.last_migration_index, + ) + .await + .map_err(WorkflowError::into_migration_err)?; + } + + Ok(()) +} + +async fn run_migrations(pool: &SqlitePool, last_index: i64) -> WorkflowResult<()> { + for (idx, file) in &*MIGRATIONS.files { + // Skip already applied migrations + if *idx <= last_index { + continue; + } + + sql_execute!( + [SqlStub {}, pool] + file.contents_utf8().unwrap(), + ) + .await + .map_err(WorkflowError::into_migration_err)?; + } + + Ok(()) +} + +fn parse_migration_index(file: &File) -> i64 { + file.path() + .file_name() + .unwrap() + .to_str() + .unwrap() + .split_once('_') + .expect("invalid migration name") + .0 + .parse() + .expect("invalid migration index") +} + +/// Stores data for all event types in one. +#[derive(Debug, sqlx::FromRow)] +pub struct AmalgamEventRow { + location: Location, + version: i64, + event_type: i64, + name: Option, + auxiliary_id: Option, + hash: Option>, + input: Option, + output: Option, + create_ts: Option, + error_count: Option, + iteration: Option, + deadline_ts: Option, + state: Option, + inner_event_type: Option, +} + +impl TryFrom for Event { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + let event_type = value + .event_type + .try_into() + .map_err(|_| WorkflowError::IntegerConversion)?; + let event_type = EventType::from_repr(event_type) + .ok_or_else(|| WorkflowError::InvalidEventType(value.event_type))?; + + Ok(Event { + coordinate: value.location.tail().cloned().expect("empty location"), + version: value + .version + .try_into() + .map_err(|_| WorkflowError::IntegerConversion)?, + data: match event_type { + EventType::Activity => EventData::Activity(value.try_into()?), + EventType::Signal => EventData::Signal(value.try_into()?), + EventType::SignalSend => EventData::SignalSend(value.try_into()?), + EventType::MessageSend => EventData::MessageSend(value.try_into()?), + EventType::SubWorkflow => EventData::SubWorkflow(value.try_into()?), + EventType::Loop => EventData::Loop(value.try_into()?), + EventType::Sleep => EventData::Sleep(value.try_into()?), + EventType::Branch => EventData::Branch, + EventType::Removed => EventData::Removed(value.try_into()?), + EventType::VersionCheck => EventData::VersionCheck, + }, + }) + } +} + +impl TryFrom for ActivityEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + Ok(ActivityEvent { + event_id: EventId::from_be_bytes( + value.name.ok_or(WorkflowError::MissingEventData)?, + value.hash.ok_or(WorkflowError::MissingEventData)?, + )?, + create_ts: value.create_ts.ok_or(WorkflowError::MissingEventData)?, + output: value.output.map(|x| x.0), + error_count: value + .error_count + .ok_or(WorkflowError::MissingEventData)? + .try_into() + .map_err(|_| WorkflowError::IntegerConversion)?, + }) + } +} + +impl TryFrom for SignalEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + Ok(SignalEvent { + name: value.name.ok_or(WorkflowError::MissingEventData)?, + body: value + .output + .map(|x| x.0) + .ok_or(WorkflowError::MissingEventData)?, + }) + } +} + +impl TryFrom for SignalSendEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + Ok(SignalSendEvent { + signal_id: value.auxiliary_id.ok_or(WorkflowError::MissingEventData)?, + name: value.name.ok_or(WorkflowError::MissingEventData)?, + }) + } +} + +impl TryFrom for MessageSendEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + Ok(MessageSendEvent { + name: value.name.ok_or(WorkflowError::MissingEventData)?, + }) + } +} + +impl TryFrom for SubWorkflowEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + Ok(SubWorkflowEvent { + sub_workflow_id: value.auxiliary_id.ok_or(WorkflowError::MissingEventData)?, + name: value.name.ok_or(WorkflowError::MissingEventData)?, + }) + } +} + +impl TryFrom for LoopEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + Ok(LoopEvent { + state: value.input.ok_or(WorkflowError::MissingEventData)?.0, + output: value.output.map(|x| x.0), + iteration: value + .iteration + .ok_or(WorkflowError::MissingEventData)? + .try_into() + .map_err(|_| WorkflowError::IntegerConversion)?, + }) + } +} + +impl TryFrom for SleepEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + let state = value.state.ok_or(WorkflowError::MissingEventData)?; + + Ok(SleepEvent { + deadline_ts: value.deadline_ts.ok_or(WorkflowError::MissingEventData)?, + state: SleepState::from_repr(state.try_into()?) + .ok_or_else(|| WorkflowError::InvalidSleepState(state))?, + }) + } +} + +impl TryFrom for RemovedEvent { + type Error = WorkflowError; + + fn try_from(value: AmalgamEventRow) -> WorkflowResult { + let event_type = value + .inner_event_type + .ok_or(WorkflowError::MissingEventData)?; + + Ok(RemovedEvent { + name: value.name, + event_type: EventType::from_repr(event_type.try_into()?) + .ok_or_else(|| WorkflowError::InvalidEventType(event_type))?, + }) + } +} + +/// Takes all workflow events (each with their own location) and combines them via enum into a hashmap of the +/// following structure: +/// +/// Given the location [1, 2, 3], 3 is the index and [1, 2] is the root location +/// +/// HashMap { +/// [1, 2]: [ +/// example signal event, +/// example activity event, +/// example sub workflow event, +/// example activity event (this is [1, 2, 3]) +/// ], +/// } +pub fn build_history( + event_rows: Vec, +) -> WorkflowResult>> { + let mut events_by_location: HashMap> = HashMap::new(); + + for event_row in event_rows { + events_by_location + .entry(event_row.location.root()) + .or_default() + .push(event_row.try_into()?); + } + + for events in events_by_location.values_mut() { + // Events are already mostly sorted themselves so this should be fairly cheap + events.sort_by_key(|event| event.coordinate().clone()); + } + + Ok(events_by_location) +} diff --git a/packages/common/chirp-workflow/core/src/db/mod.rs b/packages/common/chirp-workflow/core/src/db/mod.rs index cad07f0bbb..e1952bc029 100644 --- a/packages/common/chirp-workflow/core/src/db/mod.rs +++ b/packages/common/chirp-workflow/core/src/db/mod.rs @@ -13,6 +13,8 @@ use crate::{ mod crdb_nats; pub use crdb_nats::DatabaseCrdbNats; +mod fdb_sqlite_nats; +pub use fdb_sqlite_nats::DatabaseFdbSqliteNats; pub type DatabaseHandle = Arc; @@ -20,6 +22,10 @@ pub type DatabaseHandle = Arc; // manually in the driver. #[async_trait::async_trait] pub trait Database: Send { + fn from_pools(pools: rivet_pools::Pools) -> Result, rivet_pools::Error> + where + Self: Sized; + /// When using a wake worker instead of a polling worker, this function will return once the worker /// should fetch the database again. async fn wake(&self) -> WorkflowResult<()> { @@ -29,7 +35,8 @@ pub trait Database: Send { ); } - /// Writes a new workflow to the database. + /// Writes a new workflow to the database. If unique is set, this should return the existing workflow ID + /// (if one exists) instead of the given workflow ID. async fn dispatch_workflow( &self, ray_id: Uuid, @@ -41,7 +48,14 @@ pub trait Database: Send { ) -> WorkflowResult; /// Retrieves a workflow with the given ID. - async fn get_workflow(&self, id: Uuid) -> WorkflowResult>; + async fn get_workflow(&self, workflow_id: Uuid) -> WorkflowResult>; + + /// Retrieves a workflow with the given name and tags. + async fn find_workflow( + &self, + workflow_name: &str, + tags: &serde_json::Value, + ) -> WorkflowResult>; /// Pulls workflows for processing by the worker. Will only pull workflows with names matching the filter. async fn pull_workflows( @@ -51,16 +65,18 @@ pub trait Database: Send { ) -> WorkflowResult>; /// Mark a workflow as completed. - async fn commit_workflow( + async fn complete_workflow( &self, workflow_id: Uuid, + workflow_name: &str, output: &serde_json::value::RawValue, ) -> WorkflowResult<()>; - /// Write a workflow failure to the database. - async fn fail_workflow( + /// Write a workflow sleep/failure to the database. + async fn commit_workflow( &self, workflow_id: Uuid, + workflow_name: &str, wake_immediate: bool, wake_deadline_ts: Option, wake_signals: &[&str], @@ -68,26 +84,6 @@ pub trait Database: Send { error: &str, ) -> WorkflowResult<()>; - /// Updates workflow tags. - async fn update_workflow_tags( - &self, - workflow_id: Uuid, - tags: &serde_json::Value, - ) -> WorkflowResult<()>; - - /// Write a workflow activity event to history. - async fn commit_workflow_activity_event( - &self, - workflow_id: Uuid, - location: &Location, - version: usize, - event_id: &EventId, - create_ts: i64, - input: &serde_json::value::RawValue, - output: Result<&serde_json::value::RawValue, &str>, - loop_location: Option<&Location>, - ) -> WorkflowResult<()>; - /// Pulls the oldest signal with the given filter. Pulls from regular and tagged signals. async fn pull_next_signal( &self, @@ -98,6 +94,14 @@ pub trait Database: Send { loop_location: Option<&Location>, ) -> WorkflowResult>; + /// Retrieves a workflow with the given ID. Can only be called from a workflow context. + async fn get_sub_workflow( + &self, + workflow_id: Uuid, + workflow_name: &str, + sub_workflow_id: Uuid, + ) -> WorkflowResult>; + /// Write a new signal to the database. async fn publish_signal( &self, @@ -161,6 +165,27 @@ pub trait Database: Send { unique: bool, ) -> WorkflowResult; + /// Updates workflow tags. + async fn update_workflow_tags( + &self, + workflow_id: Uuid, + workflow_name: &str, + tags: &serde_json::Value, + ) -> WorkflowResult<()>; + + /// Write a workflow activity event to history. + async fn commit_workflow_activity_event( + &self, + workflow_id: Uuid, + location: &Location, + version: usize, + event_id: &EventId, + create_ts: i64, + input: &serde_json::value::RawValue, + output: Result<&serde_json::value::RawValue, &str>, + loop_location: Option<&Location>, + ) -> WorkflowResult<()>; + /// Writes a message send event to history. async fn commit_workflow_message_send_event( &self, @@ -180,6 +205,7 @@ pub trait Database: Send { location: &Location, version: usize, iteration: usize, + state: &serde_json::value::RawValue, output: Option<&serde_json::value::RawValue>, loop_location: Option<&Location>, ) -> WorkflowResult<()>; @@ -252,7 +278,6 @@ pub struct PulledWorkflow { pub create_ts: i64, pub ray_id: Uuid, pub input: Box, - pub wake_deadline_ts: Option, pub events: HashMap>, } diff --git a/packages/common/chirp-workflow/core/src/error.rs b/packages/common/chirp-workflow/core/src/error.rs index 86b97089c2..02fc805e76 100644 --- a/packages/common/chirp-workflow/core/src/error.rs +++ b/packages/common/chirp-workflow/core/src/error.rs @@ -1,5 +1,6 @@ use std::time::{SystemTime, UNIX_EPOCH}; +use foundationdb as fdb; use global_error::GlobalError; use tokio::time::Instant; use uuid::Uuid; @@ -78,6 +79,12 @@ pub enum WorkflowError { #[error("deserialize message: {0}")] DeserializeMessage(serde_json::Error), + #[error("serialize loop state: {0}")] + SerializeLoopState(serde_json::Error), + + #[error("deserialize loop state: {0}")] + DeserializeLoopState(serde_json::Error), + #[error("failed to serialize cjson tags: {0:?}")] CjsonSerializeTags(cjson::Error), @@ -87,8 +94,8 @@ pub enum WorkflowError { #[error("failed to deserialize tags: {0}")] DeserializeTags(serde_json::Error), - #[error("tags must be a json object")] - InvalidTags, + #[error("invalid tags: {0}")] + InvalidTags(String), #[error("failed to serialize loop output: {0}")] SerializeLoopOutput(serde_json::Error), @@ -138,6 +145,9 @@ pub enum WorkflowError { #[error("sql error: {0}")] Sqlx(#[from] sqlx::Error), + #[error("fdb error: {0}")] + Fdb(#[from] fdb::FdbBindingError), + #[error("max sql retries, last error: {0}")] MaxSqlRetries(sqlx::Error), @@ -173,6 +183,15 @@ pub enum WorkflowError { #[error("invalid version: {0}")] InvalidVersion(String), + + #[error("tagged signals are disabled for this workflow engine's database driver. use workflow directed signals instead")] + TaggedSignalsDisabled, + + #[error("failed to acquire migration lock (workflow id {0}): {1}")] + MigrationLock(Uuid, String), + + #[error("migration failed: {0}")] + Migration(sqlx::Error), } impl WorkflowError { @@ -242,4 +261,11 @@ impl WorkflowError { None } } + + pub(crate) fn into_migration_err(self) -> Self { + match self { + WorkflowError::Sqlx(err) => WorkflowError::Migration(err), + _ => self, + } + } } diff --git a/packages/common/chirp-workflow/core/src/history/event.rs b/packages/common/chirp-workflow/core/src/history/event.rs index d2a364e44d..e7ad15820d 100644 --- a/packages/common/chirp-workflow/core/src/history/event.rs +++ b/packages/common/chirp-workflow/core/src/history/event.rs @@ -163,6 +163,7 @@ pub struct SubWorkflowEvent { #[derive(Debug)] pub struct LoopEvent { + pub(crate) state: Box, /// If the loop completes, this will be some. pub(crate) output: Option>, pub iteration: usize, @@ -226,7 +227,7 @@ impl EventId { } } - pub fn from_bytes(name: String, input_hash: Vec) -> WorkflowResult { + pub fn from_le_bytes(name: String, input_hash: Vec) -> WorkflowResult { Ok(EventId { name, input_hash: u64::from_le_bytes( @@ -236,4 +237,15 @@ impl EventId { ), }) } + + pub fn from_be_bytes(name: String, input_hash: Vec) -> WorkflowResult { + Ok(EventId { + name, + input_hash: u64::from_be_bytes( + input_hash + .try_into() + .map_err(|_| WorkflowError::IntegerConversion)?, + ), + }) + } } diff --git a/packages/common/chirp-workflow/core/src/history/location.rs b/packages/common/chirp-workflow/core/src/history/location.rs index 1712482648..7978b3f96a 100644 --- a/packages/common/chirp-workflow/core/src/history/location.rs +++ b/packages/common/chirp-workflow/core/src/history/location.rs @@ -156,3 +156,67 @@ impl Deref for Coordinate { &self.0 } } + +// Implements sqlx types for `Location` +mod sqlx { + use super::Location; + + impl sqlx::Type for Location + where + DB: sqlx::Database, + serde_json::Value: sqlx::Type, + { + fn type_info() -> DB::TypeInfo { + >::type_info() + } + + fn compatible(ty: &DB::TypeInfo) -> bool { + >::compatible(ty) + } + } + + impl<'q, DB> sqlx::Encode<'q, DB> for Location + where + serde_json::Value: sqlx::Encode<'q, DB>, + DB: sqlx::Database, + { + fn encode_by_ref( + &self, + buf: &mut DB::ArgumentBuffer<'q>, + ) -> Result { + >::encode(serialize_location(self), buf) + } + } + + impl<'r, DB> sqlx::Decode<'r, DB> for Location + where + sqlx::types::Json]>>: sqlx::Decode<'r, DB>, + DB: sqlx::Database, + { + fn decode(value: DB::ValueRef<'r>) -> Result { + let value = + ]>> as sqlx::Decode<'r, DB>>::decode(value)?; + + Ok(IntoIterator::into_iter(value.0).collect()) + } + } + + // TODO: Implement serde serialize and deserialize for `Location` + /// Convert location to json as `number[][]`. + fn serialize_location(location: &Location) -> serde_json::Value { + serde_json::Value::Array( + location + .as_ref() + .iter() + .map(|coord| { + serde_json::Value::Array( + coord + .iter() + .map(|x| serde_json::Value::Number((*x).into())) + .collect(), + ) + }) + .collect(), + ) + } +} diff --git a/packages/common/chirp-workflow/core/src/message.rs b/packages/common/chirp-workflow/core/src/message.rs index 21220f82d0..89ea1c8c6c 100644 --- a/packages/common/chirp-workflow/core/src/message.rs +++ b/packages/common/chirp-workflow/core/src/message.rs @@ -5,8 +5,6 @@ use uuid::Uuid; use crate::error::{WorkflowError, WorkflowResult}; -pub const WORKER_WAKE_SUBJECT: &str = "chirp.workflow.worker.wake"; - pub trait Message: Debug + Send + Sync + Serialize + DeserializeOwned + 'static { const NAME: &'static str; const TAIL_TTL: std::time::Duration; @@ -42,7 +40,7 @@ impl AsTags for serde_json::Value { fn as_tags(&self) -> WorkflowResult { match self { serde_json::Value::Object(_) => Ok(self.clone()), - _ => Err(WorkflowError::InvalidTags), + _ => Err(WorkflowError::InvalidTags("must be an object".to_string())), } } @@ -51,7 +49,7 @@ impl AsTags for serde_json::Value { serde_json::Value::Object(_) => { cjson::to_string(&self).map_err(WorkflowError::CjsonSerializeTags) } - _ => Err(WorkflowError::InvalidTags), + _ => Err(WorkflowError::InvalidTags("must be an object".to_string())), } } } diff --git a/packages/common/chirp-workflow/core/tests/common.rs b/packages/common/chirp-workflow/core/tests/common.rs new file mode 100644 index 0000000000..54c9537c77 --- /dev/null +++ b/packages/common/chirp-workflow/core/tests/common.rs @@ -0,0 +1,132 @@ +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Once, + }, + time::Duration, +}; + +use tokio::process::Command; +use tracing_subscriber::prelude::*; + +pub async fn start_nats() { + Command::new("docker") + .arg("rm") + .arg("test-nats") + .arg("--force") + .status() + .await + .unwrap(); + + let status = Command::new("docker") + .arg("run") + .arg("--rm") + .arg("-p") + .arg("4222:4222") + .arg("--name") + .arg("test-nats") + .arg("nats:latest") + .status() + .await + .unwrap(); + + assert!(status.success()); +} + +pub async fn start_redis() { + Command::new("docker") + .arg("rm") + .arg("test-redis") + .arg("--force") + .status() + .await + .unwrap(); + + let status = Command::new("docker") + .arg("run") + .arg("--rm") + .arg("-p") + .arg("6379:6379") + .arg("--name") + .arg("test-redis") + .arg("redis:latest") + .status() + .await + .unwrap(); + + assert!(status.success()); +} + +pub async fn start_fdb() { + Command::new("docker") + .arg("rm") + .arg("test-fdb") + .arg("--force") + .status() + .await + .unwrap(); + + let status = Command::new("docker") + .arg("run") + .arg("--rm") + .arg("-p") + .arg("4500:4500") + .arg("--name") + .arg("test-fdb") + .arg("-e") + .arg("FDB_CLUSTER_FILE_CONTENTS=fdb:fdb@127.0.0.1:4500") + // See docs-internal/infrastructure/fdb/AVX.md + .arg("foundationdb/foundationdb:7.1.60") + .status() + .await + .unwrap(); + + assert!(status.success()); +} + +pub async fn create_fdb_db() { + loop { + // Create db + let status = Command::new("docker") + .arg("exec") + .arg("test-fdb") + .arg("fdbcli") + .arg("--exec") + .arg(r#"configure new single ssd"#) + .status() + .await + .unwrap(); + + if status.success() { + break; + } else { + tracing::error!("failed to create fdb database"); + } + + tokio::time::sleep(Duration::from_secs(1)).await; + } +} + +static SETUP_DEPENDENCIES: AtomicBool = AtomicBool::new(false); + +pub async fn setup_dependencies() { + if !SETUP_DEPENDENCIES.swap(true, Ordering::SeqCst) { + tokio::spawn(start_nats()); + tokio::spawn(start_redis()); + tokio::spawn(start_fdb()); + create_fdb_db().await; + } +} + +static SETUP_TRACING: Once = Once::new(); +pub fn setup_tracing() { + SETUP_TRACING.call_once(|| { + tracing_subscriber::registry() + .with( + tracing_logfmt::builder() + .layer() + .with_filter(tracing_subscriber::filter::LevelFilter::INFO), + ) + .init(); + }); +} diff --git a/packages/common/chirp-workflow/core/tests/integration.rs b/packages/common/chirp-workflow/core/tests/integration.rs new file mode 100644 index 0000000000..d542063c6c --- /dev/null +++ b/packages/common/chirp-workflow/core/tests/integration.rs @@ -0,0 +1,149 @@ +use chirp_workflow::db::Database; +use chirp_workflow::prelude::*; +use serde_json::json; +use uuid::Uuid; + +mod common; +use common::*; + +#[tokio::test(flavor = "multi_thread")] +async fn fdb_sqlite_nats_driver() { + setup_tracing(); + setup_dependencies().await; + + let ctx = chirp_workflow::prelude::TestCtx::from_env::( + "fdb_sqlite_nats_driver", + true, + ) + .await; + let config = ctx.config().clone(); + let pools = ctx.pools().clone(); + + // // CLEAR DB + // pools + // .fdb() + // .unwrap() + // .run(|tx, _mc| async move { + // tx.clear_range(&[0], &[255]); + // Ok(()) + // }) + // .await + // .unwrap(); + // tokio::time::sleep(std::time::Duration::from_millis(250)).await; + + let mut reg = Registry::new(); + reg.register_workflow::().unwrap(); + + let db = db::DatabaseFdbSqliteNats::from_pools(pools.clone()).unwrap(); + + // let workflow_id = Uuid::new_v4(); + // let input = serde_json::value::RawValue::from_string("null".to_string()).unwrap(); + + // db.dispatch_workflow( + // Uuid::new_v4(), + // workflow_id, + // "workflow_name", + // Some(&json!({ "bald": "eagle" })), + // &input, + // false, + // ) + // .await.unwrap(); + + // let res = db.find_workflow("workflow_name", &json!({ + // "bald": "eagle", + // "fat": "man" + // })).await.unwrap(); + // tracing::info!(?res); + + // db.update_workflow_tags(workflow_id, "workflow_name", &json!({ + // "bald": "eagle", + // "fat": "man" + // })) + // .await + // .unwrap(); + + // let res = db.find_workflow("workflow_name", &json!({ + // "bald": "eagle", + // "fat": "man" + // })).await.unwrap(); + // tracing::info!(?res); + + let worker = Worker::new(reg.handle(), db); + + tokio::spawn(async move { + ctx.workflow(def::Input {}) + .tag("foo", "bar") + .dispatch() + .await + .unwrap(); + }) + .await + .unwrap(); + + // Start worker + tokio::select! { + res = worker.poll_start(config, pools) => res.unwrap(), + res = tokio::signal::ctrl_c() => res.unwrap(), + } +} + +mod def { + use chirp_workflow::prelude::*; + use futures_util::FutureExt; + + #[derive(Debug, Serialize, Deserialize)] + pub struct Input {} + + #[workflow] + pub async fn test(ctx: &mut WorkflowCtx, input: &Input) -> GlobalResult<()> { + tracing::info!("hello from workflow"); + + ctx.activity(TestActivityInput { + foo: "bar".to_string(), + }) + .await?; + + let workflow_id = ctx.workflow_id(); + ctx.signal(MySignal { + test: Uuid::new_v4(), + }) + .to_workflow(workflow_id) + .send() + .await?; + + ctx.repeat(|ctx| { + async move { + let sig = ctx.listen::().await?; + tracing::info!(?sig); + + tracing::info!("eepy"); + ctx.sleep(12000).await?; + tracing::info!("eeped"); + + Ok(Loop::<()>::Continue) + } + .boxed() + }) + .await?; + + Ok(()) + } + + #[derive(Debug, Serialize, Deserialize, Hash)] + struct TestActivityInput { + foo: String, + } + + #[activity(TestActivity)] + async fn test_activity(ctx: &ActivityCtx, input: &TestActivityInput) -> GlobalResult<()> { + tracing::info!(?input.foo, "hello from activity"); + + Ok(()) + } + + #[signal("my_signal")] + #[derive(Debug)] + struct MySignal { + test: Uuid, + } +} diff --git a/packages/common/chirp-workflow/macros/src/lib.rs b/packages/common/chirp-workflow/macros/src/lib.rs index 8b1cb8ae81..98902aa825 100644 --- a/packages/common/chirp-workflow/macros/src/lib.rs +++ b/packages/common/chirp-workflow/macros/src/lib.rs @@ -395,21 +395,24 @@ pub fn message(attr: TokenStream, item: TokenStream) -> TokenStream { #[proc_macro_attribute] pub fn workflow_test(_attr: TokenStream, item: TokenStream) -> TokenStream { - let input = syn::parse_macro_input!(item as syn::ItemFn); + let item_fn = syn::parse_macro_input!(item as syn::ItemFn); - let test_ident = &input.sig.ident; - let body = &input.block; + let test_ident = &item_fn.sig.ident; + let body = &item_fn.block; // Check if is async - if input.sig.asyncness.is_none() { + if item_fn.sig.asyncness.is_none() { return error( - input.sig.span(), + item_fn.sig.span(), "the async keyword is missing from the function declaration", ); } // Parse args - let ctx_input = match input.sig.inputs.first().unwrap() { + let Some(first_arg) = item_fn.sig.inputs.first() else { + return error(item_fn.span(), "must have ctx argument"); + }; + let ctx_input = match first_arg { syn::FnArg::Receiver(recv) => { return error(recv.span(), "cannot have receiver argument"); } @@ -479,7 +482,7 @@ fn parse_config(attrs: &[syn::Attribute]) -> syn::Result { .base10_parse()?; } else if ident != "doc" { return Err(syn::Error::new( - ident.span(), + name_value.span(), format!("Unknown config property `{ident}`"), )); } @@ -504,7 +507,7 @@ fn parse_msg_config(attrs: &[syn::Attribute]) -> syn::Result { .base10_parse()?; } else if ident != "doc" { return Err(syn::Error::new( - ident.span(), + name_value.span(), format!("Unknown config property `{ident}`"), )); } @@ -523,7 +526,7 @@ fn parse_empty_config(attrs: &[syn::Attribute]) -> syn::Result<()> { if ident != "doc" { return Err(syn::Error::new( - ident.span(), + name_value.span(), format!("Unknown config property `{ident}`"), )); } diff --git a/packages/common/config/src/config/server/mod.rs b/packages/common/config/src/config/server/mod.rs index db6bf5a5f5..c33026452c 100644 --- a/packages/common/config/src/config/server/mod.rs +++ b/packages/common/config/src/config/server/mod.rs @@ -34,6 +34,8 @@ pub struct Server { pub clickhouse: Option, #[serde(default)] pub prometheus: Option, + #[serde(default)] + pub fdb: Fdb, // Services #[serde(default)] @@ -443,3 +445,17 @@ impl ClickHouseUserRole { } } } + +#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub struct Fdb { + pub connection: String, +} + +impl Default for Fdb { + fn default() -> Self { + Self { + connection: "fdb:fdb@127.0.0.1:4500".to_string(), + } + } +} diff --git a/packages/common/config/src/lib.rs b/packages/common/config/src/lib.rs index a95be24a97..abcaf16c15 100644 --- a/packages/common/config/src/lib.rs +++ b/packages/common/config/src/lib.rs @@ -42,6 +42,10 @@ impl Config { Ok(Self(Arc::new(ConfigData { config }))) } + + pub fn from_root(config: config::Root) -> Self { + Self(Arc::new(ConfigData { config })) + } } impl Deref for Config { diff --git a/packages/common/connection/src/lib.rs b/packages/common/connection/src/lib.rs index c3f09e7480..da43d8a19f 100644 --- a/packages/common/connection/src/lib.rs +++ b/packages/common/connection/src/lib.rs @@ -58,6 +58,10 @@ impl Connection { self.cache.clone() } + pub fn pools(&self) -> &rivet_pools::Pools { + &self.pools + } + pub async fn nats(&self) -> Result { self.pools.nats() } diff --git a/packages/common/fdb-util/Cargo.toml b/packages/common/fdb-util/Cargo.toml new file mode 100644 index 0000000000..12d78b122f --- /dev/null +++ b/packages/common/fdb-util/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "fdb-util" +version.workspace = true +authors.workspace = true +license.workspace = true +edition.workspace = true + +[dependencies] +anyhow = "1.0" +foundationdb.workspace = true +lazy_static = "1.4" +tokio = { version = "1.40.0", features = ["full"] } +tracing = "0.1.40" diff --git a/packages/common/fdb-util/src/lib.rs b/packages/common/fdb-util/src/lib.rs new file mode 100644 index 0000000000..bbb36a0e14 --- /dev/null +++ b/packages/common/fdb-util/src/lib.rs @@ -0,0 +1,51 @@ +use std::{path::{PathBuf, Path}, result::Result::Ok, time::Duration}; + +use anyhow::*; +use foundationdb::{self as fdb, options::DatabaseOption}; + +lazy_static::lazy_static! { + /// Must only be created once per program and must not be dropped until the program is over otherwise all + /// FDB calls will fail with error code 1100. + static ref FDB_NETWORK: fdb::api::NetworkAutoStop = unsafe { fdb::boot() }; +} + +pub fn handle(fdb_cluster_path: &Path) -> Result { + let db = fdb::Database::from_path( + &fdb_cluster_path + .to_str() + .context("bad fdb_cluster_path")? + .to_string(), + ) + .context("failed to create FDB database")?; + db.set_option(DatabaseOption::TransactionRetryLimit(10))?; + + Ok(db) +} + +/// Starts the network thread and spawns a health check task. +pub fn init(fdb_cluster_path: &Path) { + // Initialize lazy static + let _network = &*FDB_NETWORK; + + tokio::spawn(fdb_health_check(fdb_cluster_path.to_path_buf())); +} + +pub async fn fdb_health_check(fdb_cluster_path: PathBuf) -> Result<()> { + let db = handle(&fdb_cluster_path)?; + + loop { + match ::tokio::time::timeout( + Duration::from_secs(3), + db.run(|trx, _maybe_committed| async move { Ok(trx.get(b"", true).await?) }), + ) + .await + { + Ok(res) => { + res?; + } + Err(_) => tracing::error!("fdb missed ping"), + } + + ::tokio::time::sleep(Duration::from_secs(3)).await; + } +} diff --git a/packages/common/operation/core/src/lib.rs b/packages/common/operation/core/src/lib.rs index 56ba99408b..7f9cf30cb3 100644 --- a/packages/common/operation/core/src/lib.rs +++ b/packages/common/operation/core/src/lib.rs @@ -240,6 +240,10 @@ where self.conn.cache_handle() } + pub fn pools(&self) -> &rivet_pools::Pools { + self.conn.pools() + } + pub async fn crdb(&self) -> Result { self.conn.crdb().await } diff --git a/packages/common/pools/Cargo.toml b/packages/common/pools/Cargo.toml index e6206d654d..3a03fd6862 100644 --- a/packages/common/pools/Cargo.toml +++ b/packages/common/pools/Cargo.toml @@ -6,34 +6,42 @@ license.workspace = true edition.workspace = true [dependencies] +anyhow = "1.0" async-nats = "0.33" clickhouse = { version = "0.11.2" } +fdb-util.workspace = true +foundationdb.workspace = true funty = "=1.1.0" # Fixes issue with sqlx dependency, see https://github.com/bitvecto-rs/bitvec/issues/105#issuecomment-778570981 +futures-util = "0.3" global-error.workspace = true +governor = "0.6" hyper = { version = "0.14" } hyper-tls = { version = "0.5.0" } lazy_static = "1.4" rand = "0.8" +rivet-config.workspace = true rivet-metrics.workspace = true +tempfile = "3.13.0" thiserror = "1.0" tokio = { version = "1.40", features = ["tracing"] } tokio-util = "0.7" tracing = "0.1" -governor = "0.6" url = "2.4" -rivet-config.workspace = true +uuid = { version = "1", features = ["v4"] } [dependencies.sqlx] workspace = true features = [ - "runtime-tokio", - "runtime-tokio-native-tls", - "postgres", - "macros", - "uuid", + "bit-vec", "ipnetwork", "json", - "bit-vec", + "macros", + "migrate", + "postgres", + "runtime-tokio-native-tls", + "runtime-tokio", + "sqlite", + "uuid", ] [dependencies.redis] diff --git a/packages/common/pools/src/db/crdb.rs b/packages/common/pools/src/db/crdb.rs index 1b68441ae3..13b5c80c75 100644 --- a/packages/common/pools/src/db/crdb.rs +++ b/packages/common/pools/src/db/crdb.rs @@ -10,8 +10,6 @@ pub async fn setup(config: Config) -> Result { let crdb = &config.server().map_err(Error::Global)?.cockroachdb; tracing::debug!("crdb connecting"); - // let client_name = client_name.clone(); - let mut opts: sqlx::postgres::PgConnectOptions = crdb.url.to_string().parse().map_err(Error::BuildSqlx)?; opts = opts.username(&crdb.username); diff --git a/packages/common/pools/src/db/fdb.rs b/packages/common/pools/src/db/fdb.rs new file mode 100644 index 0000000000..23b92762a6 --- /dev/null +++ b/packages/common/pools/src/db/fdb.rs @@ -0,0 +1,44 @@ +use std::{ops::Deref, sync::Arc}; + +use foundationdb as fdb; +use rivet_config::Config; +use tokio::fs; + +use crate::Error; + +#[derive(Clone)] +pub struct FdbPool { + db: Arc, + // Prevent dropping temp file + _connection_file: Arc, +} + +impl Deref for FdbPool { + type Target = Arc; + + fn deref(&self) -> &Self::Target { + &self.db + } +} + +#[tracing::instrument(skip(config))] +pub async fn setup(config: Config) -> Result { + let connection = &config.server().map_err(Error::Global)?.fdb.connection; + + let temp_file = tempfile::NamedTempFile::new().map_err(Error::BuildFdbConnectionFile)?; + fs::write(temp_file.path(), connection.as_bytes()) + .await + .map_err(Error::BuildFdbConnectionFile)?; + + // Start network + fdb_util::init(temp_file.path()); + + let fdb_handle = fdb_util::handle(&temp_file.path()).map_err(Error::BuildFdb)?; + + tracing::debug!(?connection, "fdb connected"); + + Ok(FdbPool { + db: Arc::new(fdb_handle), + _connection_file: Arc::new(temp_file), + }) +} diff --git a/packages/common/pools/src/db/mod.rs b/packages/common/pools/src/db/mod.rs index f43c0d83eb..725ad6eaad 100644 --- a/packages/common/pools/src/db/mod.rs +++ b/packages/common/pools/src/db/mod.rs @@ -1,4 +1,6 @@ pub mod clickhouse; pub mod crdb; +pub mod fdb; pub mod nats; pub mod redis; +pub mod sqlite; diff --git a/packages/common/pools/src/db/sqlite.rs b/packages/common/pools/src/db/sqlite.rs new file mode 100644 index 0000000000..0201274bd9 --- /dev/null +++ b/packages/common/pools/src/db/sqlite.rs @@ -0,0 +1,112 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use futures_util::{StreamExt, TryStreamExt}; +use sqlx::{ + migrate::MigrateDatabase, + sqlite::{SqliteConnectOptions, SqlitePoolOptions}, + Sqlite, +}; +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::Error; + +pub type SqlitePool = sqlx::SqlitePool; + +#[derive(Clone)] +pub struct SqlitePoolManager { + // TODO: Somehow remove old pools + pools: Arc>>, +} + +impl SqlitePoolManager { + pub fn new() -> Self { + SqlitePoolManager { + pools: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// Get or creates an sqlite pool for the given key + pub async fn get(&self, key: Uuid) -> Result { + let mut pools_guard = self.pools.lock().await; + + let pool = if let Some(pool) = pools_guard.get(&key) { + pool.clone() + } else { + // TODO: Hardcoded for testing + let db_url = format!("sqlite:///home/rivet/rivet-ee/oss/packages/common/chirp-workflow/core/tests/db/{key}.db"); + + tracing::debug!(?key, "sqlite connecting"); + + // Init if doesn't exist + if !Sqlite::database_exists(&db_url) + .await + .map_err(Error::BuildSqlx)? + { + Sqlite::create_database(&db_url) + .await + .map_err(Error::BuildSqlx)?; + } + + let opts: SqliteConnectOptions = db_url.parse().map_err(Error::BuildSqlx)?; + + let pool = SqlitePoolOptions::new() + .max_lifetime_jitter(Duration::from_secs(90)) + // Open connection immediately on startup + .min_connections(1) + .connect_with(opts) + .await + .map_err(Error::BuildSqlx)?; + + // Run at the start of every connection + setup_pragma(&pool).await.map_err(Error::BuildSqlx)?; + + pools_guard.insert(key, pool.clone()); + + tracing::debug!(?key, "sqlite connected"); + + pool + }; + + Ok(pool) + } +} + +async fn setup_pragma(pool: &SqlitePool) -> Result<(), sqlx::Error> { + // Has to be String instead of static str due to a weird compiler bug. This crate will compile just fine + // but chirp-workflow will not and the error has nothing to do with this code + let settings = [ + // Set the journal mode to Write-Ahead Logging 2 for concurrency + "PRAGMA journal_mode = WAL2".to_string(), + // Set synchronous mode to NORMAL for performance and data safety balance + "PRAGMA synchronous = NORMAL".to_string(), + // Set busy timeout to 5 seconds to avoid "database is locked" errors + "PRAGMA busy_timeout = 5000".to_string(), + // Enable foreign key constraint enforcement + "PRAGMA foreign_keys = ON".to_string(), + // Enable auto vacuuming and set it to incremental mode for gradual space reclaiming + "PRAGMA auto_vacuum = INCREMENTAL".to_string(), + ]; + + futures_util::stream::iter(settings) + .map(|setting| { + let pool = pool.clone(); + async move { + // Attempt to use an existing connection + let mut conn = if let Some(conn) = pool.try_acquire() { + conn + } else { + // Create a new connection + pool.acquire().await? + }; + // Result::<_, sqlx::Error>::Ok(()) + + sqlx::query(&setting).execute(&mut *conn).await + } + }) + .buffer_unordered(16) + .try_collect::>() + .await?; + + Ok(()) +} diff --git a/packages/common/pools/src/error.rs b/packages/common/pools/src/error.rs index 24d09534fd..46bb4113cf 100644 --- a/packages/common/pools/src/error.rs +++ b/packages/common/pools/src/error.rs @@ -12,6 +12,9 @@ pub enum Error { #[error("missing clickhouse pool")] MissingClickHousePool, + #[error("missing fdb pool")] + MissingFdbPool, + #[error("tokio join: {0}")] TokioJoin(tokio::task::JoinError), @@ -30,6 +33,12 @@ pub enum Error { #[error("build redis url: {0}")] BuildRedisUrl(url::ParseError), + #[error("build fdb: {0}")] + BuildFdb(anyhow::Error), + + #[error("build fdb connection file: {0}")] + BuildFdbConnectionFile(std::io::Error), + #[error("modify redis url")] ModifyRedisUrl, diff --git a/packages/common/pools/src/lib.rs b/packages/common/pools/src/lib.rs index b5fa5c74ab..2170a296aa 100644 --- a/packages/common/pools/src/lib.rs +++ b/packages/common/pools/src/lib.rs @@ -6,6 +6,6 @@ pub mod prelude; pub mod utils; pub use crate::{ - db::clickhouse::ClickHousePool, db::crdb::CrdbPool, db::nats::NatsPool, db::redis::RedisPool, - error::Error, pools::Pools, + db::clickhouse::ClickHousePool, db::crdb::CrdbPool, db::fdb::FdbPool, db::nats::NatsPool, + db::redis::RedisPool, db::sqlite::SqlitePool, error::Error, pools::Pools, }; diff --git a/packages/common/pools/src/pools.rs b/packages/common/pools/src/pools.rs index 2a91119b68..49ba5cf188 100644 --- a/packages/common/pools/src/pools.rs +++ b/packages/common/pools/src/pools.rs @@ -1,18 +1,25 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + use global_error::{ensure_with, prelude::*, GlobalResult}; use rivet_config::Config; -use std::{collections::HashMap, sync::Arc, time::Duration}; use tokio_util::sync::{CancellationToken, DropGuard}; +use uuid::Uuid; -use crate::{ClickHousePool, CrdbPool, Error, NatsPool, RedisPool}; +use crate::{ + db::sqlite::SqlitePoolManager, ClickHousePool, CrdbPool, Error, FdbPool, NatsPool, RedisPool, + SqlitePool, +}; // TODO: Automatically shutdown all pools on drop pub(crate) struct PoolsInner { pub(crate) _guard: DropGuard, - config: rivet_config::Config, pub(crate) nats: Option, pub(crate) crdb: Option, pub(crate) redis: HashMap, pub(crate) clickhouse: Option, + pub(crate) fdb: Option, + pub(crate) sqlite: SqlitePoolManager, + clickhouse_enabled: bool, } #[derive(Clone)] @@ -25,20 +32,26 @@ impl Pools { let client_name = "rivet".to_string(); let token = CancellationToken::new(); - let (nats, crdb, redis) = tokio::try_join!( + let (nats, crdb, redis, fdb) = tokio::try_join!( crate::db::nats::setup(config.clone(), client_name.clone()), crate::db::crdb::setup(config.clone()), crate::db::redis::setup(config.clone()), + crate::db::fdb::setup(config.clone()), )?; let clickhouse = crate::db::clickhouse::setup(config.clone())?; let pool = Pools(Arc::new(PoolsInner { _guard: token.clone().drop_guard(), - config, nats: Some(nats), crdb: Some(crdb), redis, clickhouse, + fdb: Some(fdb), + sqlite: SqlitePoolManager::new(), + clickhouse_enabled: config + .server + .as_ref() + .map_or(false, |x| x.clickhouse.is_some()), })); pool.clone().start(token); @@ -49,6 +62,43 @@ impl Pools { Ok(pool) } + + // Only for tests + #[tracing::instrument(skip(config))] + pub async fn test(config: Config) -> Result { + // TODO: Choose client name for this service + let client_name = "rivet".to_string(); + let token = CancellationToken::new(); + + let (nats, redis, fdb) = tokio::try_join!( + crate::db::nats::setup(config.clone(), client_name.clone()), + crate::db::redis::setup(config.clone()), + crate::db::fdb::setup(config.clone()), + )?; + + let pool = Pools(Arc::new(PoolsInner { + _guard: token.clone().drop_guard(), + nats: Some(nats), + crdb: None, + redis, + clickhouse: None, + fdb: Some(fdb), + sqlite: SqlitePoolManager::new(), + clickhouse_enabled: config + .server + .as_ref() + .map_or(false, |x| x.clickhouse.is_some()), + })); + pool.clone().start(token); + + tokio::task::Builder::new() + .name("rivet_pools::runtime") + .spawn(runtime(pool.clone(), client_name.clone())) + .map_err(Error::TokioSpawn)?; + + Ok(pool) + } + /// Spawn background tasks required to operate the pool. pub(crate) fn start(self, token: CancellationToken) { let spawn_res = tokio::task::Builder::new() @@ -100,11 +150,7 @@ impl Pools { } pub fn clickhouse_enabled(&self) -> bool { - self.0 - .config - .server - .as_ref() - .map_or(false, |x| x.clickhouse.is_some()) + self.0.clickhouse_enabled } pub fn clickhouse(&self) -> GlobalResult { @@ -118,6 +164,14 @@ impl Pools { Ok(ch) } + pub fn fdb(&self) -> Result { + self.0.fdb.clone().ok_or(Error::MissingFdbPool) + } + + pub async fn sqlite(&self, key: Uuid) -> Result { + self.0.sqlite.get(key).await + } + #[tracing::instrument(skip_all)] async fn record_metrics_loop(self, token: CancellationToken) { let cancelled = token.cancelled(); diff --git a/packages/common/pools/src/prelude.rs b/packages/common/pools/src/prelude.rs index 504f81e6a5..02035c454b 100644 --- a/packages/common/pools/src/prelude.rs +++ b/packages/common/pools/src/prelude.rs @@ -4,6 +4,7 @@ pub use redis; pub use sqlx; pub use crate::{ - ClickHousePool, CrdbPool, NatsPool, RedisPool, __sql_query, __sql_query_as, __sql_query_as_raw, - sql_execute, sql_fetch, sql_fetch_all, sql_fetch_many, sql_fetch_one, sql_fetch_optional, + ClickHousePool, CrdbPool, FdbPool, NatsPool, RedisPool, SqlitePool, __sql_query, + __sql_query_as, __sql_query_as_raw, sql_execute, sql_fetch, sql_fetch_all, sql_fetch_many, + sql_fetch_one, sql_fetch_optional, }; diff --git a/packages/common/pools/src/utils/sql_query_macros.rs b/packages/common/pools/src/utils/sql_query_macros.rs index 3b6926b3df..6fa87f1000 100644 --- a/packages/common/pools/src/utils/sql_query_macros.rs +++ b/packages/common/pools/src/utils/sql_query_macros.rs @@ -90,7 +90,7 @@ macro_rules! __sql_query_metrics_finish { /// than make 4 RTT queries over 8 existing connections. #[macro_export] macro_rules! __sql_acquire { - ($ctx:expr, $crdb:expr) => {{ + ($ctx:expr, $driver:expr) => {{ let location = concat!(file!(), ":", line!(), ":", column!()); let mut tries = 0; @@ -98,7 +98,7 @@ macro_rules! __sql_acquire { tries += 1; // Attempt to use an existing connection - if let Some(conn) = $crdb.try_acquire() { + if let Some(conn) = $driver.try_acquire() { break (conn, "try_acquire"); } else { // Check if we can create a new connection @@ -107,7 +107,7 @@ macro_rules! __sql_acquire { .is_ok() { // Create a new connection - break ($crdb.acquire().await?, "acquire"); + break ($driver.acquire().await?, "acquire"); } else { // TODO: Backoff tokio::time::sleep(std::time::Duration::from_millis(1)).await; @@ -128,7 +128,7 @@ macro_rules! __sql_acquire { #[macro_export] macro_rules! __sql_query { - ([$ctx:expr, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { + ([$ctx:expr, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { async { use sqlx::Acquire; @@ -139,7 +139,7 @@ macro_rules! __sql_query { // Acquire connection $crate::__sql_query_metrics_acquire!(_acquire); - let crdb = $crdb; + let crdb = $driver; let mut conn = $crate::__sql_acquire!($ctx, crdb); // Execute query @@ -173,7 +173,7 @@ macro_rules! __sql_query { #[macro_export] macro_rules! __sql_query_as { - ([$ctx:expr, $rv:ty, $action:ident, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { + ([$ctx:expr, $rv:ty, $action:ident, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { async { use sqlx::Acquire; @@ -184,7 +184,7 @@ macro_rules! __sql_query_as { // Acquire connection $crate::__sql_query_metrics_acquire!(_acquire); - let crdb = $crdb; + let crdb = $driver; let mut conn = $crate::__sql_acquire!($ctx, crdb); // Execute query @@ -220,13 +220,13 @@ macro_rules! __sql_query_as { /// Used for the `fetch` function. #[macro_export] macro_rules! __sql_query_as_raw { - ([$ctx:expr, $rv:ty, $action:ident, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { + ([$ctx:expr, $rv:ty, $action:ident, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { // We can't record metrics for this because we can't move the `await` in to this macro sqlx::query_as::<_, $rv>($crate::__opt_indoc!($sql)) $( .bind($bind) )* - .$action($crdb) + .$action($driver) }; // TODO: This doesn't work with `fetch` ([$ctx:expr, $rv:ty, $action:ident] $sql:expr, $($bind:expr),* $(,)?) => { @@ -244,18 +244,18 @@ macro_rules! sql_execute { #[macro_export] macro_rules! sql_fetch { - ([$ctx:expr, $rv:ty, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as_raw!([$ctx, $rv, fetch, $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as_raw!([$ctx, $rv, fetch, $driver] $sql, $($bind),*) }; } #[macro_export] macro_rules! sql_fetch_all { - ([$ctx:expr, $rv:ty, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_all, $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_all, $driver] $sql, $($bind),*) }; - ([$ctx:expr, $rv:ty, @tx $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_all, @tx $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, @tx $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_all, @tx $driver] $sql, $($bind),*) }; ([$ctx:expr, $rv:ty] $sql:expr, $($bind:expr),* $(,)?) => { __sql_query_as!([$ctx, $rv, fetch_all] $sql, $($bind),*) @@ -264,11 +264,11 @@ macro_rules! sql_fetch_all { #[macro_export] macro_rules! sql_fetch_many { - ([$ctx:expr, $rv:ty, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_many, $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_many, $driver] $sql, $($bind),*) }; - ([$ctx:expr, $rv:ty, @tx $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_many, @tx $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, @tx $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_many, @tx $driver] $sql, $($bind),*) }; ([$ctx:expr, $rv:ty] $sql:expr, $($bind:expr),* $(,)?) => { __sql_query_as!([$ctx, $rv, fetch_many] $sql, $($bind),*) @@ -277,11 +277,11 @@ macro_rules! sql_fetch_many { #[macro_export] macro_rules! sql_fetch_one { - ([$ctx:expr, $rv:ty, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_one, $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_one, $driver] $sql, $($bind),*) }; - ([$ctx:expr, $rv:ty, @tx $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_one, @tx $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, @tx $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_one, @tx $driver] $sql, $($bind),*) }; ([$ctx:expr, $rv:ty] $sql:expr, $($bind:expr),* $(,)?) => { __sql_query_as!([$ctx, $rv, fetch_one] $sql, $($bind),*) @@ -290,11 +290,11 @@ macro_rules! sql_fetch_one { #[macro_export] macro_rules! sql_fetch_optional { - ([$ctx:expr, $rv:ty, $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_optional, $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_optional, $driver] $sql, $($bind),*) }; - ([$ctx:expr, $rv:ty, @tx $crdb:expr] $sql:expr, $($bind:expr),* $(,)?) => { - __sql_query_as!([$ctx, $rv, fetch_optional, @tx $crdb] $sql, $($bind),*) + ([$ctx:expr, $rv:ty, @tx $driver:expr] $sql:expr, $($bind:expr),* $(,)?) => { + __sql_query_as!([$ctx, $rv, fetch_optional, @tx $driver] $sql, $($bind),*) }; ([$ctx:expr, $rv:ty] $sql:expr, $($bind:expr),* $(,)?) => { __sql_query_as!([$ctx, $rv, fetch_optional] $sql, $($bind),*) diff --git a/packages/common/sqlite-util/Cargo.toml b/packages/common/sqlite-util/Cargo.toml new file mode 100644 index 0000000000..e04f99b07d --- /dev/null +++ b/packages/common/sqlite-util/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "sqlite-util" +version.workspace = true +authors.workspace = true +license.workspace = true +edition.workspace = true + +[dependencies] +tokio.workspace = true + +[dependencies.sqlx] +workspace = true +features = [ + "runtime-tokio", + "migrate", + "sqlite", + "uuid", + "json", + "ipnetwork", + "derive", +] diff --git a/packages/common/sqlite-util/src/lib.rs b/packages/common/sqlite-util/src/lib.rs new file mode 100644 index 0000000000..821558bbb4 --- /dev/null +++ b/packages/common/sqlite-util/src/lib.rs @@ -0,0 +1,67 @@ +use std::{ + future::Future, + ops::{Deref, DerefMut}, +}; + +use sqlx::sqlite::SqliteQueryResult; +use sqlx::{Executor, SqliteConnection}; + +pub trait SqliteConnectionExt { + fn begin_immediate(&mut self) -> impl Future>; +} + +impl SqliteConnectionExt for SqliteConnection { + async fn begin_immediate(&mut self) -> sqlx::Result { + let conn = &mut *self; + + conn.execute("BEGIN IMMEDIATE;").await?; + + Ok(Transaction { + conn, + is_open: true, + }) + } +} + +pub struct Transaction<'c> { + conn: &'c mut SqliteConnection, + /// is the transaction open? + is_open: bool, +} + +impl<'c> Transaction<'c> { + pub async fn commit(mut self) -> sqlx::Result { + let res = self.conn.execute("COMMIT;").await; + + if res.is_ok() { + self.is_open = false; + } + + res + } +} + +impl<'c> Drop for Transaction<'c> { + fn drop(&mut self) { + if self.is_open { + let handle = tokio::runtime::Handle::current(); + handle.block_on(async move { + let _ = self.execute("ROLLBACK").await; + }); + } + } +} + +impl<'c> Deref for Transaction<'c> { + type Target = SqliteConnection; + + fn deref(&self) -> &Self::Target { + self.conn + } +} + +impl<'c> DerefMut for Transaction<'c> { + fn deref_mut(&mut self) -> &mut Self::Target { + self.conn + } +} diff --git a/packages/infra/client/actor-kv/Cargo.toml b/packages/infra/client/actor-kv/Cargo.toml index b4bf0da971..31ee67f63b 100644 --- a/packages/infra/client/actor-kv/Cargo.toml +++ b/packages/infra/client/actor-kv/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [dependencies] anyhow.workspace = true deno_core.workspace = true -foundationdb = {version = "0.9.1", features = [ "fdb-7_1", "embedded-fdb-include" ] } +foundationdb.workspace = true futures-util = { version = "0.3" } indexmap = { version = "2.0" } prost = "0.13.3" diff --git a/packages/infra/client/actor-kv/src/lib.rs b/packages/infra/client/actor-kv/src/lib.rs index 0f4037c0b8..dfc434d528 100644 --- a/packages/infra/client/actor-kv/src/lib.rs +++ b/packages/infra/client/actor-kv/src/lib.rs @@ -191,6 +191,8 @@ impl ActorKv { let list_range = list_range.clone(); async move { + // TODO: Make this use `RangeOption.limit` instead of the weird custom limit impl with an + // error // Get all sub keys in the key subspace let stream = tx.get_ranges_keyvalues_owned( fdb::RangeOption { diff --git a/packages/infra/client/actor-kv/src/utils.rs b/packages/infra/client/actor-kv/src/utils.rs index 5e1e7c7fed..10e1d5b508 100644 --- a/packages/infra/client/actor-kv/src/utils.rs +++ b/packages/infra/client/actor-kv/src/utils.rs @@ -11,7 +11,7 @@ use crate::{ }; pub trait TransactionExt { - /// Owned version of `Transaction.get_ranges`. + /// Owned version of `Transaction.get_ranges` (self is owned). fn get_ranges_owned<'a>( self, opt: fdb::RangeOption<'a>, @@ -22,7 +22,7 @@ pub trait TransactionExt { + Unpin + 'a; - /// Owned version of `Transaction.get_ranges_keyvalues`. + /// Owned version of `Transaction.get_ranges_keyvalues` (self is owned). fn get_ranges_keyvalues_owned<'a>( self, opt: fdb::RangeOption<'a>, diff --git a/packages/infra/client/isolate-v8-runner/Cargo.toml b/packages/infra/client/isolate-v8-runner/Cargo.toml index f3f1de66bf..b84150604f 100644 --- a/packages/infra/client/isolate-v8-runner/Cargo.toml +++ b/packages/infra/client/isolate-v8-runner/Cargo.toml @@ -13,18 +13,19 @@ path = "src/main.rs" anyhow.workspace = true deno_ast = "0.42.1" deno_core.workspace = true -foundationdb = {version = "0.9.1", features = [ "fdb-7_1", "embedded-fdb-include" ] } +fdb-util.workspace = true +foundationdb.workspace = true futures-util = { version = "0.3" } netif = "0.1.6" nix.workspace = true serde = { version = "1.0.195", features = ["derive"] } serde_json = "1.0.111" signal-hook = "0.3.17" -tokio.workspace = true tokio-tungstenite = "0.23.1" -tracing.workspace = true +tokio.workspace = true tracing-logfmt.workspace = true tracing-subscriber.workspace = true +tracing.workspace = true twox-hash = "1.6.3" uuid = { version = "1.6.1", features = ["v4"] } diff --git a/packages/infra/client/isolate-v8-runner/src/isolate.rs b/packages/infra/client/isolate-v8-runner/src/isolate.rs index f43acadbed..ef5d8020f8 100644 --- a/packages/infra/client/isolate-v8-runner/src/isolate.rs +++ b/packages/infra/client/isolate-v8-runner/src/isolate.rs @@ -141,7 +141,10 @@ pub async fn run_inner( tracing::info!(?actor_id, "starting isolate"); // Init KV store (create or open) - let mut kv = ActorKv::new(utils::fdb_handle(&config)?, actor_config.owner.clone()); + let mut kv = ActorKv::new( + fdb_util::handle(&config.fdb_cluster_path)?, + actor_config.owner.clone(), + ); kv.init().await?; tracing::info!(?actor_id, "isolate kv initialized"); @@ -563,8 +566,7 @@ mod tests { ]); // Start FDB network thread - let _network = unsafe { fdb::boot() }; - tokio::spawn(utils::fdb_health_check(config.clone())); + fdb_util::init(&config.fdb_cluster_path); // For receiving the terminate handle let (terminate_tx, _terminate_rx) = diff --git a/packages/infra/client/isolate-v8-runner/src/main.rs b/packages/infra/client/isolate-v8-runner/src/main.rs index 489baf3124..c3806333c9 100644 --- a/packages/infra/client/isolate-v8-runner/src/main.rs +++ b/packages/infra/client/isolate-v8-runner/src/main.rs @@ -10,7 +10,6 @@ use std::{ use anyhow::*; use deno_core::{v8_set_flags, JsRuntime}; use deno_runtime::worker::MainWorkerTerminateHandle; -use foundationdb as fdb; use futures_util::{stream::SplitStream, SinkExt, StreamExt}; use pegboard::protocol; use pegboard_actor_kv::ActorKv; @@ -61,8 +60,7 @@ async fn main() -> Result<()> { let config = serde_json::from_str::(&config_data)?; // Start FDB network thread - let _network = unsafe { fdb::boot() }; - tokio::spawn(utils::fdb_health_check(config.clone())); + fdb_util::init(&config.fdb_cluster_path); tracing::info!(pid=%std::process::id(), "starting"); @@ -292,7 +290,7 @@ async fn watch_thread( // Remove state if !persist_storage { - let db = match utils::fdb_handle(&config) { + let db = match fdb_util::handle(&config.fdb_cluster_path) { Ok(db) => db, Err(err) => { tracing::error!(?err, ?actor_id, "failed to create fdb handle"); diff --git a/packages/infra/client/isolate-v8-runner/src/utils.rs b/packages/infra/client/isolate-v8-runner/src/utils.rs index dd1bc596b0..6592296375 100644 --- a/packages/infra/client/isolate-v8-runner/src/utils.rs +++ b/packages/infra/client/isolate-v8-runner/src/utils.rs @@ -1,43 +1,3 @@ -use std::{result::Result::Ok, time::Duration}; - -use anyhow::*; -use foundationdb::{self as fdb, options::DatabaseOption}; -use pegboard_config::isolate_runner::Config; - -pub fn fdb_handle(config: &Config) -> Result { - let db = fdb::Database::from_path( - &config - .fdb_cluster_path - .to_str() - .context("bad fdb_cluster_path")? - .to_string(), - ) - .context("failed to create FDB database")?; - db.set_option(DatabaseOption::TransactionRetryLimit(10))?; - - Ok(db) -} - -pub async fn fdb_health_check(config: Config) -> Result<()> { - let db = fdb_handle(&config)?; - - loop { - match ::tokio::time::timeout( - Duration::from_secs(3), - db.run(|trx, _maybe_committed| async move { Ok(trx.get(b"", false).await?) }), - ) - .await - { - Ok(res) => { - res?; - } - Err(_) => tracing::error!("fdb missed ping"), - } - - ::tokio::time::sleep(Duration::from_secs(3)).await; - } -} - pub mod tokio { use anyhow::*; use deno_core::unsync::MaskFutureAsSend; diff --git a/packages/infra/client/manager/Cargo.toml b/packages/infra/client/manager/Cargo.toml index e9bba804dc..ad0bf9323c 100644 --- a/packages/infra/client/manager/Cargo.toml +++ b/packages/infra/client/manager/Cargo.toml @@ -28,15 +28,16 @@ rand_chacha = "0.3.1" reqwest = { version = "0.11", default-features = false, features = ["stream", "rustls-tls", "json"] } serde = { version = "1.0.195", features = ["derive"] } serde_json = "1.0.111" +sqlite-util.workspace = true sysinfo = "0.31.4" tempfile = "3.2" thiserror = "1.0" tokio = { workspace = true, default-features = false, features = ["fs", "process", "macros", "rt", "rt-multi-thread"] } tokio-tungstenite = "0.23.1" tokio-util = { version = "0.7", default-features = false, features = ["io-util"] } -tracing.workspace = true tracing-logfmt.workspace = true tracing-subscriber.workspace = true +tracing.workspace = true url = "2.4" uuid = { version = "1.6.1", features = ["v4"] } diff --git a/packages/infra/client/manager/src/ctx.rs b/packages/infra/client/manager/src/ctx.rs index 2ce6e71b71..fb20556985 100644 --- a/packages/infra/client/manager/src/ctx.rs +++ b/packages/infra/client/manager/src/ctx.rs @@ -18,6 +18,7 @@ use pegboard::{protocol, system_info::SystemInfo}; use pegboard_config::{ isolate_runner::Config as IsolateRunnerConfig, runner_protocol, Client, Config, }; +use sqlite_util::SqliteConnectionExt; use sqlx::{pool::PoolConnection, Sqlite, SqlitePool}; use tokio::{ fs, @@ -35,12 +36,8 @@ use url::Url; use uuid::Uuid; use crate::{ - actor::Actor, - event_sender::EventSender, - metrics, - pull_addr_handler::PullAddrHandler, - runner, - utils::{self, sql::SqliteConnectionExt}, + actor::Actor, event_sender::EventSender, metrics, pull_addr_handler::PullAddrHandler, runner, + utils, }; const PING_INTERVAL: Duration = Duration::from_secs(1); @@ -130,7 +127,7 @@ impl Ctx { // Fetch next idx let index = utils::sql::query(|| async { let mut conn = self.sql().await?; - let mut txn = conn.begin_immediate().await?; + let mut tx = conn.begin_immediate().await?; let (index,) = sqlx::query_as::<_, (i64,)>(indoc!( " @@ -139,7 +136,7 @@ impl Ctx { RETURNING last_event_idx ", )) - .fetch_one(&mut *txn) + .fetch_one(&mut *tx) .await?; sqlx::query(indoc!( @@ -155,10 +152,10 @@ impl Ctx { .bind(index) .bind(&event_json) .bind(utils::now()) - .execute(&mut *txn) + .execute(&mut *tx) .await?; - txn.commit().await?; + tx.commit().await?; Ok(index) }) diff --git a/packages/infra/client/manager/src/utils/sql.rs b/packages/infra/client/manager/src/utils/sql.rs index 1c98163145..04511b8d63 100644 --- a/packages/infra/client/manager/src/utils/sql.rs +++ b/packages/infra/client/manager/src/utils/sql.rs @@ -1,92 +1,29 @@ use std::{ future::Future, - ops::{Deref, DerefMut}, result::Result::{Err, Ok}, time::Duration, }; use anyhow::*; -use sqlx::sqlite::SqliteQueryResult; -use sqlx::{Executor, SqliteConnection}; use crate::metrics; const MAX_QUERY_RETRIES: usize = 16; const QUERY_RETRY: Duration = Duration::from_millis(500); -pub(crate) trait SqliteConnectionExt { - fn begin_immediate(&mut self) -> impl Future>; -} - -impl SqliteConnectionExt for SqliteConnection { - async fn begin_immediate(&mut self) -> sqlx::Result { - let conn = &mut *self; - - conn.execute("BEGIN IMMEDIATE;").await?; - - Ok(Transaction { - conn, - is_open: true, - }) - } -} - -pub(crate) struct Transaction<'c> { - conn: &'c mut SqliteConnection, - /// is the transaction open? - is_open: bool, -} - -impl<'c> Transaction<'c> { - pub(crate) async fn commit(mut self) -> sqlx::Result { - let res = self.conn.execute("COMMIT;").await; - - if res.is_ok() { - self.is_open = false; - } - - res - } -} - -impl<'c> Drop for Transaction<'c> { - fn drop(&mut self) { - if self.is_open { - let handle = tokio::runtime::Handle::current(); - handle.block_on(async move { - let _ = self.execute("ROLLBACK").await; - }); - } - } -} - -impl<'c> Deref for Transaction<'c> { - type Target = SqliteConnection; - - fn deref(&self) -> &Self::Target { - self.conn - } -} - -impl<'c> DerefMut for Transaction<'c> { - fn deref_mut(&mut self) -> &mut Self::Target { - self.conn - } -} - /// Executes queries and explicitly handles retry errors. pub async fn query<'a, F, Fut, T>(mut cb: F) -> Result where F: FnMut() -> Fut, - Fut: std::future::Future> + 'a, + Fut: Future> + 'a, T: 'a, { let mut i = 0; loop { match cb().await { - std::result::Result::Ok(x) => return Ok(x), - std::result::Result::Err(err) => { + Ok(x) => return Ok(x), + Err(err) => { use sqlx::Error::*; metrics::SQL_ERROR diff --git a/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_gg_tls.sh b/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_gg_tls.sh index fcc649f218..a24a1d676b 100644 --- a/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_gg_tls.sh +++ b/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_gg_tls.sh @@ -70,7 +70,7 @@ EOF # Create systemd timer file cat << 'EOF' > /etc/systemd/system/rivet_fetch_gg_tls.timer [Unit] -Description=Runs TLS fetch every minute +Description=Runs TLS fetch every hour Requires=network-online.target After=network-online.target diff --git a/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_tunnel_tls.sh b/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_tunnel_tls.sh index e8be783ab3..467e4eef8a 100644 --- a/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_tunnel_tls.sh +++ b/packages/services/cluster/src/workflows/server/install/install_scripts/files/rivet_fetch_tunnel_tls.sh @@ -68,7 +68,7 @@ EOF # Create systemd timer file cat << 'EOF' > /etc/systemd/system/rivet_fetch_tunnel_tls.timer [Unit] -Description=Runs TLS fetch every minute +Description=Runs TLS fetch every hour Requires=network-online.target After=network-online.target diff --git a/site/src/content/docs/self-hosting/server-spec.json b/site/src/content/docs/self-hosting/server-spec.json index e653bb2fd8..951ae5c006 100644 --- a/site/src/content/docs/self-hosting/server-spec.json +++ b/site/src/content/docs/self-hosting/server-spec.json @@ -37,6 +37,16 @@ } ] }, + "fdb": { + "default": { + "connection": "fdb:fdb@127.0.0.1:4500" + }, + "allOf": [ + { + "$ref": "#/definitions/Fdb" + } + ] + }, "hcaptcha": { "default": null, "allOf": [ @@ -916,6 +926,18 @@ "cloudflare" ] }, + "Fdb": { + "type": "object", + "required": [ + "connection" + ], + "properties": { + "connection": { + "type": "string" + } + }, + "additionalProperties": false + }, "FirewallRule": { "type": "object", "required": [