@atproto/pds 0.4.117 → 0.4.118

package/src/scripts/sequencer-recovery/recoverer.ts

@@ -0,0 +1,289 @@
+import { rmIfExists } from '@atproto/common'
+import { Secp256k1Keypair } from '@atproto/crypto'
+import {
+  BlockMap,
+  CidSet,
+  CommitData,
+  WriteOpAction,
+  cborToLexRecord,
+  parseDataKey,
+  readCar,
+} from '@atproto/repo'
+import {
+  AccountManager,
+  AccountStatus,
+} from '../../account-manager/account-manager'
+import { ActorStore } from '../../actor-store/actor-store'
+import { ActorStoreTransactor } from '../../actor-store/actor-store-transactor'
+import { countAll } from '../../db'
+import {
+  PreparedWrite,
+  prepareCreate,
+  prepareDelete,
+  prepareUpdate,
+} from '../../repo'
+import { AccountEvt, CommitEvt, SeqEvt, Sequencer } from '../../sequencer'
+import { RecoveryDb } from './recovery-db'
+import { UserQueues } from './user-queues'
+
+export type RecovererContextNoDb = {
+  sequencer: Sequencer
+  accountManager: AccountManager
+  actorStore: ActorStore
+}
+
+export type RecovererContext = RecovererContextNoDb & {
+  recoveryDb: RecoveryDb
+}
+
+const PAGE_SIZE = 5000
+
+export class Recoverer {
+  queues: UserQueues
+  failed: Set<string>
+
+  constructor(
+    public ctx: RecovererContext,
+    opts: { concurrency: number },
+  ) {
+    this.queues = new UserQueues(opts.concurrency)
+    this.failed = new Set()
+  }
+
+  async run(startCursor = 0) {
+    const failed = await this.ctx.recoveryDb.db
+      .selectFrom('failed')
+      .select('did')
+      .execute()
+    for (const row of failed) {
+      this.failed.add(row.did)
+    }
+
+    const totalRes = await this.ctx.sequencer.db.db
+      .selectFrom('repo_seq')
+      .select(countAll.as('count'))
+      .executeTakeFirstOrThrow()
+    const totalEvts = totalRes.count
+    let completed = 0
+
+    let cursor: number | undefined = startCursor
+    while (cursor !== undefined) {
+      const page = await this.ctx.sequencer.requestSeqRange({
+        earliestSeq: cursor,
+        limit: PAGE_SIZE,
+      })
+      page.forEach((evt) => this.processEvent(evt))
+      cursor = page.at(-1)?.seq
+
+      await this.queues.onEmpty()
+
+      completed += PAGE_SIZE
+      const percentComplete = (completed / totalEvts) * 100
+      console.log(`${percentComplete.toFixed(2)}% - ${cursor}`)
+    }
+
+    await this.queues.processAll()
+  }
+
+  async processAll() {
+    await this.queues.processAll()
+  }
+
+  async destroy() {
+    await this.queues.destroy()
+  }
+
+  processEvent(evt: SeqEvt) {
+    const did = didFromEvt(evt)
+    if (!did) {
+      return
+    }
+    this.queues.addToUser(did, async () => {
+      if (this.failed.has(did)) {
+        return
+      }
+      await processSeqEvt(this.ctx, evt).catch(async (err) => {
+        this.failed.add(did)
+        await trackFailure(this.ctx.recoveryDb, did, err)
+      })
+    })
+  }
+}
+
+export const processSeqEvt = async (ctx: RecovererContext, evt: SeqEvt) => {
+  // only need to process commits & tombstones
+  if (evt.type === 'account') {
+    await processAccountEvt(ctx, evt.evt)
+  }
+  if (evt.type === 'commit') {
+    await processCommit(ctx, evt.evt).catch()
+  }
+}
+
+const processCommit = async (ctx: RecovererContext, evt: CommitEvt) => {
+  const did = evt.repo
+  const { writes, blocks } = await parseCommitEvt(evt)
+  if (evt.since === null) {
+    const actorExists = await ctx.actorStore.exists(did)
+    if (!actorExists) {
+      await processRepoCreation(ctx, evt, writes, blocks)
+      return
+    }
+  }
+  await ctx.actorStore.transact(did, async (actorTxn) => {
+    const root = await actorTxn.repo.storage.getRootDetailed()
+    if (root.rev >= evt.rev) {
+      return
+    }
+    const commit = await actorTxn.repo.formatCommit(writes)
+    commit.newBlocks = blocks
+    commit.cid = evt.commit
+    commit.rev = evt.rev
+    await Promise.all([
+      actorTxn.repo.storage.applyCommit(commit),
+      actorTxn.repo.indexWrites(writes, commit.rev),
+      trackBlobs(actorTxn, writes),
+    ])
+  })
+}
+
+const processRepoCreation = async (
+  ctx: RecovererContext,
+  evt: CommitEvt,
+  writes: PreparedWrite[],
+  blocks: BlockMap,
+) => {
+  const did = evt.repo
+  const keypair = await Secp256k1Keypair.create({ exportable: true })
+  await ctx.actorStore.create(did, keypair)
+  const commit: CommitData = {
+    cid: evt.commit,
+    rev: evt.rev,
+    since: evt.since,
+    prev: null,
+    newBlocks: blocks,
+    relevantBlocks: new BlockMap(),
+    removedCids: new CidSet(),
+  }
+  await ctx.actorStore.transact(did, (actorTxn) =>
+    Promise.all([
+      actorTxn.repo.storage.applyCommit(commit, true),
+      actorTxn.repo.indexWrites(writes, commit.rev),
+      actorTxn.repo.blob.processWriteBlobs(commit.rev, writes),
+    ]),
+  )
+  await trackNewAccount(ctx.recoveryDb, did)
+}
+
+const processAccountEvt = async (ctx: RecovererContext, evt: AccountEvt) => {
+  // do not need to process deactivation/takedowns because we backup account DB as well
+  if (evt.status !== AccountStatus.Deleted) {
+    return
+  }
+  const { directory } = await ctx.actorStore.getLocation(evt.did)
+  await rmIfExists(directory, true)
+  await ctx.accountManager.deleteAccount(evt.did)
+}
+
+const trackBlobs = async (
+  store: ActorStoreTransactor,
+  writes: PreparedWrite[],
+) => {
+  await store.repo.blob.deleteDereferencedBlobs(writes)
+
+  for (const write of writes) {
+    if (
+      write.action === WriteOpAction.Create ||
+      write.action === WriteOpAction.Update
+    ) {
+      for (const blob of write.blobs) {
+        await store.repo.blob.insertBlobMetadata(blob)
+        await store.repo.blob.associateBlob(blob, write.uri)
+      }
+    }
+  }
+}
+
+const trackFailure = async (
+  recoveryDb: RecoveryDb,
+  did: string,
+  err: unknown,
+) => {
+  await recoveryDb.db
+    .insertInto('failed')
+    .values({
+      did,
+      error: err?.toString(),
+      fixed: 0,
+    })
+    .onConflict((oc) => oc.doNothing())
+    .execute()
+}
+
+const trackNewAccount = async (recoveryDb: RecoveryDb, did: string) => {
+  await recoveryDb.db
+    .insertInto('new_account')
+    .values({
+      did,
+      published: 0,
+    })
+    .onConflict((oc) => oc.doNothing())
+    .execute()
+}
+
+const parseCommitEvt = async (
+  evt: CommitEvt,
+): Promise<{
+  writes: PreparedWrite[]
+  blocks: BlockMap
+}> => {
+  const did = evt.repo
+  const evtCar = await readCar(evt.blocks)
+  const writesUnfiltered = await Promise.all(
+    evt.ops.map(async (op) => {
+      const { collection, rkey } = parseDataKey(op.path)
+      if (op.action === 'delete') {
+        return prepareDelete({ did, collection, rkey })
+      }
+      if (!op.cid) return undefined
+      const recordBytes = evtCar.blocks.get(op.cid)
+      if (!recordBytes) return undefined
+      const record = cborToLexRecord(recordBytes)
+
+      if (op.action === 'create') {
+        return prepareCreate({
+          did,
+          collection,
+          rkey,
+          record,
+          validate: false,
+        })
+      } else {
+        return prepareUpdate({
+          did,
+          collection,
+          rkey,
+          record,
+          validate: false,
+        })
+      }
+    }),
+  )
+  const writes = writesUnfiltered.filter(
+    (w) => w !== undefined,
+  ) as PreparedWrite[]
+  return {
+    writes,
+    blocks: evtCar.blocks,
+  }
+}
+
+const didFromEvt = (evt: SeqEvt): string | null => {
+  if (evt.type === 'account') {
+    return evt.evt.did
+  } else if (evt.type === 'commit') {
+    return evt.evt.repo
+  } else {
+    return null
+  }
+}

