diff --git a/app/api/rpc/[command]/push.ts b/app/api/rpc/[command]/push.ts index cbd4b375..7c257ef3 100644 --- a/app/api/rpc/[command]/push.ts +++ b/app/api/rpc/[command]/push.ts @@ -1,15 +1,13 @@ -import { serverMutationContext } from "src/replicache/serverMutationContext"; import { mutations } from "src/replicache/mutations"; -import { eq } from "drizzle-orm"; +import { eq, sql } from "drizzle-orm"; import { permission_token_rights, replicache_clients } from "drizzle/schema"; import { getClientGroup } from "src/replicache/utils"; import { makeRoute } from "../lib"; import { z } from "zod"; import type { Env } from "./route"; +import { cachedServerMutationContext } from "src/replicache/cachedServerMutationContext"; import { drizzle } from "drizzle-orm/node-postgres"; - import { Pool } from "pg"; - import { attachDatabasePool } from "@vercel/functions"; import { DbPool } from "@vercel/functions/db-connections"; @@ -57,6 +55,10 @@ const pool = new Pool({ // Attach the pool to ensure idle connections close before suspension attachDatabasePool(pool as DbPool); +import { Lock } from "src/utils/lock"; + +let locks = new Map(); + export const push = makeRoute({ route: "push", input: z.object({ @@ -73,8 +75,13 @@ export const push = makeRoute({ let client = await pool.connect(); const db = drizzle(client); - let channel = supabase.channel(`rootEntity:${rootEntity}`); + let lock = locks.get(token.id); + if (!lock) { + lock = new Lock(); + locks.set(token.id, lock); + } + let release = await lock.lock(); try { await db.transaction(async (tx) => { let clientGroup = await getClientGroup(tx, pushRequest.clientGroupID); @@ -82,6 +89,14 @@ export const push = makeRoute({ .select() .from(permission_token_rights) .where(eq(permission_token_rights.token, token.id)); + let { getContext, flush } = cachedServerMutationContext( + tx, + token.id, + token_rights, + ); + + let lastMutations = new Map(); + console.log(pushRequest.mutations.map((m) => m.name)); for (let mutation of pushRequest.mutations) { let lastMutationID = clientGroup[mutation.clientID] || 0; if (mutation.id <= lastMutationID) continue; @@ -91,10 +106,8 @@ export const push = makeRoute({ continue; } try { - await mutations[name]( - mutation.args as any, - serverMutationContext(tx, token.id, token_rights), - ); + let ctx = getContext(mutation.clientID, mutation.id); + await mutations[name](mutation.args as any, ctx); } catch (e) { console.log( `Error occured while running mutation: ${name}`, @@ -102,18 +115,25 @@ export const push = makeRoute({ JSON.stringify(mutation, null, 2), ); } + lastMutations.set(mutation.clientID, mutation.id); + } + + let lastMutationIdsUpdate = Array.from(lastMutations.entries()).map( + (entries) => ({ + client_group: pushRequest.clientGroupID, + client_id: entries[0], + last_mutation: entries[1], + }), + ); + if (lastMutationIdsUpdate.length > 0) await tx .insert(replicache_clients) - .values({ - client_group: pushRequest.clientGroupID, - client_id: mutation.clientID, - last_mutation: mutation.id, - }) + .values(lastMutationIdsUpdate) .onConflictDoUpdate({ target: replicache_clients.client_id, - set: { last_mutation: mutation.id }, + set: { last_mutation: sql`excluded.last_mutation` }, }); - } + await flush(); }); await channel.send({ @@ -121,8 +141,11 @@ export const push = makeRoute({ event: "poke", payload: { message: "poke" }, }); + } catch (e) { + console.log(e); } finally { client.release(); + release(); supabase.removeChannel(channel); return { result: undefined } as const; } diff --git a/src/replicache/cachedServerMutationContext.ts b/src/replicache/cachedServerMutationContext.ts new file mode 100644 index 00000000..525b4778 --- /dev/null +++ b/src/replicache/cachedServerMutationContext.ts @@ -0,0 +1,252 @@ +import { PgTransaction } from "drizzle-orm/pg-core"; +import { Fact, PermissionToken } from "."; +import { MutationContext } from "./mutations"; +import { supabaseServerClient } from "supabase/serverClient"; +import { entities, facts } from "drizzle/schema"; +import * as driz from "drizzle-orm"; +import { Attribute, Attributes, FilterAttributes } from "./attributes"; +import { v7 } from "uuid"; +import * as base64 from "base64-js"; +import * as Y from "yjs"; +import { DeepReadonly } from "replicache"; + +type WriteCacheEntry = + | { type: "put"; fact: Fact } + | { type: "del"; fact: { id: string } }; +export function cachedServerMutationContext( + tx: PgTransaction, + permission_token_id: string, + token_rights: PermissionToken["permission_token_rights"], +) { + let writeCache: WriteCacheEntry[] = []; + let eavCache = new Map>[]>(); + let permissionsCache: { [key: string]: boolean } = {}; + let entitiesCache: { set: string; id: string }[] = []; + let deleteEntitiesCache: string[] = []; + let textAttributeWriteCache = {} as { + [entityAttribute: string]: { [clientID: string]: string }; + }; + + const scanIndex = { + async eav(entity: string, attribute: A) { + let cached = eavCache.get(`${entity}-${attribute}`) as DeepReadonly< + Fact + >[]; + let baseFacts: DeepReadonly>[]; + if (deleteEntitiesCache.includes(entity)) return []; + if (cached) baseFacts = cached; + else { + cached = (await tx + .select({ + id: facts.id, + data: facts.data, + entity: facts.entity, + attribute: facts.attribute, + }) + .from(facts) + .where( + driz.and( + driz.eq(facts.attribute, attribute), + driz.eq(facts.entity, entity), + ), + )) as DeepReadonly>[]; + } + cached = cached.filter( + (f) => + !writeCache.find((wc) => wc.type === "del" && wc.fact.id === f.id), + ); + let newlyWrittenFacts = writeCache.filter( + (f) => + f.type === "put" && + f.fact.attribute === attribute && + f.fact.entity === entity, + ); + return [ + ...cached, + ...newlyWrittenFacts.map((f) => f.fact as Fact), + ].filter( + (f) => + !( + (f.data.type === "reference" || + f.data.type === "ordered-reference" || + f.data.type === "spatial-reference") && + deleteEntitiesCache.includes(f.data.value) + ), + ) as DeepReadonly>[]; + }, + }; + let getContext = (clientID: string, mutationID: number) => { + let ctx: MutationContext & { + checkPermission: (entity: string) => Promise; + } = { + scanIndex, + permission_token_id, + async runOnServer(cb) { + return cb({ supabase: supabaseServerClient }); + }, + async checkPermission(entity: string) { + if (deleteEntitiesCache.includes(entity)) return false; + let cachedEntity = entitiesCache.find((e) => e.id === entity); + if (cachedEntity) { + return !!token_rights.find( + (r) => r.entity_set === cachedEntity?.set && r.write === true, + ); + } + if (permissionsCache[entity] !== undefined) + return permissionsCache[entity]; + let [permission_set] = await tx + .select({ entity_set: entities.set }) + .from(entities) + .where(driz.eq(entities.id, entity)); + let hasPermission = + !!permission_set && + !!token_rights.find( + (r) => + r.entity_set === permission_set.entity_set && r.write == true, + ); + permissionsCache[entity] = hasPermission; + return hasPermission; + }, + async runOnClient(_cb) {}, + async createEntity({ entityID, permission_set }) { + if ( + !token_rights.find( + (r) => r.entity_set === permission_set && r.write === true, + ) + ) { + return false; + } + if (!entitiesCache.find((e) => e.id === entityID)) + entitiesCache.push({ set: permission_set, id: entityID }); + deleteEntitiesCache = deleteEntitiesCache.filter((e) => e === entityID); + return true; + }, + async deleteEntity(entity) { + if (!(await this.checkPermission(entity))) return; + deleteEntitiesCache.push(entity); + entitiesCache = entitiesCache.filter((e) => e.id === entity); + }, + async assertFact(f) { + if (!f.entity) return; + let attribute = Attributes[f.attribute as Attribute]; + if (!attribute) return; + let id = f.id || v7(); + let data = { ...f.data }; + if (!(await this.checkPermission(f.entity))) return; + if (attribute.cardinality === "one") { + let existingFact = await scanIndex.eav(f.entity, f.attribute); + if (existingFact[0]) { + id = existingFact[0].id; + if (attribute.type === "text") { + let c = + textAttributeWriteCache[`${f.entity}-${f.attribute}`] || {}; + textAttributeWriteCache[`${f.entity}-${f.attribute}`] = { + ...c, + [clientID]: ( + data as Fact>["data"] + ).value, + }; + } + } + } + writeCache = writeCache.filter((f) => f.fact.id !== id); + writeCache.push({ + type: "put", + fact: { + id: id, + entity: f.entity, + data: data, + attribute: f.attribute, + }, + }); + }, + async retractFact(factID) { + writeCache = writeCache.filter((f) => f.fact.id !== factID); + writeCache.push({ type: "del", fact: { id: factID } }); + }, + }; + return ctx; + }; + let flush = async () => { + if (entitiesCache.length > 0) + await tx + .insert(entities) + .values(entitiesCache.map((e) => ({ set: e.set, id: e.id }))); + let factWrites = writeCache.flatMap((f) => + f.type === "del" ? [] : [f.fact], + ); + if (factWrites.length > 0) + await tx + .insert(facts) + .values( + await Promise.all( + factWrites.map(async (f) => { + let attribute = Attributes[f.attribute as Attribute]; + let data = f.data; + if ( + attribute.type === "text" && + attribute.cardinality === "one" + ) { + let values = Object.values( + textAttributeWriteCache[`${f.entity}-${f.attribute}`] || {}, + ); + if (values.length > 0) { + let existingFact = await scanIndex.eav(f.entity, f.attribute); + if (existingFact[0]) values.push(existingFact[0].data.value); + let updateBytes = Y.mergeUpdates( + values.map((v) => base64.toByteArray(v)), + ); + data.value = base64.fromByteArray(updateBytes); + } + } + + return { + id: f.id, + entity: f.entity, + data: driz.sql`${data}::jsonb`, + attribute: f.attribute, + }; + }), + ), + ) + .onConflictDoUpdate({ + target: facts.id, + set: { data: driz.sql`excluded.data` }, + }); + if (deleteEntitiesCache.length > 0) + await tx + .delete(entities) + .where(driz.inArray(entities.id, deleteEntitiesCache)); + let factDeletes = writeCache.flatMap((f) => + f.type === "put" ? [] : [f.fact.id], + ); + if (factDeletes.length > 0 || deleteEntitiesCache.length > 0) { + const conditions = []; + if (factDeletes.length > 0) { + conditions.push(driz.inArray(facts.id, factDeletes)); + } + if (deleteEntitiesCache.length > 0) { + conditions.push( + driz.and( + driz.sql`(data->>'type' = 'ordered-reference' or data->>'type' = 'reference' or data->>'type' = 'spatial-reference')`, + driz.inArray(driz.sql`data->>'value'`, deleteEntitiesCache), + ), + ); + } + if (conditions.length > 0) { + await tx.delete(facts).where(driz.or(...conditions)); + } + } + + writeCache = []; + eavCache.clear(); + permissionsCache = {}; + entitiesCache = []; + deleteEntitiesCache = []; + }; + + return { + getContext, + flush, + }; +} diff --git a/src/replicache/index.tsx b/src/replicache/index.tsx index 900d8717..321c137b 100644 --- a/src/replicache/index.tsx +++ b/src/replicache/index.tsx @@ -130,18 +130,24 @@ export function ReplicacheProvider(props: { ) as ReplicacheMutators, licenseKey: "l381074b8d5224dabaef869802421225a", pusher: async (pushRequest) => { + const batchSize = 250; let smolpushRequest = { ...pushRequest, - mutations: pushRequest.mutations.slice(0, 250), + mutations: pushRequest.mutations.slice(0, batchSize), } as PushRequest; + let response = ( + await callRPC("push", { + pushRequest: smolpushRequest, + token: props.token, + rootEntity: props.name, + }) + ).result; + if (pushRequest.mutations.length > batchSize) + setTimeout(() => { + newRep.push(); + }, 50); return { - response: ( - await callRPC("push", { - pushRequest: smolpushRequest, - token: props.token, - rootEntity: props.name, - }) - ).result, + response, httpRequestInfo: { errorMessage: "", httpStatusCode: 200 }, }; }, @@ -158,10 +164,10 @@ export function ReplicacheProvider(props: { name: props.name, indexes: { eav: { jsonPointer: "/indexes/eav", allowEmpty: true }, - aev: { jsonPointer: "/indexes/aev", allowEmpty: true }, vae: { jsonPointer: "/indexes/vae", allowEmpty: true }, }, }); + setRep(newRep); let channel: RealtimeChannel | null = null; if (!props.disablePull) { diff --git a/src/utils/lock.ts b/src/utils/lock.ts new file mode 100644 index 00000000..5c303a89 --- /dev/null +++ b/src/utils/lock.ts @@ -0,0 +1,30 @@ +// Taken from https://github.com/rocicorp/lock/blob/main/src/lock.ts +export class Lock { + private _lockP: Promise | null = null; + + async lock(): Promise<() => void> { + const previous = this._lockP; + const { promise, resolve } = resolver(); + this._lockP = promise; + await previous; + return resolve; + } + async withLock(f: () => Promise) { + let release = await this.lock(); + try { + return await f(); + } finally { + release(); + } + } +} + +export function resolver() { + let resolve!: (v: void) => void; + let reject!: () => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +}