@atproto/ozone 0.1.69 → 0.1.70

package/src/background.ts

@@ -1,35 +1,161 @@
 import PQueue from 'p-queue'
 import { Database } from './db'
 import { dbLogger } from './logger'
+import { boundAbortController, isCausedBySignal, startInterval } from './util'
 
-// A simple queue for in-process, out-of-band/backgrounded work
+type Task = (db: Database, signal: AbortSignal) => Promise<void>
 
+/**
+ * A simple queue for in-process, out-of-band/backgrounded work
+ */
 export class BackgroundQueue {
-  queue = new PQueue({ concurrency: 20 })
-  destroyed = false
-  constructor(public db: Database) {}
+  private abortController = new AbortController()
+  private queue = new PQueue({ concurrency: 20 })
 
-  add(task: Task) {
+  public get signal() {
+    return this.abortController.signal
+  }
+
+  public get destroyed() {
+    return this.signal.aborted
+  }
+
+  constructor(protected db: Database) {}
+
+  getStats() {
+    return {
+      runningCount: this.queue.pending,
+      waitingCount: this.queue.size,
+    }
+  }
+
+  /**
+   * Add a task that will be executed at some point in the future.
+   *
+   * The task will be executed even if the backgroundQueue is destroyed, unless
+   * the provided `signal` is aborted.
+   *
+   * The `signal` provided to the task will be aborted whenever either the
+   * backgroundQueue is destroyed or the provided `signal` is aborted.
+   */
+  async add(task: Task, signal?: AbortSignal): Promise<void> {
     if (this.destroyed) {
       return
     }
-    this.queue
-      .add(() => task(this.db))
-      .catch((err) => {
-        dbLogger.error(err, 'background queue task failed')
-      })
+
+    const abortController = boundAbortController(this.signal, signal)
+
+    return this.queue.add<void>(async () => {
+      try {
+        // Do not run the task if the signal provided to the task has become
+        // aborted. Do not use `abortController.signal` here since we do not
+        // want to abort the task if the backgroundQueue is being destroyed.
+        if (signal?.aborted) return
+
+        // The task will receive a "combined signal" allowing it to abort if
+        // either the backgroundQueue is destroyed or the provided signal is
+        // aborted.
+        await task(this.db, abortController.signal)
+      } catch (err) {
+        if (!isCausedBySignal(err, abortController.signal)) {
+          dbLogger.error(err, 'background queue task failed')
+        }
+      } finally {
+        abortController.abort()
+      }
+    })
   }
 
   async processAll() {
     await this.queue.onIdle()
   }
 
-  // On destroy we stop accepting new tasks, but complete all pending/in-progress tasks.
-  // The application calls this only once http connections have drained (tasks no longer being added).
+  /**
+   * On destroy we stop accepting new tasks, but complete all
+   * pending/in-progress tasks. Tasks can decide to abort their current
+   * operation based on the signal they received. The application calls this
+   * only once http connections have drained (tasks no longer being added).
+   */
   async destroy() {
-    this.destroyed = true
+    this.abortController.abort()
     await this.queue.onIdle()
   }
 }
 
-type Task = (db: Database) => Promise<void>
+/**
+ * A simple periodic background task runner. This class will schedule a task to
+ * run through a provided {@link BackgroundQueue} at a fixed interval. The task
+ * will never run more than once concurrently, and will wait at least `interval`
+ * milliseconds between the end of one run and the start of the next.
+ */
+export class PeriodicBackgroundTask {
+  private abortController: AbortController
+
+  private intervalPromise?: Promise<void>
+  private runningPromise?: Promise<void>
+
+  public get signal() {
+    return this.abortController.signal
+  }
+
+  public get destroyed() {
+    return this.signal.aborted
+  }
+
+  constructor(
+    protected backgroundQueue: BackgroundQueue,
+    protected interval: number,
+    protected task: Task,
+  ) {
+    if (!Number.isFinite(interval) || interval <= 0) {
+      throw new TypeError('interval must be a positive number')
+    }
+
+    // Bind this class's signal to the backgroundQueue's signal (destroying this
+    // instance if the backgroundQueue is destroyed)
+    this.abortController = boundAbortController(backgroundQueue.signal)
+  }
+
+  public run(signal?: AbortSignal): Promise<void> {
+    // `startInterval` already ensures that only one run is in progress at a
+    // time. However, we want to be able to expose a `run()` method that can be
+    // used to force a run, which could cause concurrent executions. We prevent
+    // this using the `runningPromise` property.
+
+    if (this.runningPromise) return this.runningPromise
+
+    // Combine the `this.signal` with the provided `signal`, if any.
+    const abortController = boundAbortController(this.signal, signal)
+
+    const promise = this.backgroundQueue.add(this.task, abortController.signal)
+
+    return (this.runningPromise = promise).finally(() => {
+      if (this.runningPromise === promise) this.runningPromise = undefined
+
+      // Cleanup the listeners added by `boundAbortController`
+      abortController.abort()
+    })
+  }
+
+  public start() {
+    // Noop if already started. Throws if this.signal is aborted (instance is
+    // destroyed).
+    this.intervalPromise ||= startInterval(
+      async (signal) => this.run(signal),
+      this.interval,
+      this.signal,
+    )
+  }
+
+  public async destroy() {
+    // @NOTE This instance does not "own" the backgroundQueue, so we do not
+    // destroy it here.
+
+    this.abortController.abort()
+    console.error('ABOOOORT')
+
+    await this.intervalPromise
+    this.intervalPromise = undefined
+    console.error('DONE -_-')
+  }
+}

package/src/config/config.ts

@@ -24,6 +24,7 @@ export const envToCfg = (env: OzoneEnvironment): OzoneConfig => {
     poolSize: env.dbPoolSize,
     poolMaxUses: env.dbPoolMaxUses,
     poolIdleTimeoutMs: env.dbPoolIdleTimeoutMs,
+    materializedViewRefreshIntervalMs: env.dbMaterializedViewRefreshIntervalMs,
   }
 
   assert(env.appviewUrl, 'appviewUrl is required')
@@ -122,6 +123,7 @@ export type DatabaseConfig = {
   poolSize?: number
   poolMaxUses?: number
   poolIdleTimeoutMs?: number
+  materializedViewRefreshIntervalMs?: number
 }
 
 export type AppviewConfig = {

package/src/config/env.ts

@@ -20,6 +20,9 @@ export const readEnv = (): OzoneEnvironment => {
     dbPoolSize: envInt('OZONE_DB_POOL_SIZE'),
     dbPoolMaxUses: envInt('OZONE_DB_POOL_MAX_USES'),
     dbPoolIdleTimeoutMs: envInt('OZONE_DB_POOL_IDLE_TIMEOUT_MS'),
+    dbMaterializedViewRefreshIntervalMs: envInt(
+      'OZONE_DB_MATERIALIZED_VIEW_REFRESH_INTERVAL_MS',
+    ),
     didPlcUrl: envStr('OZONE_DID_PLC_URL'),
     didCacheStaleTTL: envInt('OZONE_DID_CACHE_STALE_TTL'),
     didCacheMaxTTL: envInt('OZONE_DID_CACHE_MAX_TTL'),
@@ -53,6 +56,7 @@ export type OzoneEnvironment = {
   dbPoolSize?: number
   dbPoolMaxUses?: number
   dbPoolIdleTimeoutMs?: number
+  dbMaterializedViewRefreshIntervalMs?: number
   didPlcUrl?: string
   didCacheStaleTTL?: number
   didCacheMaxTTL?: number

package/src/daemon/context.ts

@@ -6,17 +6,20 @@ import { OzoneConfig, OzoneSecrets } from '../config'
 import { Database } from '../db'
 import { EventPusher } from './event-pusher'
 import { EventReverser } from './event-reverser'
-import { ModerationService, ModerationServiceCreator } from '../mod-service'
+import { ModerationService } from '../mod-service'
 import { BackgroundQueue } from '../background'
 import { getSigningKeyId } from '../util'
+import { MaterializedViewRefresher } from './materialized-view-refresher'
+import { allFulfilled } from '@atproto/common'
 
 export type DaemonContextOptions = {
   db: Database
   cfg: OzoneConfig
-  modService: ModerationServiceCreator
+  backgroundQueue: BackgroundQueue
   signingKey: Keypair
   eventPusher: EventPusher
   eventReverser: EventReverser
+  materializedViewRefresher: MaterializedViewRefresher
 }
 
 export class DaemonContext {
@@ -67,13 +70,19 @@ export class DaemonContext {
 
     const eventReverser = new EventReverser(db, modService)
 
+    const materializedViewRefresher = new MaterializedViewRefresher(
+      backgroundQueue,
+      cfg.db.materializedViewRefreshIntervalMs,
+    )
+
     return new DaemonContext({
       db,
       cfg,
-      modService,
+      backgroundQueue,
       signingKey,
       eventPusher,
       eventReverser,
+      materializedViewRefresher,
       ...(overrides ?? {}),
     })
   }
@@ -86,8 +95,8 @@ export class DaemonContext {
     return this.opts.cfg
   }
 
-  get modService(): ModerationServiceCreator {
-    return this.opts.modService
+  get backgroundQueue(): BackgroundQueue {
+    return this.opts.backgroundQueue
   }
 
   get eventPusher(): EventPusher {
@@ -97,6 +106,35 @@ export class DaemonContext {
   get eventReverser(): EventReverser {
     return this.opts.eventReverser
   }
+
+  get materializedViewRefresher(): MaterializedViewRefresher {
+    return this.opts.materializedViewRefresher
+  }
+
+  async start() {
+    this.eventPusher.start()
+    this.eventReverser.start()
+    this.materializedViewRefresher.start()
+  }
+
+  async processAll() {
+    // Sequential because the materialized view values depend on the events.
+    await this.eventPusher.processAll()
+    await this.materializedViewRefresher.run()
+  }
+
+  async destroy() {
+    try {
+      await allFulfilled([
+        this.eventReverser.destroy(),
+        this.eventPusher.destroy(),
+        this.materializedViewRefresher.destroy(),
+      ])
+    } finally {
+      await this.backgroundQueue.destroy()
+      await this.db.close()
+    }
+  }
 }
 
 export default DaemonContext

package/src/daemon/index.ts

@@ -18,17 +18,14 @@ export class OzoneDaemon {
   }
 
   async start() {
-    this.ctx.eventPusher.start()
-    this.ctx.eventReverser.start()
+    await this.ctx.start()
   }
 
   async processAll() {
-    await this.ctx.eventPusher.processAll()
+    await this.ctx.processAll()
   }
 
   async destroy() {
-    await this.ctx.eventReverser.destroy()
-    await this.ctx.eventPusher.destroy()
-    await this.ctx.db.close()
+    await this.ctx.destroy()
   }
 }

package/src/daemon/materialized-view-refresher.ts

@@ -0,0 +1,27 @@
+import { MINUTE } from '@atproto/common'
+import { sql } from 'kysely'
+import { BackgroundQueue, PeriodicBackgroundTask } from '../background'
+
+export class MaterializedViewRefresher extends PeriodicBackgroundTask {
+  constructor(backgroundQueue: BackgroundQueue, interval = 30 * MINUTE) {
+    super(backgroundQueue, interval, async ({ db }, signal) => {
+      for (const view of [
+        'account_events_stats',
+        'record_events_stats',
+        'account_record_events_stats',
+        'account_record_status_stats',
+      ]) {
+        if (signal.aborted) break
+
+        // Kysely does not provide a way to cancel a running query. Because of
+        // this, killing the process during a refresh will cause the process to
+        // wait for the current refresh to finish before exiting. This is not
+        // ideal, but it is the best we can do until Kysely provides a way to
+        // cancel a query.
+        await sql`REFRESH MATERIALIZED VIEW CONCURRENTLY ${sql.id(view)}`.execute(
+          db,
+        )
+      }
+    })
+  }
+}

package/src/db/migrations/20241220T144630860Z-stats-materialized-views.ts

@@ -0,0 +1,218 @@
+import { Kysely, sql } from 'kysely'
+
+import { REASONAPPEAL } from '../../lexicon/types/com/atproto/moderation/defs'
+import { DatabaseSchemaType } from '../schema'
+
+import {
+  REVIEWESCALATED,
+  REVIEWOPEN,
+} from '../../lexicon/types/tools/ozone/moderation/defs'
+import * as modEvent from '../schema/moderation_event'
+import * as modStatus from '../schema/moderation_subject_status'
+import * as recordEventsStats from '../schema/record_events_stats'
+
+export async function up(db: Kysely<any>): Promise<void> {
+  // Used by "tools.ozone.moderation.queryStatuses". Reduces query cost by two
+  // order of magnitudes when sorting using "reportedRecordsCount" or
+  // "takendownRecordsCount" and filtering by "reviewState".
+  await db.schema
+    .createIndex('moderation_subject_status_did_id_review_state_idx')
+    .on('moderation_subject_status')
+    .column('did')
+    .expression(sql`"id" ASC NULLS FIRST`)
+    .column('reviewState')
+    .execute()
+
+  // ~6sec for 16M events
+  await db.schema
+    .createView('account_events_stats')
+    .materialized()
+    .ifNotExists()
+    .as(
+      (db as Kysely<modEvent.PartialDB>)
+        .selectFrom('moderation_event')
+        .where('subjectType', '=', 'com.atproto.admin.defs#repoRef')
+        .where('subjectUri', 'is', null)
+        .select('subjectDid')
+        .select([
+          (eb) =>
+            sql<number>`COUNT(*) FILTER(
+              WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventTakedown'
+              AND ${eb.ref('durationInHours')} IS NULL
+            )`.as('takedownCount'),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER(
+              WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventTakedown'
+              AND ${eb.ref('durationInHours')} IS NOT NULL
+            )`.as('suspendCount'),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER(
+              WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventEscalate'
+            )`.as('escalateCount'),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER(
+              WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventReport'
+              AND ${eb.ref('meta')} ->> 'reportType' != ${REASONAPPEAL}
+            )`.as('reportCount'),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER(
+              WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventReport'
+              AND ${eb.ref('meta')} ->> 'reportType' = ${REASONAPPEAL}
+            )`.as('appealCount'),
+        ])
+        .groupBy('subjectDid'),
+    )
+    .execute()
+
+  await db.schema
+    .createIndex('account_events_stats_did_idx')
+    .unique()
+    .on('account_events_stats')
+    .column('subjectDid')
+    .execute()
+
+  await db.schema
+    .createIndex('account_events_stats_suspend_count_idx')
+    .on('account_events_stats')
+    .expression(sql`"suspendCount" ASC NULLS FIRST`)
+    .column('subjectDid')
+    .execute()
+
+  // ~50sec for 16M events
+  await db.schema
+    .createView('record_events_stats')
+    .materialized()
+    .ifNotExists()
+    .as(
+      (db as Kysely<modEvent.PartialDB>)
+        .selectFrom('moderation_event')
+        .select([
+          'subjectDid',
+          'subjectUri',
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventEscalate')`.as(
+              'escalateCount',
+            ),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventReport' AND ${eb.ref('meta')} ->> 'reportType' != 'com.atproto.moderation.defs#reasonAppeal')`.as(
+              'reportCount',
+            ),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('action')} = 'tools.ozone.moderation.defs#modEventReport' AND ${eb.ref('meta')} ->> 'reportType' = 'com.atproto.moderation.defs#reasonAppeal')`.as(
+              'appealCount',
+            ),
+        ])
+        .where('subjectType', '=', 'com.atproto.repo.strongRef')
+        .where('subjectUri', 'is not', null)
+        .groupBy(['subjectDid', 'subjectUri']),
+    )
+    .execute()
+
+  await db.schema
+    .createIndex('record_events_stats_uri_idx')
+    .unique()
+    .on('record_events_stats')
+    .column('subjectUri')
+    .execute()
+
+  await db.schema
+    .createIndex('record_events_stats_did_idx')
+    .on('record_events_stats')
+    .column('subjectDid')
+    .execute()
+
+  await db.schema
+    .createView('account_record_events_stats')
+    .materialized()
+    .ifNotExists()
+    .as(
+      (db as Kysely<recordEventsStats.PartialDB>)
+        .selectFrom('record_events_stats')
+        .select([
+          'subjectDid',
+          (eb) =>
+            // Casting to "bigint" because "numeric" gets casted to a string
+            // by default by postgres-node.
+            sql<number>`SUM(${eb.ref('reportCount')})::bigint`.as(
+              'totalReports',
+            ),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('reportCount')} > 0)`.as(
+              'reportedCount',
+            ),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('escalateCount')} > 0)`.as(
+              'escalatedCount',
+            ),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('appealCount')} > 0)`.as(
+              'appealedCount',
+            ),
+        ])
+        .groupBy('subjectDid'),
+    )
+    .execute()
+
+  await db.schema
+    .createIndex('account_record_events_stats_did_idx')
+    .unique()
+    .on('account_record_events_stats')
+    .column('subjectDid')
+    .execute()
+
+  await db.schema
+    .createIndex('account_record_events_stats_reported_count_idx')
+    .on('account_record_events_stats')
+    .expression(sql`"reportedCount" ASC NULLS FIRST`)
+    .column('subjectDid')
+    .execute()
+
+  await db.schema
+    .createView('account_record_status_stats')
+    .materialized()
+    .ifNotExists()
+    .as(
+      (db as Kysely<modStatus.PartialDB>)
+        .selectFrom('moderation_subject_status')
+        .select('did')
+        .select([
+          sql<number>`COUNT(*)`.as('subjectCount'),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('reviewState')} IN (${REVIEWOPEN}, ${REVIEWESCALATED}))`.as(
+              'pendingCount',
+            ),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('reviewState')} NOT IN (${REVIEWOPEN}, ${REVIEWESCALATED}))`.as(
+              'processedCount',
+            ),
+          (eb) =>
+            sql<number>`COUNT(*) FILTER (WHERE ${eb.ref('takendown')})`.as(
+              'takendownCount',
+            ),
+        ])
+        .where('recordPath', '!=', '')
+        .groupBy('did'),
+    )
+    .execute()
+
+  await db.schema
+    .createIndex('account_record_status_stats_did_idx')
+    .unique()
+    .on('account_record_status_stats')
+    .column('did')
+    .execute()
+
+  await db.schema
+    .createIndex('account_record_status_stats_takendown_count_idx')
+    .on('account_record_status_stats')
+    .expression(sql`"takendownCount" ASC NULLS FIRST`)
+    .column('did')
+    .execute()
+}
+
+export async function down(db: Kysely<DatabaseSchemaType>): Promise<void> {
+  db.schema.dropView('account_record_status_stats').materialized().execute()
+  db.schema.dropView('account_record_events_stats').materialized().execute()
+  db.schema.dropView('record_events_stats').materialized().execute()
+  db.schema.dropView('account_events_stats').materialized().execute()
+}