package/src/scripts/sequencer-recovery/recovery-db.ts

@@ -0,0 +1,64 @@
+import path from 'node:path'
+import { Kysely } from 'kysely'
+import { Database, Migrator } from '../../db'
+
+export interface NewAccount {
+  did: string
+  published: 0 | 1
+}
+
+export interface Failed {
+  did: string
+  error: string | null
+  fixed: 0 | 1
+}
+
+export type RecoveryDbSchema = {
+  new_account: NewAccount
+  failed: Failed
+}
+
+export type RecoveryDb = Database<RecoveryDbSchema>
+
+export const getRecoveryDbFromSequencerLoc = (
+  sequencerLoc: string,
+): Promise<RecoveryDb> => {
+  const recoveryDbLoc = path.join(path.dirname(sequencerLoc), 'recovery.sqlite')
+  return getAndMigrateRecoveryDb(recoveryDbLoc)
+}
+
+export const getAndMigrateRecoveryDb = async (
+  location: string,
+  disableWalAutoCheckpoint = false,
+): Promise<RecoveryDb> => {
+  const pragmas: Record<string, string> = disableWalAutoCheckpoint
+    ? { wal_autocheckpoint: '0' }
+    : {}
+  const db = Database.sqlite(location, pragmas)
+  const migrator = new Migrator(db.db, migrations)
+  await migrator.migrateToLatestOrThrow()
+  return db
+}
+
+const migrations = {
+  '001': {
+    up: async (db: Kysely<unknown>) => {
+      await db.schema
+        .createTable('new_account')
+        .addColumn('did', 'varchar', (col) => col.primaryKey())
+        .addColumn('published', 'int2', (col) => col.notNull())
+        .execute()
+
+      await db.schema
+        .createTable('failed')
+        .addColumn('did', 'varchar', (col) => col.primaryKey())
+        .addColumn('error', 'varchar')
+        .addColumn('fixed', 'int2', (col) => col.notNull())
+        .execute()
+    },
+    down: async (db: Kysely<unknown>) => {
+      await db.schema.dropTable('new_account').execute()
+      await db.schema.dropTable('failed').execute()
+    },
+  },
+}

