From d97ee744bf0640e5051043327eab6db179770cfb Mon Sep 17 00:00:00 2001
From: Hong Minhee <hong@minhee.org>
Date: Wed, 27 Nov 2024 02:32:32 +0900
Subject: [PATCH] Instrument actor key pairs dispatcher

---
 docs/manual/opentelemetry.md | 31 ++++++++++++++++---------------
 src/federation/middleware.ts | 23 +++++++++++++++++++++--
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/docs/manual/opentelemetry.md b/docs/manual/opentelemetry.md
index ab049e6..d183e8d 100644
--- a/docs/manual/opentelemetry.md
+++ b/docs/manual/opentelemetry.md
@@ -117,21 +117,22 @@ Instrumented spans
 Fedify automatically instruments the following operations with OpenTelemetry
 spans:
 
-| Span name                        | [Span kind] | Description                           |
-|----------------------------------|-------------|---------------------------------------|
-| `{method} {template}`            | Server      | Serves the incoming HTTP request.     |
-| `activitypub.dispatch_actor`     | Server      | Dispatches the ActivityPub actor.     |
-| `activitypub.get_actor_handle`   | Client      | Resolves the actor handle.            |
-| `activitypub.lookup_object`      | Client      | Looks up the Activity Streams object. |
-| `activitypub.parse_object`       | Internal    | Parses the Activity Streams object.   |
-| `http_signatures.sign`           | Internal    | Signs the HTTP request.               |
-| `http_signatures.verify`         | Internal    | Verifies the HTTP request signature.  |
-| `ld_signatures.sign`             | Internal    | Makes the Linked Data signature.      |
-| `ld_signatures.verify`           | Internal    | Verifies the Linked Data signature.   |
-| `object_integrity_proofs.sign`   | Internal    | Makes the object integrity proof.     |
-| `object_integrity_proofs.verify` | Internal    | Verifies the object integrity proof.  |
-| `webfinger.handle`               | Server      | Handles the WebFinger request.        |
-| `webfinger.lookup`               | Client      | Looks up the WebFinger resource.      |
+| Span name                              | [Span kind] | Description                                 |
+|----------------------------------------|-------------|---------------------------------------------|
+| `{method} {template}`                  | Server      | Serves the incoming HTTP request.           |
+| `activitypub.dispatch_actor`           | Server      | Dispatches the ActivityPub actor.           |
+| `activitypub.dispatch_actor_key_pairs` | Server      | Dispatches the ActivityPub actor key pairs. |
+| `activitypub.get_actor_handle`         | Client      | Resolves the actor handle.                  |
+| `activitypub.lookup_object`            | Client      | Looks up the Activity Streams object.       |
+| `activitypub.parse_object`             | Internal    | Parses the Activity Streams object.         |
+| `http_signatures.sign`                 | Internal    | Signs the HTTP request.                     |
+| `http_signatures.verify`               | Internal    | Verifies the HTTP request signature.        |
+| `ld_signatures.sign`                   | Internal    | Makes the Linked Data signature.            |
+| `ld_signatures.verify`                 | Internal    | Verifies the Linked Data signature.         |
+| `object_integrity_proofs.sign`         | Internal    | Makes the object integrity proof.           |
+| `object_integrity_proofs.verify`       | Internal    | Verifies the object integrity proof.        |
+| `webfinger.handle`                     | Server      | Handles the WebFinger request.              |
+| `webfinger.lookup`                     | Client      | Looks up the WebFinger resource.            |
 
 More operations will be instrumented in the future releases.
 
diff --git a/src/federation/middleware.ts b/src/federation/middleware.ts
index 69a0bb7..bd00fc8 100644
--- a/src/federation/middleware.ts
+++ b/src/federation/middleware.ts
@@ -1062,8 +1062,27 @@ export class FederationImpl<TContextData> implements Federation<TContextData> {
     };
     this.actorCallbacks = callbacks;
     const setters: ActorCallbackSetters<TContextData> = {
-      setKeyPairsDispatcher(dispatcher: ActorKeyPairsDispatcher<TContextData>) {
-        callbacks.keyPairsDispatcher = dispatcher;
+      setKeyPairsDispatcher: (
+        dispatcher: ActorKeyPairsDispatcher<TContextData>,
+      ) => {
+        callbacks.keyPairsDispatcher = (ctx, identifier) =>
+          this.#getTracer().startActiveSpan(
+            "activitypub.dispatch_actor_key_pairs",
+            { kind: SpanKind.SERVER },
+            async (span) => {
+              try {
+                return await dispatcher(ctx, identifier);
+              } catch (e) {
+                span.setStatus({
+                  code: SpanStatusCode.ERROR,
+                  message: String(e),
+                });
+                throw e;
+              } finally {
+                span.end();
+              }
+            },
+          );
         return setters;
       },
       mapHandle(mapper: ActorHandleMapper<TContextData>) {