Effect Calls Cache (#633)

* Use persistence module in load layer instead of storage

* The most hacky cache persistence wip

* Make the cache work for uni v4 locally

* Improve cache dump logic

* Track Effect Cache size

* Fix cache count retrieval

* Use TSV, await dump and restore, upload cache only on the clean start, support dev console sync cache

* Expose GraphQl endpoint in TUI

* Dynamically choose psql exec

* Allow -alpha sufixed versions

* Fix wrong assumption about schema existence

* Call psql with host, user and password

* Fixes

* More env variables

* Better cache restore and dump handling

* Don't dump hash on HS, allow cache minification via schema, fix psql for local dev

* Fix create effect schema internally

* Enable BigInt json serialization for effect output

* Use default host from express

* Remove debug logs from init

* Write cache once per batch

* Fix effect cache performance downsides

* Use postgres:17.5

* Fix tests

* Fixes after review

Author: DZakh

codegenerator/cli/npm/envio/index.d.ts

@@ -13,6 +13,7 @@ import type {
 } from "./src/Envio.gen.ts";
 
 import { schema as bigDecimalSchema } from "./src/bindings/BigDecimal.gen.ts";
+import { schema as bigintSchema } from "./src/bindings/BigInt.gen.ts";
 import * as Sury from "rescript-schema";
 
 type UnknownToOutput<T> = T extends Sury.Schema<unknown>
@@ -35,26 +36,6 @@ type UnknownToOutput<T> = T extends Sury.Schema<unknown>
     >
   : T;
 
-type UnknownToInput<T> = T extends Sury.Schema<unknown>
-  ? Sury.Input<T>
-  : T extends (...args: any[]) => any
-  ? T
-  : T extends unknown[]
-  ? { [K in keyof T]: UnknownToInput<T[K]> }
-  : T extends { [k in keyof T]: unknown }
-  ? Flatten<
-      {
-        [k in keyof T as HasUndefined<UnknownToInput<T[k]>> extends true
-          ? k
-          : never]?: UnknownToInput<T[k]>;
-      } & {
-        [k in keyof T as HasUndefined<UnknownToInput<T[k]>> extends true
-          ? never
-          : k]: UnknownToInput<T[k]>;
-      }
-    >
-  : T;
-
 type HasUndefined<T> = [T] extends [undefined]
   ? true
   : undefined extends T
@@ -92,7 +73,7 @@ type Flatten<T> = T extends object
 export function experimental_createEffect<
   IS,
   OS,
-  I = UnknownToInput<IS>,
+  I = UnknownToOutput<IS>,
   O = UnknownToOutput<OS>,
   // A hack to enforce that the inferred return type
   // matches the output schema type
@@ -121,7 +102,7 @@ export declare namespace S {
   export const boolean: typeof Sury.boolean;
   export const int32: typeof Sury.int32;
   export const number: typeof Sury.number;
-  export const bigint: typeof Sury.bigint;
+  export const bigint: typeof bigintSchema;
   export const never: typeof Sury.never;
   export const union: typeof Sury.union;
   export const object: typeof Sury.object;
@@ -130,6 +111,7 @@ export declare namespace S {
   // Don't expose recursive for now, since it's too advanced
   // export const recursive: typeof Sury.recursive;
   export const transform: typeof Sury.transform;
+  export const shape: typeof Sury.shape;
   export const refine: typeof Sury.refine;
   export const schema: typeof Sury.schema;
   export const record: typeof Sury.record;

codegenerator/cli/npm/envio/index.js

@@ -13,7 +13,7 @@ exports.S = {
   boolean: Sury.boolean,
   int32: Sury.int32,
   number: Sury.number,
-  bigint: Sury.bigint,
+  bigint: require("./src/bindings/BigInt.res.js").schema,
   never: Sury.never,
   union: Sury.union,
   object: Sury.object,
@@ -22,6 +22,7 @@ exports.S = {
   // Don't expose recursive for now, since it's too advanced
   // recursive: Sury.recursive,
   transform: Sury.transform,
+  shape: Sury.shape,
   refine: Sury.refine,
   schema: Sury.schema,
   record: Sury.record,

codegenerator/cli/npm/envio/src/Envio.res

@@ -42,6 +42,8 @@ let experimental_createEffect = (
   handler: effectArgs<'input> => promise<'output>,
 ) => {
   Prometheus.EffectCallsCount.set(~callsCount=0, ~effectName=options.name)
+  let outputSchema =
+    S.schema(_ => options.output)->(Utils.magic: S.t<S.t<'output>> => S.t<Internal.effectOutput>)
   {
     name: options.name,
     handler: handler->(
@@ -50,6 +52,28 @@ let experimental_createEffect = (
       >
     ),
     callsCount: 0,
-    cache: options.cache->Belt.Option.getWithDefault(false),
+    // This is the way to make the createEffect API
+    // work without the need for users to call S.schema themselves,
+    // but simply pass the desired object/tuple/etc.
+    // If they pass a schem, it'll also work.
+    input: S.schema(_ => options.input)->(
+      Utils.magic: S.t<S.t<'input>> => S.t<Internal.effectInput>
+    ),
+    output: outputSchema,
+    cache: switch options.cache {
+    | Some(true) =>
+      let itemSchema = S.schema((s): Internal.effectCacheItem => {
+        id: s.matches(S.string),
+        output: s.matches(outputSchema),
+      })
+      Some({
+        table: Internal.makeCacheTable(~effectName=options.name),
+        rowsSchema: S.array(itemSchema),
+        itemSchema,
+      })
+    | None
+    | Some(false) =>
+      None
+    },
   }->(Utils.magic: Internal.effect => effect<'input, 'output>)
 }

codegenerator/cli/npm/envio/src/Internal.res

@@ -194,12 +194,31 @@ type effectArgs = {
   context: effectContext,
   cacheKey: string,
 }
+type effectCacheItem = {id: string, output: effectOutput}
+type effectCacheMeta = {
+  itemSchema: S.t<effectCacheItem>,
+  rowsSchema: S.t<array<effectCacheItem>>,
+  table: Table.table,
+}
 type effect = {
   name: string,
   handler: effectArgs => promise<effectOutput>,
-  cache: bool,
+  cache: option<effectCacheMeta>,
+  output: S.t<effectOutput>,
+  input: S.t<effectInput>,
   mutable callsCount: int,
 }
+let cacheTablePrefix = "envio_effect_"
+let makeCacheTable = (~effectName) => {
+  Table.mkTable(
+    cacheTablePrefix ++ effectName,
+    ~fields=[
+      Table.mkField("id", Text, ~fieldSchema=S.string, ~isPrimaryKey=true),
+      Table.mkField("output", JsonB, ~fieldSchema=S.json(~validate=false)),
+    ],
+    ~compositeIndices=[],
+  )
+}
 
 @genType.import(("./Types.ts", "Invalid"))
 type noEventFilters

codegenerator/cli/npm/envio/src/Persistence.res

@@ -5,14 +5,12 @@
 // Currently there are quite many code spread across
 // DbFunctions, Db, Migrations, InMemoryStore modules which use codegen code directly.
 
-// The type reflects an effect cache table in the db
+// The type reflects an cache table in the db
 // It might be present even if the effect is not used in the application
-type effectCache = {
-  name: string,
+type effectCacheRecord = {
+  effectName: string,
   // Number of rows in the table
-  mutable size: int,
-  // Lazily attached table definition when effect is used in the application
-  mutable table: option<Table.table>,
+  mutable count: int,
 }
 
 type operator = [#">" | #"="]
@@ -28,7 +26,6 @@ type storage = {
     ~generalTables: array<Table.table>=?,
     ~enums: array<Internal.enumConfig<Internal.enum>>=?,
   ) => promise<unit>,
-  loadEffectCaches: unit => promise<array<effectCache>>,
   @raises("StorageError")
   loadByIdsOrThrow: 'item. (
     ~ids: array<string>,
@@ -50,14 +47,26 @@ type storage = {
     ~table: Table.table,
     ~itemSchema: S.t<'item>,
   ) => promise<unit>,
+  @raises("StorageError")
+  setEffectCacheOrThrow: (
+    ~effect: Internal.effect,
+    ~items: array<Internal.effectCacheItem>,
+    ~initialize: bool,
+  ) => promise<unit>,
+  // This is to download cache from the database to .envio/cache
+  dumpEffectCache: unit => promise<unit>,
+  // This is not good, but the function does two things:
+  // - Gets info about existing cache tables
+  // - if withUpload is true, it also populates the cache from .envio/cache to the database
+  restoreEffectCache: (~withUpload: bool) => promise<array<effectCacheRecord>>,
 }
 
 exception StorageError({message: string, reason: exn})
 
 type storageStatus =
   | Unknown
   | Initializing(promise<unit>)
-  | Ready({cleanRun: bool, effectCaches: dict<effectCache>})
+  | Ready({cleanRun: bool, cache: dict<effectCacheRecord>})
 
 type t = {
   userEntities: array<Internal.entityConfig>,
@@ -66,7 +75,6 @@ type t = {
   allEnums: array<Internal.enumConfig<Internal.enum>>,
   mutable storageStatus: storageStatus,
   storage: storage,
-  onStorageInitialize: option<unit => promise<unit>>,
 }
 
 let entityHistoryActionEnumConfig: Internal.enumConfig<EntityHistory.RowAction.t> = {
@@ -83,7 +91,6 @@ let make = (
   ~allEnums,
   ~staticTables,
   ~storage,
-  ~onStorageInitialize=?,
 ) => {
   let allEntities = userEntities->Js.Array2.concat([dcRegistryEntityConfig])
   let allEnums =
@@ -95,62 +102,70 @@ let make = (
     allEnums,
     storageStatus: Unknown,
     storage,
-    onStorageInitialize,
   }
 }
 
-let init = async (persistence, ~reset=false) => {
-  try {
-    let shouldRun = switch persistence.storageStatus {
-    | Unknown => true
-    | Initializing(promise) => {
-        await promise
-        reset
-      }
-    | Ready(_) => reset
-    }
-    if shouldRun {
-      let resolveRef = ref(%raw(`null`))
-      let promise = Promise.make((resolve, _) => {
-        resolveRef := resolve
-      })
-      persistence.storageStatus = Initializing(promise)
-      if reset || !(await persistence.storage.isInitialized()) {
-        let _ = await persistence.storage.initialize(
-          ~entities=persistence.allEntities,
-          ~generalTables=persistence.staticTables,
-          ~enums=persistence.allEnums,
-        )
+let init = {
+  let loadInitialCache = async (persistence, ~withUpload) => {
+    let effectCacheRecords = await persistence.storage.restoreEffectCache(~withUpload)
+    let cache = Js.Dict.empty()
+    effectCacheRecords->Js.Array2.forEach(record => {
+      Prometheus.EffectCacheCount.set(~count=record.count, ~effectName=record.effectName)
+      cache->Js.Dict.set(record.effectName, record)
+    })
+    cache
+  }
 
-        persistence.storageStatus = Ready({
-          cleanRun: true,
-          effectCaches: Js.Dict.empty(),
-        })
-        switch persistence.onStorageInitialize {
-        | Some(onStorageInitialize) => await onStorageInitialize()
-        | None => ()
+  async (persistence, ~reset=false) => {
+    try {
+      let shouldRun = switch persistence.storageStatus {
+      | Unknown => true
+      | Initializing(promise) => {
+          await promise
+          reset
         }
-      } else if (
-        // In case of a race condition,
-        // we want to set the initial status to Ready only once.
-        switch persistence.storageStatus {
-        | Initializing(_) => true
-        | _ => false
-        }
-      ) {
-        let effectCaches = Js.Dict.empty()
-        (await persistence.storage.loadEffectCaches())->Js.Array2.forEach(effectCache => {
-          effectCaches->Js.Dict.set(effectCache.name, effectCache)
-        })
-        persistence.storageStatus = Ready({
-          cleanRun: false,
-          effectCaches,
+      | Ready(_) => reset
+      }
+      if shouldRun {
+        let resolveRef = ref(%raw(`null`))
+        let promise = Promise.make((resolve, _) => {
+          resolveRef := resolve
         })
+        persistence.storageStatus = Initializing(promise)
+        if reset || !(await persistence.storage.isInitialized()) {
+          Logging.info(`Initializing the indexer storage...`)
+
+          await persistence.storage.initialize(
+            ~entities=persistence.allEntities,
+            ~generalTables=persistence.staticTables,
+            ~enums=persistence.allEnums,
+          )
+
+          Logging.info(`The indexer storage is ready. Uploading cache...`)
+          persistence.storageStatus = Ready({
+            cleanRun: true,
+            cache: await loadInitialCache(persistence, ~withUpload=true),
+          })
+        } else if (
+          // In case of a race condition,
+          // we want to set the initial status to Ready only once.
+          switch persistence.storageStatus {
+          | Initializing(_) => true
+          | _ => false
+          }
+        ) {
+          Logging.info(`The indexer storage is ready.`)
+          persistence.storageStatus = Ready({
+            cleanRun: false,
+            cache: await loadInitialCache(persistence, ~withUpload=false),
+          })
+        }
+        resolveRef.contents()
       }
-      resolveRef.contents()
+    } catch {
+    | exn =>
+      exn->ErrorHandling.mkLogAndRaise(~msg=`EE800: Failed to initialize the indexer storage.`)
     }
-  } catch {
-  | exn => exn->ErrorHandling.mkLogAndRaise(~msg=`EE800: Failed to initialize the indexer storage.`)
   }
 }
 
@@ -162,3 +177,27 @@ let getInitializedStorageOrThrow = persistence => {
   | Ready(_) => persistence.storage
   }
 }
+
+let setEffectCacheOrThrow = async (persistence, ~effect: Internal.effect, ~items) => {
+  switch persistence.storageStatus {
+  | Unknown
+  | Initializing(_) =>
+    Js.Exn.raiseError(`Failed to access the indexer storage. The Persistence layer is not initialized.`)
+  | Ready({cache}) => {
+      let storage = persistence.storage
+      let effectName = effect.name
+      let effectCacheRecord = switch cache->Utils.Dict.dangerouslyGetNonOption(effectName) {
+      | Some(c) => c
+      | None => {
+          let c = {effectName, count: 0}
+          cache->Js.Dict.set(effectName, c)
+          c
+        }
+      }
+      let initialize = effectCacheRecord.count === 0
+      await storage.setEffectCacheOrThrow(~effect, ~items, ~initialize)
+      effectCacheRecord.count = effectCacheRecord.count + items->Js.Array2.length
+      Prometheus.EffectCacheCount.set(~count=effectCacheRecord.count, ~effectName)
+    }
+  }
+}