package/src/db/migrations/index.ts

@@ -17,3 +17,4 @@ export * as _20241001T205730722Z from './20241001T205730722Z-subject-status-revi
 export * as _20241008T205730722Z from './20241008T205730722Z-sets'
 export * as _20241018T205730722Z from './20241018T205730722Z-setting'
 export * as _20241026T205730722Z from './20241026T205730722Z-add-hosting-status-to-subject-status'
+export * as _20241220T144630860Z from './20241220T144630860Z-stats-materialized-views'

package/src/db/schema/account_events_stats.ts

@@ -0,0 +1,16 @@
+import { GeneratedAlways, Selectable } from 'kysely'
+
+export const tableName = 'account_events_stats'
+
+export type AccountEventsStats = {
+  subjectDid: GeneratedAlways<string>
+  takedownCount: GeneratedAlways<number>
+  suspendCount: GeneratedAlways<number>
+  escalateCount: GeneratedAlways<number>
+  reportCount: GeneratedAlways<number>
+  appealCount: GeneratedAlways<number>
+}
+
+export type AccountEventsStatsRow = Selectable<AccountEventsStats>
+
+export type PartialDB = { [tableName]: AccountEventsStats }

package/src/db/schema/account_record_events_stats.ts

@@ -0,0 +1,15 @@
+import { GeneratedAlways, Selectable } from 'kysely'
+
+export const tableName = 'account_record_events_stats'
+
+type AccountRecordEventsStats = {
+  subjectDid: GeneratedAlways<string>
+  totalReports: GeneratedAlways<number>
+  reportedCount: GeneratedAlways<number>
+  escalatedCount: GeneratedAlways<number>
+  appealedCount: GeneratedAlways<number>
+}
+
+export type AccountRecordEventsStatsRow = Selectable<AccountRecordEventsStats>
+
+export type PartialDB = { [tableName]: AccountRecordEventsStats }

