@senzops/apm-node 1.2.8 → 1.3.0

package/src/instrumentation/dataloader.ts

@@ -0,0 +1,291 @@
+import { SenzorOptions } from '../core/types';
+import { hookRequire } from './hook';
+import { patchMethod } from './patch';
+import { runWithCapturedSpan, startCapturedSpan } from './span';
+
+// ---------------------------------------------------------------------------
+// DataLoader Instrumentation
+//
+// Instruments the `dataloader` package — the standard batching/caching
+// utility for solving N+1 queries in GraphQL resolvers, REST APIs,
+// and any data-fetching layer.
+//
+// Patches DataLoader.prototype methods:
+//   - load(key)        — single key lookup (batched)
+//   - loadMany(keys)   — multi-key lookup (batched)
+//   - prime(key, val)  — cache priming
+//   - clear(key)       — single cache eviction
+//   - clearAll()       — full cache clear
+//
+// The key insight: load() calls are batched by DataLoader and dispatched
+// together via the batch function. We instrument both the individual
+// load() calls AND the batch dispatch to show:
+//   1. Per-key latency (includes batching wait time)
+//   2. Batch execution performance
+//
+// Captured attributes:
+//   - dataloader.operation: load, loadMany, prime, clear, clearAll, batch
+//   - dataloader.key_count: number of keys in the batch
+//   - dataloader.name: DataLoader instance name (if available)
+//   - dataloader.cache_hit: whether result came from cache
+// ---------------------------------------------------------------------------
+
+// ---------------------------------------------------------------------------
+// DataLoader prototype patching
+// ---------------------------------------------------------------------------
+
+const patchDataLoader = (DataLoader: any, options?: SenzorOptions) => {
+  const proto = DataLoader?.prototype;
+  if (!proto) return;
+
+  // --- load(key) → Promise<value> ---
+  patchMethod(
+    proto,
+    'load',
+    'senzor.dataloader.load',
+    (original) =>
+      function patchedLoad(this: any, key: any) {
+        const loaderName = this.name || this._name || 'DataLoader';
+
+        const span = startCapturedSpan(
+          `${loaderName} load`,
+          'function',
+          {
+            'dataloader.operation': 'load',
+            'dataloader.name': loaderName,
+            library: 'dataloader',
+          },
+          options
+        );
+
+        if (!span) return original.call(this, key);
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            const result = original.call(this, key);
+
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (value: any) => {
+                  span.end(0);
+                  return value;
+                },
+                (error: any) => {
+                  span.end(500, {
+                    'error.message': error?.message,
+                    'error.type': error?.name || 'Error',
+                  });
+                  throw error;
+                }
+              );
+            }
+
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, { 'error.message': error?.message });
+            throw error;
+          }
+        });
+      }
+  );
+
+  // --- loadMany(keys) → Promise<value[]> ---
+  patchMethod(
+    proto,
+    'loadMany',
+    'senzor.dataloader.loadMany',
+    (original) =>
+      function patchedLoadMany(this: any, keys: any[]) {
+        const loaderName = this.name || this._name || 'DataLoader';
+        const keyCount = Array.isArray(keys) ? keys.length : 0;
+
+        const span = startCapturedSpan(
+          `${loaderName} loadMany`,
+          'function',
+          {
+            'dataloader.operation': 'loadMany',
+            'dataloader.name': loaderName,
+            'dataloader.key_count': keyCount,
+            library: 'dataloader',
+          },
+          options
+        );
+
+        if (!span) return original.call(this, keys);
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            const result = original.call(this, keys);
+
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (values: any) => {
+                  // Count errors in results
+                  const errorCount = Array.isArray(values)
+                    ? values.filter((v: any) => v instanceof Error).length
+                    : 0;
+
+                  span.end(errorCount > 0 ? 207 : 0, {
+                    'dataloader.key_count': keyCount,
+                    'dataloader.error_count': errorCount,
+                  });
+                  return values;
+                },
+                (error: any) => {
+                  span.end(500, { 'error.message': error?.message });
+                  throw error;
+                }
+              );
+            }
+
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, { 'error.message': error?.message });
+            throw error;
+          }
+        });
+      }
+  );
+
+  // --- Wrap the batch function dispatch ---
+  // DataLoader calls the batch function internally. We intercept it
+  // by wrapping the constructor or the internal _batchLoadFn.
+  // Since we can't easily wrap the constructor via patchMethod,
+  // we patch _batchScheduleFn or override the dispatch mechanism.
+
+  // Patch the internal _batch method if it exists
+  const batchMethodName = '_dispatchBatch' in proto
+    ? '_dispatchBatch'
+    : '_dispatch' in proto
+      ? '_dispatch'
+      : null;
+
+  if (batchMethodName && typeof proto[batchMethodName] === 'function') {
+    patchMethod(
+      proto,
+      batchMethodName,
+      `senzor.dataloader.${batchMethodName}`,
+      (original) =>
+        function patchedDispatch(this: any, ...args: any[]) {
+          const loaderName = this.name || this._name || 'DataLoader';
+          // The batch queue holds pending keys
+          const batchSize = this._queue?.length || this._batch?.length || 0;
+
+          const span = startCapturedSpan(
+            `${loaderName} batch dispatch`,
+            'function',
+            {
+              'dataloader.operation': 'batch',
+              'dataloader.name': loaderName,
+              'dataloader.batch_size': batchSize,
+              library: 'dataloader',
+            },
+            options
+          );
+
+          if (!span) return original.apply(this, args);
+
+          return runWithCapturedSpan(span, () => {
+            try {
+              const result = original.apply(this, args);
+              // Batch dispatch is sync; the batch function's promise
+              // is resolved internally. End span after dispatch.
+              span.end(0, { 'dataloader.batch_size': batchSize });
+              return result;
+            } catch (error: any) {
+              span.end(500, { 'error.message': error?.message });
+              throw error;
+            }
+          });
+        }
+    );
+  }
+
+  // --- prime(key, value) ---
+  if (typeof proto.prime === 'function') {
+    patchMethod(
+      proto,
+      'prime',
+      'senzor.dataloader.prime',
+      (original) =>
+        function patchedPrime(this: any, key: any, value: any) {
+          const loaderName = this.name || this._name || 'DataLoader';
+
+          const span = startCapturedSpan(
+            `${loaderName} prime`,
+            'function',
+            {
+              'dataloader.operation': 'prime',
+              'dataloader.name': loaderName,
+              library: 'dataloader',
+            },
+            options
+          );
+
+          if (!span) return original.call(this, key, value);
+
+          try {
+            const result = original.call(this, key, value);
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, { 'error.message': error?.message });
+            throw error;
+          }
+        }
+    );
+  }
+
+  // --- clearAll() ---
+  if (typeof proto.clearAll === 'function') {
+    patchMethod(
+      proto,
+      'clearAll',
+      'senzor.dataloader.clearAll',
+      (original) =>
+        function patchedClearAll(this: any) {
+          const loaderName = this.name || this._name || 'DataLoader';
+
+          const span = startCapturedSpan(
+            `${loaderName} clearAll`,
+            'function',
+            {
+              'dataloader.operation': 'clearAll',
+              'dataloader.name': loaderName,
+              library: 'dataloader',
+            },
+            options
+          );
+
+          if (!span) return original.call(this);
+
+          try {
+            const result = original.call(this);
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, { 'error.message': error?.message });
+            throw error;
+          }
+        }
+    );
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+export const instrumentDataloader = (options?: SenzorOptions) => {
+  hookRequire('dataloader', (exports: any) => {
+    // dataloader exports the constructor directly
+    patchDataLoader(exports, options);
+
+    // Handle default export (ESM interop)
+    if (exports?.default?.prototype) {
+      patchDataLoader(exports.default, options);
+    }
+  });
+};

