openAgent

1import { Buffer } from "node:buffer"
2import { FirehoseClient, PutRecordBatchCommand } from "@aws-sdk/client-firehose"
3import { Effect, Layer } from "effect"
4import * as Context from "effect/Context"
5import { Resource } from "sst/resource"
6 
7const MAX_FIREHOSE_BATCH_SIZE = 500
8const MAX_FIREHOSE_ATTEMPTS = 3
9const LAKE_TYPE = /^([A-Za-z0-9_]+)\.([A-Za-z0-9_]+)$/
10 
11type IngestEvent = Record<string, unknown>
12type LakeRoute = { database: string; table: string }
13type FirehoseRecord = { Data: Uint8Array }
14 
15export class IngestError extends Error {
16  readonly _tag = "IngestError"
17  readonly failed: number
18 
19  constructor(input: { message: string; failed: number; cause?: unknown }) {
20    super(input.message, { cause: input.cause })
21    this.name = "IngestError"
22    this.failed = input.failed
23  }
24}
25 
26export declare namespace Ingest {
27  export interface Service {
28    readonly write: (events: unknown[]) => Effect.Effect<{ records: number }, IngestError>
29  }
30}
31 
32export class Ingest extends Context.Service<Ingest, Ingest.Service>()("@opencode/stats/Ingest") {
33  static readonly layer: Layer.Layer<Ingest> = Layer.effect(
34    Ingest,
35    Effect.sync(() => {
36      const client = new FirehoseClient({})
37 
38      const write = Effect.fn("Ingest.write")(function* (events: unknown[]) {
39        if (events.length === 0) return { records: 0 }
40        const counts = countRoutedEvents(events)
41        if (counts.unsupported > 0) {
42          yield* Effect.logWarning(
43            `lake ingest rejected ${JSON.stringify({ records: counts.records, unsupported: counts.unsupported })}`,
44          )
45          return yield* Effect.fail(
46            new IngestError({
47              message: "Unsupported lake event type",
48              failed: counts.unsupported,
49            }),
50          )
51        }
52        if (counts.records === 0) return { records: 0 }
53 
54        let batch: FirehoseRecord[] = []
55        let batches = 0
56        let failed = 0
57 
58        for (const event of events) {
59          if (!isRecord(event)) continue
60          const route = routeEvent(event)
61          if (!route) continue
62          batch.push(toFirehoseRecord(event, route))
63          if (batch.length < MAX_FIREHOSE_BATCH_SIZE) continue
64          failed += yield* putRecords(client, Resource.LakeIngestConfig.streamName, batch)
65          batches++
66          batch = []
67        }
68 
69        if (batch.length > 0) {
70          failed += yield* putRecords(client, Resource.LakeIngestConfig.streamName, batch)
71          batches++
72        }
73 
74        if (failed > 0) {
75          yield* Effect.logWarning(`lake ingest incomplete ${JSON.stringify({ records: counts.records, failed })}`)
76          return yield* Effect.fail(new IngestError({ message: "Failed to ingest all lake records", failed }))
77        }
78 
79        yield* Effect.logInfo(`lake ingest complete ${JSON.stringify({ records: counts.records, batches })}`)
80        return { records: counts.records }
81      })
82 
83      return Ingest.of({ write })
84    }),
85  )
86}
87 
88const putRecords: (
89  client: FirehoseClient,
90  streamName: string,
91  records: FirehoseRecord[],
92  attempt?: number,
93) => Effect.Effect<number, IngestError> = Effect.fn("Ingest.putRecords")(function* (
94  client,
95  streamName,
96  records,
97  attempt = 1,
98) {
99  const result = yield* Effect.tryPromise({
100    try: () => client.send(new PutRecordBatchCommand({ DeliveryStreamName: streamName, Records: records })),
101    catch: (cause) =>
102      new IngestError({ message: "Failed to write lake records to Firehose", failed: records.length, cause }),
103  }).pipe(
104    Effect.tapError(() =>
105      Effect.logWarning(`firehose batch write failed ${JSON.stringify({ records: records.length, attempt })}`),
106    ),
107  )
108  const failed =
109    result.RequestResponses?.flatMap((item, index) => {
110      const record = records[index]
111      if (!item.ErrorCode || !record) return []
112      return [record]
113    }) ?? []
114 
115  if (failed.length === 0) return 0
116  if (attempt >= MAX_FIREHOSE_ATTEMPTS) {
117    yield* Effect.logWarning(
118      `firehose batch failed ${JSON.stringify({ records: failed.length, attempts: MAX_FIREHOSE_ATTEMPTS })}`,
119    )
120    return failed.length
121  }
122 
123  yield* Effect.logWarning(
124    `firehose batch retrying ${JSON.stringify({ records: failed.length, attempt: attempt + 1 })}`,
125  )
126  yield* Effect.sleep(`${250 * 2 ** (attempt - 1)} millis`)
127  return yield* putRecords(client, streamName, failed, attempt + 1)
128})
129 
130function countRoutedEvents(events: unknown[]) {
131  let records = 0
132  let unsupported = 0
133  for (const event of events) {
134    if (!isRecord(event)) continue
135    if (routeEvent(event)) records++
136    else unsupported++
137  }
138  return { records, unsupported }
139}
140 
141function isRecord(item: unknown): item is IngestEvent {
142  return Boolean(item) && typeof item === "object" && !Array.isArray(item)
143}
144 
145function routeEvent(event: IngestEvent): LakeRoute | undefined {
146  if (typeof event._datalake_key !== "string") return
147  const match = event._datalake_key.match(LAKE_TYPE)
148  if (!match?.[1] || !match[2]) return
149  return {
150    database: match[1],
151    table: match[2],
152  }
153}
154 
155function toFirehoseRecord(event: IngestEvent, route: LakeRoute): FirehoseRecord {
156  return {
157    Data: Buffer.from(
158      JSON.stringify({
159        ...Object.fromEntries(Object.entries(event).filter(([key]) => key !== "_datalake_key")),
160        _lake_database: route.database,
161        _lake_table: route.table,
162        _lake_operation: "insert" as const,
163      }),
164    ),
165  }
166}
167