protocol: add tracing

Author: fracek

change/@apibara-protocol-7963b847-a986-44cc-a679-3c6a33da87a2.json

@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "protocol: add tracing",
+  "packageName": "@apibara/protocol",
+  "email": "francesco@ceccon.me",
+  "dependentChangeType": "patch"
+}

packages/protocol/package.json

@@ -55,6 +55,7 @@
     "vitest": "^1.6.0"
   },
   "dependencies": {
+    "@opentelemetry/api": "^1.9.0",
     "consola": "^3.4.2",
     "long": "^5.2.3",
     "nice-grpc": "^2.1.8",

packages/protocol/src/rpc/data-stream.ts

@@ -4,9 +4,12 @@ import type { StreamDataRequest, StreamDataResponse } from "../stream";
 import { type ChainTracker, createChainTracker } from "./chain-tracker";
 import type { RpcStreamConfig } from "./config";
 import { blockInfoToCursor } from "./helpers";
+import { createTracer } from "./otel";
 
 const DEFAULT_HEARTBEAT_INTERVAL_MS = 30_000;
 
+const tracer = createTracer();
+
 type State<TFilter, TBlock> = {
   // The network-specific config.
   config: RpcStreamConfig<TFilter, TBlock>;
@@ -118,53 +121,93 @@ async function* dataStreamLoop<TFilter, TBlock>(
   state: State<TFilter, TBlock>,
 ): AsyncGenerator<StreamDataResponse<TBlock>> {
   while (shouldContinue(state)) {
-    const { cursor, chainTracker } = state;
-
-    // Always check for heartbeats first to ensure we don't miss any.
-    if (shouldSendHeartbeat(state)) {
-      state.lastHeartbeat = Date.now();
-      yield { _tag: "heartbeat" };
-    }
-
-    if (shouldRefreshFinalized(state)) {
-      const finalizedInfo = await state.config.fetchCursor({
-        blockTag: "finalized",
-      });
-
-      if (finalizedInfo === null) {
-        throw new Error("Failed to fetch finalized cursor");
-      }
-
-      const finalized = blockInfoToCursor(finalizedInfo);
-      const finalizedChanged =
-        state.chainTracker.updateFinalized(finalizedInfo);
-
-      // Only send finalized if it's needed.
-      if (finalizedChanged && state.cursor.orderKey > finalized.orderKey) {
-        yield { _tag: "finalize", finalize: { cursor: finalized } };
-      }
-
-      state.lastFinalizedRefresh = Date.now();
-    }
-
-    const finalized = chainTracker.finalized();
+    const messages = await tracer.startActiveSpan(
+      "data-stream.loop",
+      async (span) => {
+        try {
+          const { cursor, chainTracker } = state;
+
+          const attributes: Record<string, string | number | boolean> = {
+            cursor: cursor.orderKey.toString(),
+            head: chainTracker.head().orderKey.toString(),
+            finalized: chainTracker.finalized().orderKey.toString(),
+            actionSendHeartbeat: false,
+            actionRefreshFinalized: false,
+            actionBackfillFinalized: false,
+            actionWaitForHeadChange: false,
+            actionProduceLiveBlocks: false,
+          };
 
-    // console.log(
-    //   `[DS] RpcLoop: c=${cursor.orderKey} f=${finalized.orderKey} h=${chainTracker.head().orderKey}`,
-    // );
+          const messages: StreamDataResponse<TBlock>[] = [];
+
+          if (shouldSendHeartbeat(state)) {
+            state.lastHeartbeat = Date.now();
+            attributes.actionSendHeartbeat = true;
+            messages.push({ _tag: "heartbeat" });
+          }
+
+          if (shouldRefreshFinalized(state)) {
+            attributes.actionRefreshFinalized = true;
+            const finalizedInfo = await state.config.fetchCursor({
+              blockTag: "finalized",
+            });
+
+            if (finalizedInfo === null) {
+              throw new Error("Failed to fetch finalized cursor");
+            }
+
+            const finalized = blockInfoToCursor(finalizedInfo);
+            const finalizedChanged =
+              state.chainTracker.updateFinalized(finalizedInfo);
+
+            if (
+              finalizedChanged &&
+              state.cursor.orderKey > finalized.orderKey
+            ) {
+              messages.push({
+                _tag: "finalize",
+                finalize: { cursor: finalized },
+              });
+            }
+
+            state.lastFinalizedRefresh = Date.now();
+          }
+
+          const finalized = chainTracker.finalized();
+
+          if (cursor.orderKey < finalized.orderKey) {
+            attributes.actionBackfillFinalized = true;
+            for await (const msg of backfillFinalizedBlocks(state)) {
+              messages.push(msg);
+            }
+          } else {
+            if (isAtHead(state)) {
+              attributes.actionWaitForHeadChange = true;
+              for await (const msg of waitForHeadChange(state)) {
+                messages.push(msg);
+              }
+            } else {
+              attributes.actionProduceLiveBlocks = true;
+              for await (const msg of produceLiveBlocks(state)) {
+                messages.push(msg);
+              }
+            }
+          }
+
+          span.setAttributes(attributes);
+
+          return messages;
+        } catch (error) {
+          span.recordException(error as Error);
+          throw error;
+        } finally {
+          span.end();
+        }
+      },
+    );
 
-    if (cursor.orderKey < finalized.orderKey) {
-      yield* backfillFinalizedBlocks(state);
-    } else {
-      // If we're at the head, wait for a change.
-      //
-      // We don't want to produce a block immediately, but re-run the loop so
-      // that it's like any other iteration.
-      if (isAtHead(state)) {
-        yield* waitForHeadChange(state);
-      } else {
-        yield* produceLiveBlocks(state);
-      }
+    for (const message of messages) {
+      yield message;
     }
   }
 }

packages/protocol/src/rpc/otel.ts

@@ -0,0 +1,5 @@
+import { trace } from "@opentelemetry/api";
+
+export function createTracer() {
+  return trace.getTracer("@apibara/protocol");
+}