package/src/db/schema/account_record_status_stats.ts

@@ -0,0 +1,15 @@
+import { GeneratedAlways, Selectable } from 'kysely'
+
+export const tableName = 'account_record_status_stats'
+
+type AccountRecordStatusStats = {
+  did: GeneratedAlways<string>
+  subjectCount: GeneratedAlways<number>
+  pendingCount: GeneratedAlways<number>
+  processedCount: GeneratedAlways<number>
+  takendownCount: GeneratedAlways<number>
+}
+
+export type AccountRecordStatusStatsRow = Selectable<AccountRecordStatusStats>
+
+export type PartialDB = { [tableName]: AccountRecordStatusStats }

package/src/db/schema/index.ts

@@ -11,6 +11,11 @@ import * as set from './ozone_set'
 import * as member from './member'
 import * as setting from './setting'
 
+import * as recordEventsStats from './record_events_stats'
+import * as accountEventsStats from './account_events_stats'
+import * as accountRecordEventsStats from './account_record_events_stats'
+import * as accountRecordStatusStats from './account_record_status_stats'
+
 export type DatabaseSchemaType = modEvent.PartialDB &
   modSubjectStatus.PartialDB &
   label.PartialDB &
@@ -21,7 +26,11 @@ export type DatabaseSchemaType = modEvent.PartialDB &
   communicationTemplate.PartialDB &
   set.PartialDB &
   member.PartialDB &
