import type { ExtractedAction } from './hydrant';
import { withHandle, type RepoHandle } from './identity';
import type { UrlEntry } from './types';
import { buildUrlsetXml, measureXmlBytes } from './xml';
// static urls have no record behind them
export type Change =
| ExtractedAction
| { type: 'upsert'; uri?: undefined; did?: undefined; claim?: undefined; url: string; entry: UrlEntry };
export interface ShardLimits {
maxUrls: number;
maxBytes: number;
}
// d1 allows 1000 queries a run and 100 params a query, so every batch goes
// over as json arrays through json_each instead of a statement per row
const CHUNK_BYTES = 512 * 1024;
const EMPTY_SHARD_BYTES = measureXmlBytes(buildUrlsetXml([]));
// the newline that joins it to the next url is counted too
export function urlBytes(entry: UrlEntry): number {
return measureXmlBytes(buildUrlsetXml([entry])) - EMPTY_SHARD_BYTES + 1;
}
function didOfUri(uri: string): string {
return uri.slice('at://'.length).split('/')[0];
}
function chunked(items: unknown[]): string[] {
const out: string[] = [];
let chunk: string[] = [];
let size = 0;
for (const item of items) {
const json = JSON.stringify(item);
if (size + json.length > CHUNK_BYTES && chunk.length > 0) {
out.push(`[${chunk.join(',')}]`);
chunk = [];
size = 0;
}
chunk.push(json);
size += json.length + 1;
}
if (chunk.length > 0) out.push(`[${chunk.join(',')}]`);
return out;
}
async function lookup<T>(db: D1Database, sql: string, keys: string[]): Promise<T[]> {
const unique = [...new Set(keys)];
const results = await Promise.all(chunked(unique).map((json) => db.prepare(sql).bind(json).all<T>()));
return results.flatMap((res) => res.results);
}
export async function getState(db: D1Database, key: string): Promise<string | undefined> {
const row = await db.prepare('SELECT value FROM state WHERE key = ?1').bind(key).first<{ value: string }>();
return row?.value;
}
function setState(db: D1Database, key: string, value: string): D1PreparedStatement {
return db
.prepare('INSERT INTO state (key, value) VALUES (?1, ?2) ON CONFLICT (key) DO UPDATE SET value = excluded.value')
.bind(key, value);
}
export async function putState(db: D1Database, key: string, value: string): Promise<void> {
await setState(db, key, value).run();
}
export async function countUrls(db: D1Database): Promise<number> {
const row = await db.prepare('SELECT count(*) AS n FROM urls').first<{ n: number }>();
return row?.n ?? 0;
}
// the newest shard takes new urls until it's full, older ones only change in place
async function activeShard(db: D1Database): Promise<{ shard: number; count: number; bytes: number }> {
const row = await db
.prepare(
`SELECT s.shard AS shard, count(u.url) AS count, coalesce(sum(u.bytes), 0) AS bytes
FROM (SELECT coalesce(max(shard), 1) AS shard FROM shards) s
LEFT JOIN urls u ON u.shard = s.shard`,
)
.first<{ shard: number; count: number; bytes: number }>();
return { shard: row?.shard ?? 1, count: row?.count ?? 0, bytes: EMPTY_SHARD_BYTES + (row?.bytes ?? 0) };
}
// applies one batch and moves the cursor in the same transaction, so a run
// that dies halfway picks up exactly where the last batch landed
export async function applyChanges(
db: D1Database,
changes: Change[],
limits: ShardLimits,
cursor: string | undefined,
now: Date,
): Promise<number> {
const uris = changes.flatMap((change) => (change.uri ? [change.uri] : []));
const recordUrl = new Map(
(await lookup<{ uri: string; url: string }>(
db,
'SELECT uri, url FROM records WHERE uri IN (SELECT value FROM json_each(?1))',
uris,
)).map((row) => [row.uri, row.url]),
);
const urlShard = new Map<string, number | null>(
(await lookup<{ url: string; shard: number }>(
db,
'SELECT url, shard FROM urls WHERE url IN (SELECT value FROM json_each(?1))',
[...changes.flatMap((change) => (change.type === 'upsert' ? [change.url] : [])), ...recordUrl.values()],
)).map((row) => [row.url, row.shard]),
);
const handleOf = new Map(
(await lookup<RepoHandle>(
db,
'SELECT did, handle FROM handles WHERE did IN (SELECT value FROM json_each(?1))',
changes.flatMap((change) => (change.type === 'upsert' && change.did ? [change.did] : [])),
)).map((row) => [row.did, row.handle]),
);
const active = await activeShard(db);
const upserts = new Map<
string,
UrlEntry & { did: string | null; record_did: string | null; claim: string | null; shard: number; bytes: number }
>();
const records = new Map<string, string | null>();
const dirty = new Set<number>();
let mutations = 0;
const remove = (url: string): boolean => {
const shard = urlShard.get(url);
if (shard == null) return false;
urlShard.set(url, null);
upserts.delete(url);
dirty.add(shard);
return true;
};
for (const change of changes) {
const previous = change.uri ? recordUrl.get(change.uri) : undefined;
if (change.type === 'delete') {
recordUrl.delete(change.uri);
records.set(change.uri, null);
if (previous && remove(previous)) mutations++;
continue;
}
if (change.uri && previous !== change.url) {
// a repo rename keeps its at-uri but changes its url
if (previous) remove(previous);
recordUrl.set(change.uri, change.url);
records.set(change.uri, change.url);
}
// sized as it renders, so the shard caps hold with the handle in
const handle = change.did ? handleOf.get(change.did) : null;
const bytes = urlBytes({ ...change.entry, loc: withHandle(change.url, change.did, handle) });
let shard = urlShard.get(change.url);
if (shard == null) {
if (active.count >= limits.maxUrls || active.bytes + bytes > limits.maxBytes) {
active.shard++;
active.count = 0;
active.bytes = EMPTY_SHARD_BYTES;
}
active.count++;
active.bytes += bytes;
shard = active.shard;
urlShard.set(change.url, shard);
}
upserts.set(change.url, {
...change.entry,
did: change.did ?? null,
record_did: change.uri ? didOfUri(change.uri) : null,
claim: change.claim ?? null,
shard,
bytes,
});
dirty.add(shard);
mutations++;
}
const deletedUrls = [...urlShard].filter(([, shard]) => shard === null).map(([url]) => url);
const statements: D1PreparedStatement[] = [
...chunked([...upserts.values()]).map((json) =>
db
.prepare(
`INSERT INTO urls (url, shard, lastmod, changefreq, priority, bytes, did, record_did, claim)
SELECT value ->> 'loc', value ->> 'shard', value ->> 'lastmod', value ->> 'changefreq',
value ->> 'priority', value ->> 'bytes', value ->> 'did', value ->> 'record_did', value ->> 'claim'
FROM json_each(?1) WHERE true
ON CONFLICT (url) DO UPDATE SET shard = excluded.shard, lastmod = excluded.lastmod, changefreq = excluded.changefreq,
priority = excluded.priority, bytes = excluded.bytes, did = excluded.did, record_did = excluded.record_did,
claim = excluded.claim`,
)
.bind(json),
),
...chunked(deletedUrls).map((json) =>
db.prepare('DELETE FROM urls WHERE url IN (SELECT value FROM json_each(?1))').bind(json),
),
...chunked([...records].filter(([, url]) => url !== null)).map((json) =>
db
.prepare(
`INSERT INTO records (uri, url) SELECT value ->> 0, value ->> 1 FROM json_each(?1) WHERE true
ON CONFLICT (uri) DO UPDATE SET url = excluded.url`,
)
.bind(json),
),
...chunked([...records].filter(([, url]) => url === null).map(([uri]) => uri)).map((json) =>
db.prepare('DELETE FROM records WHERE uri IN (SELECT value FROM json_each(?1))').bind(json),
),
...chunked([...dirty]).map((json) =>
db
.prepare(
`INSERT INTO shards (shard, lastmod, dirty) SELECT value, ?2, 1 FROM json_each(?1) WHERE true
ON CONFLICT (shard) DO UPDATE SET dirty = dirty + 1`,
)
.bind(json, now.toISOString().slice(0, 10)),
),
];
if (cursor !== undefined) statements.push(setState(db, 'cursor', cursor));
if (statements.length > 0) await db.batch(statements);
return mutations;
}
export async function dirtyShards(db: D1Database): Promise<{ shard: number; dirty: number }[]> {
const res = await db.prepare('SELECT shard, dirty FROM shards WHERE dirty > 0 ORDER BY shard').all<{ shard: number; dirty: number }>();
return res.results;
}
export async function shardEntries(db: D1Database, shard: number): Promise<UrlEntry[]> {
const res = await db
.prepare(
`SELECT u.url AS loc, u.did, h.handle, u.lastmod, u.changefreq, u.priority
FROM urls u LEFT JOIN handles h ON h.did = u.did
WHERE u.shard = ?1 AND NOT EXISTS (SELECT 1 FROM inactive i WHERE i.did IN (u.did, u.record_did))
AND (u.claim IS NULL OR EXISTS (SELECT 1 FROM declarations d WHERE d.uri = u.claim AND d.repo = u.record_did))
ORDER BY u.url`,
)
.bind(shard)
.all<{
loc: string;
did: string | null;
handle: string | null;
lastmod: string;
changefreq: UrlEntry['changefreq'] | null;
priority: number | null;
}>();
return res.results.map((row) => ({
loc: withHandle(row.loc, row.did, row.handle),
lastmod: row.lastmod,
changefreq: row.changefreq ?? undefined,
priority: row.priority ?? undefined,
}));
}
// only clears the shard if nothing changed it while it was being rendered
export async function markRendered(db: D1Database, shard: number, dirty: number, lastmod: string): Promise<void> {
await db
.prepare('UPDATE shards SET dirty = dirty - ?2, lastmod = ?3 WHERE shard = ?1')
.bind(shard, dirty, lastmod)
.run();
}
export async function allShards(db: D1Database): Promise<{ shard: number; lastmod: string }[]> {
const res = await db.prepare('SELECT shard, lastmod FROM shards ORDER BY shard').all<{ shard: number; lastmod: string }>();
return res.results;
}
// an identity event can mean a new handle, so those get read again
export async function markStale(db: D1Database, dids: string[]): Promise<void> {
const statements = chunked([...new Set(dids)]).map((json) =>
db.prepare('UPDATE handles SET stale = 1 WHERE did IN (SELECT value FROM json_each(?1))').bind(json),
);
if (statements.length > 0) await db.batch(statements);
}
// owners with no handle read yet, or a stale one, in did order after `after`
export async function staleOwners(db: D1Database, after: string, limit: number): Promise<string[]> {
const res = await db
.prepare(
`SELECT DISTINCT u.did FROM urls u LEFT JOIN handles h ON h.did = u.did
WHERE u.did > ?1 AND (h.did IS NULL OR h.stale = 1)
ORDER BY u.did LIMIT ?2`,
)
.bind(after, limit)
.all<{ did: string }>();
return res.results.map((row) => row.did);
}
// a changed handle resizes its owner's urls and rerenders their shards
export async function saveHandles(db: D1Database, found: RepoHandle[]): Promise<void> {
const previous = new Map(
(await lookup<RepoHandle>(
db,
'SELECT did, handle FROM handles WHERE did IN (SELECT value FROM json_each(?1))',
found.map((repo) => repo.did),
)).map((row) => [row.did, row.handle]),
);
const changed = found.filter((repo) => repo.handle !== (previous.get(repo.did) ?? null));
// dids and handles are plain ascii, so characters are bytes
const growth = changed.map((repo) => [
repo.did,
(repo.handle ?? repo.did).length - (previous.get(repo.did) ?? repo.did).length,
]);
const statements: D1PreparedStatement[] = [
...chunked(found.map((repo) => [repo.did, repo.handle])).map((json) =>
db
.prepare(
`INSERT INTO handles (did, handle, stale) SELECT value ->> 0, value ->> 1, 0 FROM json_each(?1) WHERE true
ON CONFLICT (did) DO UPDATE SET handle = excluded.handle, stale = 0`,
)
.bind(json),
),
...chunked(growth).map((json) =>
db
.prepare('UPDATE urls SET bytes = bytes + (j.value ->> 1) FROM json_each(?1) AS j WHERE urls.did = j.value ->> 0')
.bind(json),
),
...chunked(changed.map((repo) => repo.did)).map((json) =>
db
.prepare(
`UPDATE shards SET dirty = dirty + 1
WHERE shard IN (SELECT shard FROM urls WHERE did IN (SELECT value FROM json_each(?1)))`,
)
.bind(json),
),
];
if (statements.length > 0) await db.batch(statements);
}
export interface AccountStatus {
did: string;
active: boolean;
status?: string;
}
export async function saveAccounts(db: D1Database, accounts: AccountStatus[]): Promise<void> {
const latest = new Map(accounts.map((account) => [account.did, account]));
if (latest.size === 0) return;
const before = new Map(
(await lookup<{ did: string; status: string | null }>(
db,
'SELECT did, status FROM inactive WHERE did IN (SELECT value FROM json_each(?1))',
[...latest.keys()],
)).map((row) => [row.did, row.status]),
);
const changed = [...latest.values()].filter((account) =>
account.active ? before.has(account.did) : before.get(account.did) !== (account.status ?? null),
);
const hidden = changed.filter((account) => !account.active);
const shown = changed.filter((account) => account.active).map((account) => account.did);
const statements: D1PreparedStatement[] = [
...chunked(hidden.map((account) => [account.did, account.status ?? null])).map((json) =>
db
.prepare(
`INSERT INTO inactive (did, status) SELECT value ->> 0, value ->> 1 FROM json_each(?1) WHERE true
ON CONFLICT (did) DO UPDATE SET status = excluded.status`,
)
.bind(json),
),
...chunked(shown).map((json) =>
db.prepare('DELETE FROM inactive WHERE did IN (SELECT value FROM json_each(?1))').bind(json),
),
...chunked(changed.map((account) => account.did)).map((json) =>
db
.prepare(
`UPDATE shards SET dirty = dirty + 1 WHERE shard IN (
SELECT shard FROM urls WHERE did IN (SELECT value FROM json_each(?1))
UNION SELECT shard FROM urls WHERE record_did IN (SELECT value FROM json_each(?1)))`,
)
.bind(json),
),
];
if (statements.length > 0) await db.batch(statements);
}
// a declaration can show or hide a repo url, and that shard has to rerender
export async function saveDeclarations(db: D1Database, declarations: { uri: string; repo: string | null }[]): Promise<void> {
const latest = new Map(declarations.map((declaration) => [declaration.uri, declaration.repo]));
if (latest.size === 0) return;
const kept = [...latest].filter(([, repo]) => repo !== null);
const gone = [...latest].filter(([, repo]) => repo === null).map(([uri]) => uri);
const statements: D1PreparedStatement[] = [
...chunked(kept).map((json) =>
db
.prepare(
`INSERT INTO declarations (uri, repo) SELECT value ->> 0, value ->> 1 FROM json_each(?1) WHERE true
ON CONFLICT (uri) DO UPDATE SET repo = excluded.repo`,
)
.bind(json),
),
...chunked(gone).map((json) =>
db.prepare('DELETE FROM declarations WHERE uri IN (SELECT value FROM json_each(?1))').bind(json),
),
...chunked([...latest.keys()]).map((json) =>
db
.prepare(
`UPDATE shards SET dirty = dirty + 1
WHERE shard IN (SELECT shard FROM urls WHERE claim IN (SELECT value FROM json_each(?1)))`,
)
.bind(json),
),
];
await db.batch(statements);
}
export async function markSynced(db: D1Database, now: Date): Promise<void> {
await setState(db, 'last_sync_at', now.toISOString()).run();
}
export async function clearDb(db: D1Database): Promise<void> {
await db.batch(['urls', 'records', 'handles', 'shards', 'state', 'inactive', 'declarations'].map((table) => db.prepare(`DELETE FROM ${table}`)));
}