package/src/instrumentation/dns.ts

@@ -0,0 +1,220 @@
+import { SenzorOptions } from '../core/types';
+import { patchMethod } from './patch';
+import { runWithCapturedSpan, startCapturedSpan } from './span';
+
+// ---------------------------------------------------------------------------
+// DNS Instrumentation
+//
+// Instruments Node.js core `dns` module:
+//   - dns.lookup()        — resolves a hostname to an address (uses OS resolver)
+//   - dns.resolve()       — resolves using DNS protocol (network call)
+//   - dns.resolve4()      — A records
+//   - dns.resolve6()      — AAAA records
+//   - dns.resolveMx()     — MX records
+//   - dns.resolveTxt()    — TXT records
+//   - dns.resolveSrv()    — SRV records
+//   - dns.resolveCname()  — CNAME records
+//   - dns.resolveNs()     — NS records
+//   - dns.reverse()       — reverse DNS lookup
+//
+// Follows OTel semantic conventions: dns.hostname, network.peer.address
+// ---------------------------------------------------------------------------
+
+/**
+ * Patch a single DNS method that accepts (hostname, [options], callback).
+ */
+const patchDnsMethod = (
+  dnsModule: any,
+  methodName: string,
+  options?: SenzorOptions
+) => {
+  patchMethod(
+    dnsModule,
+    methodName,
+    `senzor.dns.${methodName}`,
+    (original) =>
+      function patchedDnsMethod(this: any, ...args: any[]) {
+        // Extract hostname (always the first argument)
+        const hostname = typeof args[0] === 'string' ? args[0] : 'unknown';
+
+        const span = startCapturedSpan(
+          `DNS ${methodName}`,
+          'custom',
+          {
+            'dns.hostname': hostname,
+            'dns.operation': methodName,
+          },
+          options
+        );
+
+        if (!span) return original.apply(this, args);
+
+        // Find and wrap the callback (always the last argument)
+        const callbackIdx = args.findIndex(
+          (arg: any, idx: number) => typeof arg === 'function' && idx === args.length - 1
+        );
+
+        if (callbackIdx >= 0) {
+          const originalCallback = args[callbackIdx];
+          args[callbackIdx] = function wrappedDnsCallback(
+            err: NodeJS.ErrnoException | null,
+            ...results: any[]
+          ) {
+            if (err) {
+              span.end(500, {
+                'error.message': err.message,
+                'error.type': err.code || err.name || 'DnsError',
+                'dns.error_code': err.code,
+              });
+            } else {
+              // For lookup, first result is address, second is family
+              const meta: Record<string, any> = {};
+              if (methodName === 'lookup' && results[0]) {
+                meta['network.peer.address'] = results[0];
+                meta['dns.address_family'] = results[1] === 6 ? 'IPv6' : 'IPv4';
+              } else if (methodName === 'resolve4' || methodName === 'resolve6') {
+                meta['dns.result_count'] = Array.isArray(results[0]) ? results[0].length : 0;
+              } else if (methodName === 'resolve') {
+                meta['dns.result_count'] = Array.isArray(results[0]) ? results[0].length : 0;
+              }
+              span.end(0, meta);
+            }
+            return originalCallback.call(this, err, ...results);
+          };
+        }
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            return original.apply(this, args);
+          } catch (error: any) {
+            span.end(500, {
+              'error.message': error?.message,
+              'error.type': error?.name || 'Error',
+            });
+            throw error;
+          }
+        });
+      }
+  );
+};
+
+/**
+ * Patch dns.promises methods (returns promises instead of using callbacks).
+ */
+const patchDnsPromisesMethod = (
+  promises: any,
+  methodName: string,
+  options?: SenzorOptions
+) => {
+  patchMethod(
+    promises,
+    methodName,
+    `senzor.dns.promises.${methodName}`,
+    (original) =>
+      function patchedDnsPromiseMethod(this: any, ...args: any[]) {
+        const hostname = typeof args[0] === 'string' ? args[0] : 'unknown';
+
+        const span = startCapturedSpan(
+          `DNS ${methodName}`,
+          'custom',
+          {
+            'dns.hostname': hostname,
+            'dns.operation': methodName,
+          },
+          options
+        );
+
+        if (!span) return original.apply(this, args);
+
+        return runWithCapturedSpan(span, async () => {
+          try {
+            const result = await original.apply(this, args);
+
+            const meta: Record<string, any> = {};
+            if (methodName === 'lookup' && result) {
+              meta['network.peer.address'] = result.address;
+              meta['dns.address_family'] = result.family === 6 ? 'IPv6' : 'IPv4';
+            } else if (Array.isArray(result)) {
+              meta['dns.result_count'] = result.length;
+            }
+
+            span.end(0, meta);
+            return result;
+          } catch (error: any) {
+            span.end(500, {
+              'error.message': error?.message,
+              'error.type': error?.code || error?.name || 'DnsError',
+              'dns.error_code': error?.code,
+            });
+            throw error;
+          }
+        });
+      }
+  );
+};
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+export const instrumentDns = (options?: SenzorOptions) => {
+  let dns: any;
+  try {
+    dns = require('dns');
+  } catch {
+    return;
+  }
+
+  if (!dns) return;
+
+  // Callback-based methods
+  const callbackMethods = [
+    'lookup',
+    'resolve',
+    'resolve4',
+    'resolve6',
+    'resolveMx',
+    'resolveTxt',
+    'resolveSrv',
+    'resolveCname',
+    'resolveNs',
+    'resolvePtr',
+    'resolveSoa',
+    'resolveNaptr',
+    'resolveCaa',
+    'reverse',
+  ];
+
+  for (const method of callbackMethods) {
+    if (typeof dns[method] === 'function') {
+      patchDnsMethod(dns, method, options);
+    }
+  }
+
+  // Promise-based methods (dns.promises / dns/promises)
+  const promises = dns.promises;
+  if (promises) {
+    const promiseMethods = [
+      'lookup',
+      'resolve',
+      'resolve4',
+      'resolve6',
+      'resolveMx',
+      'resolveTxt',
+      'resolveSrv',
+      'resolveCname',
+      'resolveNs',
+      'resolvePtr',
+      'resolveSoa',
+      'resolveNaptr',
+      'resolveCaa',
+      'reverse',
+    ];
+
+    for (const method of promiseMethods) {
+      if (typeof promises[method] === 'function') {
+        patchDnsPromisesMethod(promises, method, options);
+      }
+    }
+  }
+};