package/src/scripts/sequencer-recovery/repair-repos.ts

@@ -0,0 +1,48 @@
+import { parseRepoSeqRows } from '../../sequencer'
+import { rebuildRepo } from '../rebuild-repo'
+import {
+  RecovererContext,
+  RecovererContextNoDb,
+  processSeqEvt,
+} from './recoverer'
+import { getRecoveryDbFromSequencerLoc } from './recovery-db'
+
+export const repairRepos = async (ctx: RecovererContextNoDb) => {
+  const recoveryDb = await getRecoveryDbFromSequencerLoc(
+    ctx.sequencer.dbLocation,
+  )
+  const repairRes = await recoveryDb.db
+    .selectFrom('failed')
+    .select('did')
+    .where('failed.fixed', '=', 0)
+    .execute()
+  const dids = repairRes.map((row) => row.did)
+  let fixed = 0
+  for (const did of dids) {
+    await rebuildRepo(ctx, did, false)
+    await recoverFromSequencer({ ...ctx, recoveryDb }, did)
+    fixed++
+    console.log(`${fixed}/${dids.length}`)
+  }
+}
+
+const recoverFromSequencer = async (ctx: RecovererContext, did: string) => {
+  const didEvts = await ctx.sequencer.db.db
+    .selectFrom('repo_seq')
+    .selectAll()
+    .where('did', '=', did)
+    .orderBy('seq', 'asc')
+    .execute()
+  const seqEvts = parseRepoSeqRows(didEvts)
+  for (const evt of seqEvts) {
+    await processSeqEvt(ctx, evt)
+  }
+  await ctx.recoveryDb.db
+    .updateTable('failed')
+    .set({
+      fixed: 1,
+      error: null,
+    })
+    .where('did', '=', did)
+    .execute()
+}

package/src/scripts/sequencer-recovery/user-queues.ts