-  setting.PartialDB
+  setting.PartialDB &
+  accountEventsStats.PartialDB &
+  recordEventsStats.PartialDB &
+  accountRecordEventsStats.PartialDB &
+  accountRecordStatusStats.PartialDB
 
 export type DatabaseSchema = Kysely<DatabaseSchemaType>
 

package/src/db/schema/record_events_stats.ts

@@ -0,0 +1,15 @@
+import { GeneratedAlways, Selectable } from 'kysely'
+
+export const tableName = 'record_events_stats'
+
+export type RecordEventsStats = {
+  subjectDid: GeneratedAlways<string>
+  subjectUri: GeneratedAlways<string>
+  escalateCount: GeneratedAlways<number>
+  reportCount: GeneratedAlways<number>
+  appealCount: GeneratedAlways<number>
+}
+
+export type RecordEventsStatsRow = Selectable<RecordEventsStats>
+
+export type PartialDB = { [tableName]: RecordEventsStats }

package/src/index.ts

@@ -114,13 +114,7 @@ export class OzoneService {
         },
         'db pool stats',
       )
-      dbLogger.info(
-        {
-          runningCount: backgroundQueue.queue.pending,
-          waitingCount: backgroundQueue.queue.size,
-        },
-        'background queue stats',
-      )
+      dbLogger.info(backgroundQueue.getStats(), 'background queue stats')
     }, 10000)
     await this.ctx.sequencer.start()
     const server = this.app.listen(this.ctx.cfg.service.port)