From 99b6aee18ec85c4f18ac36e502987304ed7e26ff Mon Sep 17 00:00:00 2001
From: devin ivy <devinivy@gmail.com>
Date: Thu, 14 Mar 2024 09:57:24 -0400
Subject: [PATCH] Dedupe labels during appview hydration (#2317)

* appview: reproduce duped labels

* appview: dedupe labels in hydration

* tidy types
---
 .../src/data-plane/server/routes/labels.ts    | 32 +++++++++++----
 packages/bsky/src/hydration/hydrator.ts       | 13 +++---
 packages/bsky/src/hydration/label.ts          | 40 +++++++++++++++----
 packages/bsky/src/hydration/util.ts           |  8 +++-
 packages/bsky/src/views/index.ts              | 14 +++----
 packages/bsky/tests/label-hydration.test.ts   | 13 ++++++
 6 files changed, 91 insertions(+), 29 deletions(-)

diff --git a/packages/bsky/src/data-plane/server/routes/labels.ts b/packages/bsky/src/data-plane/server/routes/labels.ts
index de5625d7784..a4dd2e6acac 100644
--- a/packages/bsky/src/data-plane/server/routes/labels.ts
+++ b/packages/bsky/src/data-plane/server/routes/labels.ts
@@ -3,28 +3,44 @@ import { noUndefinedVals } from '@atproto/common'
 import { ServiceImpl } from '@connectrpc/connect'
 import { Service } from '../../../proto/bsky_connect'
 import { Database } from '../db'
+import { Selectable } from 'kysely'
+import { Label } from '../db/tables/label'
+
+type LabelRow = Selectable<Label>
 
 export default (db: Database): Partial<ServiceImpl<typeof Service>> => ({
   async getLabels(req) {
     const { subjects, issuers } = req
     if (subjects.length === 0 || issuers.length === 0) {
-      return { records: [] }
+      return { labels: [] }
     }
-    const res = await db.db
+    const res: LabelRow[] = await db.db
       .selectFrom('label')
       .where('uri', 'in', subjects)
       .where('src', 'in', issuers)
       .selectAll()
       .execute()
 
-    const labels = res.map((l) => {
-      const formatted = noUndefinedVals({
-        ...l,
-        cid: l.cid === '' ? undefined : l.cid,
-        neg: l.neg === true ? true : undefined,
+    const labelsBySubject = new Map<string, LabelRow[]>()
+    res.forEach((l) => {
+      const labels = labelsBySubject.get(l.uri) ?? []
+      labels.push(l)
+      labelsBySubject.set(l.uri, labels)
+    })
+
+    // intentionally duplicate label results, appview frontend should be defensive to this
+    const labels = subjects.flatMap((sub) => {
+      const labelsForSub = labelsBySubject.get(sub) ?? []
+      return labelsForSub.map((l) => {
+        const formatted = noUndefinedVals({
+          ...l,
+          cid: l.cid === '' ? undefined : l.cid,
+          neg: l.neg === true ? true : undefined,
+        })
+        return ui8.fromString(JSON.stringify(formatted), 'utf8')
       })
-      return ui8.fromString(JSON.stringify(formatted), 'utf8')
     })
+
     return { labels }
   },
 })
diff --git a/packages/bsky/src/hydration/hydrator.ts b/packages/bsky/src/hydration/hydrator.ts
index d72c6477bad..5a16607b93f 100644
--- a/packages/bsky/src/hydration/hydrator.ts
+++ b/packages/bsky/src/hydration/hydrator.ts
@@ -29,7 +29,13 @@ import {
   Labelers,
   Labels,
 } from './label'
-import { HydrationMap, RecordInfo, didFromUri, urisByCollection } from './util'
+import {
+  HydrationMap,
+  Merges,
+  RecordInfo,
+  didFromUri,
+  urisByCollection,
+} from './util'
 import {
   FeedGenAggs,
   FeedGens,
@@ -770,10 +776,7 @@ export const mergeStates = (
   }
 }
 
-const mergeMaps = <T>(
-  mapA?: HydrationMap<T>,
-  mapB?: HydrationMap<T>,
-): HydrationMap<T> | undefined => {
+const mergeMaps = <M extends Merges>(mapA?: M, mapB?: M): M | undefined => {
   if (!mapA) return mapB
   if (!mapB) return mapA
   return mapA.merge(mapB)
diff --git a/packages/bsky/src/hydration/label.ts b/packages/bsky/src/hydration/label.ts
index f32b35a5496..7ca5b1b5b51 100644
--- a/packages/bsky/src/hydration/label.ts
+++ b/packages/bsky/src/hydration/label.ts
@@ -3,6 +3,7 @@ import { Label } from '../lexicon/types/com/atproto/label/defs'
 import { Record as LabelerRecord } from '../lexicon/types/app/bsky/labeler/service'
 import {
   HydrationMap,
+  Merges,
   RecordInfo,
   parseJsonBytes,
   parseRecord,
@@ -16,10 +17,36 @@ export type { Label } from '../lexicon/types/com/atproto/label/defs'
 
 export type SubjectLabels = {
   isTakendown: boolean
-  labels: Label[]
+  labels: HydrationMap<Label> // src + val -> label
 }
 
-export type Labels = HydrationMap<SubjectLabels>
+export class Labels extends HydrationMap<SubjectLabels> implements Merges {
+  static key(label: Label) {
+    return `${label.src}::${label.val}`
+  }
+  merge(map: Labels): this {
+    map.forEach((theirs, key) => {
+      if (!theirs) return
+      const mine = this.get(key)
+      if (mine) {
+        mine.isTakendown = mine.isTakendown || theirs.isTakendown
+        mine.labels = mine.labels.merge(theirs.labels)
+      } else {
+        this.set(key, theirs)
+      }
+    })
+    return this
+  }
+  getBySubject(sub: string): Label[] {
+    const it = this.get(sub)?.labels.values()
+    if (!it) return []
+    const labels: Label[] = []
+    for (const label of it) {
+      if (label) labels.push(label)
+    }
+    return labels
+  }
+}
 
 export type LabelerAgg = {
   likes: number
@@ -43,8 +70,7 @@ export class LabelHydrator {
     subjects: string[],
     labelers: ParsedLabelers,
   ): Promise<Labels> {
-    if (!subjects.length || !labelers.dids.length)
-      return new HydrationMap<SubjectLabels>()
+    if (!subjects.length || !labelers.dids.length) return new Labels()
     const res = await this.dataplane.getLabels({
       subjects,
       issuers: labelers.dids,
@@ -57,11 +83,11 @@ export class LabelHydrator {
       if (!entry) {
         entry = {
           isTakendown: false,
-          labels: [],
+          labels: new HydrationMap(),
         }
         acc.set(label.uri, entry)
       }
-      entry.labels.push(label)
+      entry.labels.set(Labels.key(label), label)
       if (
         TAKEDOWN_LABELS.includes(label.val) &&
         !label.neg &&
@@ -70,7 +96,7 @@ export class LabelHydrator {
         entry.isTakendown = true
       }
       return acc
-    }, new HydrationMap<SubjectLabels>())
+    }, new Labels())
   }
 
   async getLabelers(
diff --git a/packages/bsky/src/hydration/util.ts b/packages/bsky/src/hydration/util.ts
index 675c0e9f3d0..7654a7bbbe9 100644
--- a/packages/bsky/src/hydration/util.ts
+++ b/packages/bsky/src/hydration/util.ts
@@ -5,8 +5,8 @@ import * as ui8 from 'uint8arrays'
 import { lexicons } from '../lexicon/lexicons'
 import { Record } from '../proto/bsky_pb'
 
-export class HydrationMap<T> extends Map<string, T | null> {
-  merge(map: HydrationMap<T>): HydrationMap<T> {
+export class HydrationMap<T> extends Map<string, T | null> implements Merges {
+  merge(map: HydrationMap<T>): this {
     map.forEach((val, key) => {
       this.set(key, val)
     })
@@ -14,6 +14,10 @@ export class HydrationMap<T> extends Map<string, T | null> {
   }
 }
 
+export interface Merges {
+  merge<T extends this>(map: T): this
+}
+
 export type RecordInfo<T> = {
   record: T
   cid: string
diff --git a/packages/bsky/src/views/index.ts b/packages/bsky/src/views/index.ts
index de1e02bc5c6..cade128ba16 100644
--- a/packages/bsky/src/views/index.ts
+++ b/packages/bsky/src/views/index.ts
@@ -136,8 +136,8 @@ export class Views {
       'self',
     ).toString()
     const labels = [
-      ...(state.labels?.get(did)?.labels ?? []),
-      ...(state.labels?.get(profileUri)?.labels ?? []),
+      ...(state.labels?.getBySubject(did) ?? []),
+      ...(state.labels?.getBySubject(profileUri) ?? []),
       ...this.selfLabels({
         uri: profileUri,
         cid: actor.profileCid?.toString(),
@@ -225,7 +225,7 @@ export class Views {
       return undefined
     }
     const listViewer = state.listViewers?.get(uri)
-    const labels = state.labels?.get(uri)?.labels ?? []
+    const labels = state.labels?.getBySubject(uri) ?? []
     const creator = new AtUri(uri).hostname
     return {
       uri,
@@ -281,7 +281,7 @@ export class Views {
 
     const uri = AtUri.make(did, ids.AppBskyLabelerService, 'self').toString()
     const labels = [
-      ...(state.labels?.get(uri)?.labels ?? []),
+      ...(state.labels?.getBySubject(uri) ?? []),
       ...this.selfLabels({
         uri,
         cid: labeler.cid.toString(),
@@ -351,7 +351,7 @@ export class Views {
     if (!creator) return
     const viewer = state.feedgenViewers?.get(uri)
     const aggs = state.feedgenAggs?.get(uri)
-    const labels = state.labels?.get(uri)?.labels ?? []
+    const labels = state.labels?.getBySubject(uri) ?? []
 
     return {
       uri,
@@ -408,7 +408,7 @@ export class Views {
       parsedUri.rkey,
     ).toString()
     const labels = [
-      ...(state.labels?.get(uri)?.labels ?? []),
+      ...(state.labels?.getBySubject(uri) ?? []),
       ...this.selfLabels({
         uri,
         cid: post.cid,
@@ -886,7 +886,7 @@ export class Views {
       recordInfo = state.follows?.get(notif.uri)
     }
     if (!recordInfo) return
-    const labels = state.labels?.get(notif.uri)?.labels ?? []
+    const labels = state.labels?.getBySubject(notif.uri) ?? []
     const selfLabels = this.selfLabels({
       uri: notif.uri,
       cid: recordInfo.cid,
diff --git a/packages/bsky/tests/label-hydration.test.ts b/packages/bsky/tests/label-hydration.test.ts
index 236fbac1e7b..efe77512bdc 100644
--- a/packages/bsky/tests/label-hydration.test.ts
+++ b/packages/bsky/tests/label-hydration.test.ts
@@ -93,6 +93,19 @@ describe('label hydration', () => {
     )
   })
 
+  it('hydrates labels without duplication', async () => {
+    AtpAgent.configure({ appLabelers: [alice] })
+    pdsAgent.configureLabelersHeader([])
+    const res = await pdsAgent.api.app.bsky.actor.getProfiles(
+      { actors: [carol, carol] },
+      { headers: sc.getHeaders(bob) },
+    )
+    const { labels = [] } = res.data.profiles[0]
+    expect(labels.map((l) => ({ val: l.val, src: l.src }))).toEqual([
+      { src: alice, val: 'spam' },
+    ])
+  })
+
   it('hydrates labels onto list views.', async () => {
     AtpAgent.configure({ appLabelers: [labelerDid] })
     pdsAgent.configureLabelersHeader([])