@@ -0,0 +1,41 @@
+import PQueue from 'p-queue'
+export class UserQueues {
+  main: PQueue
+  queues = new Map<string, PQueue>()
+
+  constructor(concurrency: number) {
+    this.main = new PQueue({ concurrency })
+  }
+
+  async addToUser(did: string, task: () => Promise<void>) {
+    if (this.main.isPaused) return
+    return this.main.add(() => {
+      return this.getQueue(did).add(task)
+    })
+  }
+
+  private getQueue(did: string) {
+    let queue = this.queues.get(did)
+    if (!queue) {
+      queue = new PQueue({ concurrency: 1 })
+      queue.once('idle', () => this.queues.delete(did))
+      this.queues.set(did, queue)
+    }
+    return queue
+  }
+
+  async onEmpty() {
+    await this.main.onEmpty()
+  }
+
+  async processAll() {
+    await this.main.onIdle()
+  }
+
+  async destroy() {
+    this.main.pause()
+    this.main.clear()
+    this.queues.forEach((q) => q.clear())
+    await this.processAll()
+  }
+}

package/src/scripts/util.ts

@@ -0,0 +1,7 @@
+export const parseIntArg = (arg: string): number => {
+  const parsed = parseInt(arg, 10)
+  if (isNaN(parsed)) {
+    throw new Error(`Invalid arg, expected number: ${arg}`)
+  }
+  return parsed
+}

package/src/sequencer/sequencer.ts

@@ -33,7 +33,7 @@ export class Sequencer extends (EventEmitter as new () => SequencerEmitter) {
   triesWithNoResults = 0
 
   constructor(
-    dbLocation: string,
+    public dbLocation: string,
     public crawlers: Crawlers,
     public lastSeen = 0,
     disableWalAutoCheckpoint = false,
@@ -126,45 +126,7 @@ export class Sequencer extends (EventEmitter as new () => SequencerEmitter) {
       return []
     }
 
-    const seqEvts: SeqEvt[] = []
-    for (const row of rows) {
-      // should never hit this because of WHERE clause
-      if (row.seq === null) {
-        continue
-      }
-      const evt = cborDecode(row.event)
-      if (row.eventType === 'append') {
-        seqEvts.push({
-          type: 'commit',
-          seq: row.seq,
-          time: row.sequencedAt,
-          evt: evt as CommitEvt,
-        })
-      } else if (row.eventType === 'sync') {
-        seqEvts.push({
-          type: 'sync',
-          seq: row.seq,
-          time: row.sequencedAt,
-          evt: evt as SyncEvt,
-        })
-      } else if (row.eventType === 'identity') {
-        seqEvts.push({
-          type: 'identity',
-          seq: row.seq,
-          time: row.sequencedAt,
-          evt: evt as IdentityEvt,
-        })
-      } else if (row.eventType === 'account') {
-        seqEvts.push({
-          type: 'account',
-          seq: row.seq,
-          time: row.sequencedAt,
-          evt: evt as AccountEvt,
-        })
-      }
-    }
-
-    return seqEvts
+    return parseRepoSeqRows(rows)
   }
 
   private async pollDb(): Promise<void> {
@@ -243,6 +205,47 @@ export class Sequencer extends (EventEmitter as new () => SequencerEmitter) {
   }
 }
 
+export const parseRepoSeqRows = (rows: RepoSeqEntry[]): SeqEvt[] => {
+  const seqEvts: SeqEvt[] = []
+  for (const row of rows) {
+    // should never hit this because of WHERE clause
+    if (row.seq === null) {
+      continue
+    }
+    const evt = cborDecode(row.event)
+    if (row.eventType === 'append') {
+      seqEvts.push({
+        type: 'commit',
+        seq: row.seq,
+        time: row.sequencedAt,
+        evt: evt as CommitEvt,
+      })
+    } else if (row.eventType === 'sync') {
+      seqEvts.push({
+        type: 'sync',
+        seq: row.seq,
+        time: row.sequencedAt,
+        evt: evt as SyncEvt,
+      })
+    } else if (row.eventType === 'identity') {
+      seqEvts.push({
+        type: 'identity',
+        seq: row.seq,
+        time: row.sequencedAt,
+        evt: evt as IdentityEvt,
+      })
+    } else if (row.eventType === 'account') {
+      seqEvts.push({
+        type: 'account',
+        seq: row.seq,
+        time: row.sequencedAt,
+        evt: evt as AccountEvt,
+      })
+    }
+  }
+  return seqEvts
+}
+
 type SeqRow = RepoSeqEntry
 
 type SequencerEvents = {