@aidra/pmx-io 0.0.0-stage → 9.13.0

package/lib/itg/nats.js

@@ -0,0 +1,531 @@
+// Copyright (c) 2026 Mike Warriner. All rights reserved. Proprietary — see LICENSE.
+'use strict'
+
+// NATS and JetStream (82 §9 ITG-202, ITG-203; common rules §4) for `nats` 2.x and the 3.x modular
+// packages (`@nats-io/nats-core`, `@nats-io/jetstream`, through `@nats-io/transport-node`).
+//  * NatsConnectionImpl.publish → PRODUCER `publish <subject>`, trace headers in a copy of
+//    `options.headers` (a new headers object from the library's own `headers()` factory);
+//    request → CLIENT `request <subject>` until the reply (or the timeout, ERROR); requestMany →
+//    CLIENT span until the returned iterator is obtained.
+//  * subscribe(subject, { callback }) → callback style; a subscription iterated with `for await`
+//    → iterator style (QueuedIteratorImpl's async iterator, for subscriptions marked here).
+//  * JetStream: JetStreamClientImpl.publish → PRODUCER covering the ack (`messaging.message.id` =
+//    the ack's `seq`, `messaging.nats.stream`); consumers' `consume()` → iterator style, or
+//    callback style with `consume({ callback })`;
+//    `fetch()` → batch pull (`receive <subject>` from the call until the fetch iterator ends,
+//    one link per message); `next()` → batch pull of one.
+// Subjects of the protocol itself (`$JS.`, `$KV.`, `$O.`, `$SYS.`, `_INBOX.`: JetStream API calls,
+// acknowledgements, flow control, replies) make no span, like pmx-io.
+
+const path = require('path')
+const c = require('./common')
+const ctx = require('../tracing/context')
+
+const SYSTEM = 'nats'
+const META = Symbol('pmx-nats-iter')
+const factories = new Map() // package dir → headers()
+
+function internal (subject) {
+  return typeof subject !== 'string' || /^(\$JS\.|\$KV\.|\$O\.|\$SYS\.|_INBOX\.)/.test(subject)
+}
+
+function peer (nc) {
+  const tags = {}
+  try {
+    const s = typeof nc.getServer === 'function' ? nc.getServer() : ''
+    if (s) {
+      const i = s.lastIndexOf(':')
+      if (i > 0) {
+        tags['net.peer.name'] = s.slice(0, i).replace(/^\[|\]$/g, '')
+        tags['net.peer.port'] = s.slice(i + 1)
+      } else tags['net.peer.name'] = s
+    }
+  } catch (_) {}
+  return tags
+}
+
+function hdrsObject (h) {
+  if (!h) return null
+  const o = {}
+  try {
+    if (typeof h[Symbol.iterator] === 'function') {
+      for (const [k, v] of h) o[String(k).toLowerCase()] = Array.isArray(v) ? v[0] : v
+      return o
+    }
+    if (typeof h.keys === 'function' && typeof h.get === 'function') {
+      for (const k of h.keys()) o[String(k).toLowerCase()] = h.get(k)
+      return o
+    }
+  } catch (_) {}
+  return null
+}
+
+function extractMsg (msg) {
+  try {
+    const o = hdrsObject(msg && msg.headers)
+    return o ? c.extract(o) : null
+  } catch (_) {
+    return null
+  }
+}
+
+function supportsHeaders (nc) {
+  try {
+    const info = nc && nc.info
+    return !(info && info.headers === false)
+  } catch (_) {
+    return true
+  }
+}
+
+/** A copy of `options` whose `headers` (a new headers object) carries the trace headers. */
+function withTrace (options, span, factory, nc) {
+  if (!span) return options
+  if (!factory || !supportsHeaders(nc)) {
+    span.setAttribute('messaging.pmx.propagation', 'unsupported')
+    return options
+  }
+  const out = Object.assign({}, options || {})
+  const h = factory()
+  const present = new Set()
+  const src = out.headers
+  try {
+    if (src && typeof src[Symbol.iterator] === 'function') {
+      for (const [k, vs] of src) {
+        present.add(String(k).toLowerCase())
+        for (const v of (Array.isArray(vs) ? vs : [vs])) h.append(k, v)
+      }
+    }
+  } catch (_) {}
+  for (const [k, v] of c.traceHeaders(span)) if (!present.has(k)) h.append(k, v)
+  out.headers = h
+  return out
+}
+
+function bodySize (data) {
+  return c.bodySize(data)
+}
+
+// -- core connection --------------------------------------------------------------------------
+
+function wrapPublish (factory) {
+  return (orig) => function publish (subject, data, options) {
+    if (!c.msgTracer() || internal(subject) || c.nested() || nestedHere()) return orig.apply(this, arguments)
+    const tags = Object.assign({}, peer(this))
+    const size = bodySize(data)
+    if (size !== undefined) tags['messaging.message.body.size'] = String(size)
+    const span = c.producer(SYSTEM, subject, tags)
+    if (!span) return orig.apply(this, arguments)
+    let opts = options
+    try {
+      opts = withTrace(options, span, factory, this)
+    } catch (_) {}
+    try {
+      const r = c.outer(orig, this, [subject, data, opts])
+      c.end(span)
+      return r
+    } catch (err) {
+      c.end(span, err)
+      throw err
+    }
+  }
+}
+
+function wrapRequest (factory, many) {
+  return (orig) => function request (subject, data, opts) {
+    if (!c.msgTracer() || internal(subject) || c.nested() || nestedHere()) return orig.apply(this, arguments)
+    const tags = Object.assign({}, peer(this))
+    const size = bodySize(data)
+    if (size !== undefined) tags['messaging.message.body.size'] = String(size)
+    const span = c.producer(SYSTEM, subject, tags, 'CLIENT', 'request')
+    if (!span) return orig.apply(this, arguments)
+    const args = Array.prototype.slice.call(arguments)
+    try {
+      args[2] = withTrace(opts, span, factory, this)
+      if (args.length < 3) args.length = 3
+    } catch (_) {}
+    let p
+    try {
+      p = c.outer(orig, this, args)
+    } catch (err) {
+      c.end(span, err)
+      throw err
+    }
+    if (c.isThenable(p)) return c.settle(p, span)
+    c.end(span)
+    return p
+  }
+}
+
+function wrapSubscribe (orig) {
+  return function subscribe (subject, opts) {
+    if (!c.msgTracer() || internal(subject)) return orig.apply(this, arguments)
+    const nc = this
+    const o = opts && typeof opts === 'object' ? opts : {}
+    const group = typeof o.queue === 'string' && o.queue ? o.queue : undefined
+    if (typeof o.callback === 'function') {
+      const cb = o.callback
+      const copy = Object.assign({}, o, {
+        callback: function (err, msg) {
+          if (err || !msg || !c.msgTracer()) return cb.apply(this, arguments)
+          const span = consumerSpan(nc, msg, group)
+          return c.runConsumer(span, cb, this, arguments)
+        }
+      })
+      return orig.call(this, subject, copy)
+    }
+    const sub = orig.apply(this, arguments)
+    mark(sub, { mode: 'iterator', nc, group })
+    return sub
+  }
+}
+
+function consumerSpan (nc, msg, group, extra) {
+  const tags = Object.assign({}, peer(nc), extra || {})
+  if (group) tags['messaging.consumer.group.name'] = group
+  const size = bodySize(msg && msg.data)
+  if (size !== undefined) tags['messaging.message.body.size'] = String(size)
+  return c.consumer(SYSTEM, msg && msg.subject, extractMsg(msg), tags)
+}
+
+// -- JetStream --------------------------------------------------------------------------------
+
+// JetStream publish runs nc.request asynchronously: a store flag suppresses that inner span
+function nestedHere () {
+  const s = ctx.current()
+  return Boolean(s && s.itg === SYSTEM)
+}
+
+function wrapJsPublish (factory) {
+  return (orig) => function publish (subject, data, opts) {
+    if (!c.msgTracer() || internal(subject) || c.nested() || nestedHere()) return orig.apply(this, arguments)
+    const nc = this.nc
+    const tags = Object.assign({}, nc ? peer(nc) : {})
+    const size = bodySize(data)
+    if (size !== undefined) tags['messaging.message.body.size'] = String(size)
+    const span = c.producer(SYSTEM, subject, tags)
+    if (!span) return orig.apply(this, arguments)
+    const args = Array.prototype.slice.call(arguments)
+    try {
+      args[2] = withTrace(opts, span, factory, nc)
+      if (args.length < 3) args.length = 3
+    } catch (_) {}
+    let p
+    try {
+      const store = ctx.current()
+      p = ctx.run({ span, bag: span.bag, server: store && store.server, itg: SYSTEM }, orig, this, args)
+    } catch (err) {
+      c.end(span, err)
+      throw err
+    }
+    if (!c.isThenable(p)) {
+      c.end(span)
+      return p
+    }
+    return c.settle(p, span, (ack) => {
+      if (ack && ack.seq !== undefined) span.setAttribute('messaging.message.id', String(ack.seq))
+      if (ack && ack.stream) span.setAttribute('messaging.nats.stream', String(ack.stream))
+    })
+  }
+}
+
+function consumerName (consumer) {
+  try {
+    return consumer.name || (consumer._info && consumer._info.name) || undefined
+  } catch (_) {
+    return undefined
+  }
+}
+
+function jsTags (msg) {
+  const tags = {}
+  try {
+    const info = msg && msg.info
+    if (info && info.stream) tags['messaging.nats.stream'] = String(info.stream)
+    if (msg && msg.seq !== undefined) tags['messaging.message.id'] = String(msg.seq)
+  } catch (_) {}
+  return tags
+}
+
+// nats 2.x ordered consumers deliver through an internal callback given to an inner consume():
+// those callbacks are marked and not traced (the app's callback or iterator is)
+const INTERNAL = Symbol('pmx-nats-internal')
+
+function wrapInternalHandler (orig) {
+  return function internalHandler () {
+    const fn = orig.apply(this, arguments)
+    if (typeof fn === 'function') {
+      try {
+        Object.defineProperty(fn, INTERNAL, { value: true })
+      } catch (_) {}
+    }
+    return fn
+  }
+}
+
+/** consume({ callback }): callback style, one CONSUMER span per message (ITG-023). */
+function tracedCallback (cb, meta) {
+  const w = function (m) {
+    if (!m || !c.msgTracer()) return cb.apply(this, arguments)
+    const span = consumerSpan(meta.nc, m, meta.group, jsTags(m))
+    return c.runConsumer(span, cb, this, arguments)
+  }
+  return c.copyMeta(w, cb)
+}
+
+function wrapConsume (orig) {
+  return function consume (opts) {
+    if (c.msgTracer() && opts && typeof opts === 'object' && typeof opts.callback === 'function' && !opts.callback[INTERNAL]) {
+      const meta = { nc: this.api && this.api.nc, group: consumerName(this) }
+      const args = Array.prototype.slice.call(arguments)
+      args[0] = Object.assign({}, opts, { callback: tracedCallback(opts.callback, meta) })
+      return orig.apply(this, args)
+    }
+    const p = orig.apply(this, arguments)
+    if (!c.msgTracer()) return p
+    const meta = { mode: 'iterator', nc: this.api && this.api.nc, group: consumerName(this), js: true }
+    if (c.isThenable(p)) return p.then((it) => { mark(it, meta); return it })
+    mark(p, meta)
+    return p
+  }
+}
+
+function wrapFetch (orig) {
+  return function fetch () {
+    if (!c.msgTracer() || c.nested()) return orig.apply(this, arguments)
+    const meta = { mode: 'fetch', nc: this.api && this.api.nc, group: consumerName(this), js: true }
+    const tags = Object.assign({}, meta.nc ? peer(meta.nc) : {})
+    if (meta.group) tags['messaging.consumer.group.name'] = meta.group
+    // created at the call (parent: the current span), kept only when messages came back
+    meta.span = c.batchConsumer(SYSTEM, undefined, null, tags, 'receive')
+    const p = orig.apply(this, arguments)
+    if (c.isThenable(p)) return p.then((it) => { mark(it, meta); return it })
+    mark(p, meta)
+    return p
+  }
+}
+
+function wrapNext (orig) {
+  return function next () {
+    if (!c.msgTracer()) return orig.apply(this, arguments)
+    const nc = this.api && this.api.nc
+    const group = consumerName(this)
+    const tags = Object.assign({}, nc ? peer(nc) : {})
+    if (group) tags['messaging.consumer.group.name'] = group
+    // created now (covers the call), kept only when a message came back (ITG-023)
+    const span = c.batchConsumer(SYSTEM, undefined, null, tags, 'receive')
+    let p
+    try {
+      // 3.x implements next() with fetch(): that inner call makes no span of its own
+      p = c.outer(orig, this, arguments)
+    } catch (err) {
+      c.end(span, err)
+      throw err
+    }
+    if (!span || !c.isThenable(p)) return p
+    return p.then((msg) => {
+      if (msg) {
+        try {
+          if (msg.subject) {
+            span.setAttribute('messaging.destination.name', msg.subject)
+            span.updateName('receive ' + msg.subject)
+          }
+          Object.entries(jsTags(msg)).forEach(([k, v]) => span.setAttribute(k, v))
+          c.addLink(span, extractMsg(msg))
+        } catch (_) {}
+        c.end(span)
+      }
+      return msg
+    }, (err) => {
+      c.end(span, err)
+      throw err
+    })
+  }
+}
+
+// -- iterators (subscriptions, consume(), fetch()) --------------------------------------------
+
+function mark (it, meta) {
+  if (!it || typeof it !== 'object') return
+  try {
+    Object.defineProperty(it, META, { value: meta, configurable: true })
+  } catch (_) {}
+}
+
+function wrapAsyncIterator (orig) {
+  return function asyncIterator () {
+    const it = orig.apply(this, arguments)
+    const meta = this && this[META]
+    if (!meta || !it || typeof it.next !== 'function' || !c.msgTracer()) return it
+    return meta.mode === 'fetch' ? fetchIterator(it, meta) : traceIterator(it, meta)
+  }
+}
+
+/** Iterator style (ITG-023): one CONSUMER span per message, current for the loop body. */
+function traceIterator (it, meta) {
+  const st = c.iterState({})
+  const out = {
+    next () {
+      st.nextRequested()
+      let p
+      try {
+        p = it.next.apply(it, arguments)
+      } catch (err) {
+        st.closed(true)
+        throw err
+      }
+      return Promise.resolve(p).then((r) => {
+        if (r && r.done) st.closed(true)
+        else if (r && r.value) {
+          try {
+            st.delivered(consumerSpan(meta.nc, r.value, meta.group, meta.js ? jsTags(r.value) : undefined))
+          } catch (_) {}
+        }
+        return r
+      }, (err) => {
+        // the library call itself failed: the open span (if any) is closed by closed()
+        st.closed(true)
+        throw err
+      })
+    },
+    return () {
+      st.closed(true)
+      return typeof it.return === 'function' ? it.return.apply(it, arguments) : Promise.resolve({ done: true, value: undefined })
+    },
+    throw () {
+      st.closed(true)
+      return typeof it.throw === 'function' ? it.throw.apply(it, arguments) : Promise.reject(arguments[0])
+    },
+    [Symbol.asyncIterator] () {
+      return this
+    }
+  }
+  return out
+}
+
+/** Batch pull over an iterator (`fetch()`): a `receive` span from the call to the end, with links. */
+function fetchIterator (it, meta) {
+  const span = meta.span
+  const subjects = new Set()
+  let count = 0
+  let finished = false
+  const finish = (err) => {
+    if (finished || !span) return
+    finished = true
+    if (err) return c.end(span, err)
+    if (count === 0) return
+    try {
+      if (subjects.size === 1) {
+        const s = Array.from(subjects)[0]
+        span.setAttribute('messaging.destination.name', s)
+        span.updateName('receive ' + s)
+      }
+      if (count >= 2) span.setAttribute('messaging.batch.message_count', String(count))
+    } catch (_) {}
+    c.end(span)
+  }
+  return {
+    next () {
+      return Promise.resolve(it.next.apply(it, arguments)).then((r) => {
+        if (r && r.done) finish()
+        else if (r && r.value) {
+          count++
+          try {
+            if (r.value.subject) subjects.add(r.value.subject)
+            c.addLink(span, extractMsg(r.value))
+          } catch (_) {}
+        }
+        return r
+      }, (err) => {
+        finish(err)
+        throw err
+      })
+    },
+    return () {
+      finish()
+      return typeof it.return === 'function' ? it.return.apply(it, arguments) : Promise.resolve({ done: true, value: undefined })
+    },
+    throw () {
+      finish()
+      return typeof it.throw === 'function' ? it.throw.apply(it, arguments) : Promise.reject(arguments[0])
+    },
+    [Symbol.asyncIterator] () {
+      return this
+    }
+  }
+}
+
+// -- patches ----------------------------------------------------------------------------------
+
+function headersFactory (dir) {
+  if (factories.has(dir)) return factories.get(dir)
+  let f = null
+  for (const rel of ['lib/nats-base-client/headers.js', 'lib/headers.js']) {
+    try {
+      const m = require(path.join(dir, rel))
+      if (typeof m.headers === 'function') {
+        f = m.headers
+        break
+      }
+    } catch (_) {}
+  }
+  factories.set(dir, f)
+  return f
+}
+
+/** nats/lib/nats-base-client/nats.js (2.x) and @nats-io/nats-core/lib/nats.js (3.x). */
+function patchConnection (exports, info) {
+  const Impl = exports && exports.NatsConnectionImpl
+  const proto = Impl && Impl.prototype
+  if (!proto || typeof proto.publish !== 'function') {
+    c.disabled(info.pkg, 'NatsConnectionImpl.publish not found')
+    return
+  }
+  const factory = headersFactory(info.dir)
+  c.wrap(proto, 'publish', wrapPublish(factory))
+  if (typeof proto.request === 'function') c.wrap(proto, 'request', wrapRequest(factory, false))
+  if (typeof proto.requestMany === 'function') c.wrap(proto, 'requestMany', wrapRequest(factory, true))
+  if (typeof proto.subscribe === 'function') c.wrap(proto, 'subscribe', wrapSubscribe)
+}
+
+/** queued_iterator.js: QueuedIteratorImpl.prototype[Symbol.asyncIterator]. */
+function patchIterator (exports) {
+  const Q = exports && exports.QueuedIteratorImpl
+  const proto = Q && Q.prototype
+  if (!proto || typeof proto[Symbol.asyncIterator] !== 'function') return
+  c.wrap(proto, Symbol.asyncIterator, wrapAsyncIterator)
+}
+
+/** nats/lib/jetstream/jsclient.js and @nats-io/jetstream/lib/jsclient.js. */
+function patchJsClient (exports, info) {
+  const Impl = exports && exports.JetStreamClientImpl
+  const proto = Impl && Impl.prototype
+  if (!proto || typeof proto.publish !== 'function') {
+    c.disabled(info.pkg, 'JetStreamClientImpl.publish not found')
+    return
+  }
+  // the headers factory lives in the core package: nats (same package) or @nats-io/nats-core
+  let factory = headersFactory(info.dir)
+  if (!factory) {
+    try {
+      const core = require.resolve('@nats-io/nats-core', { paths: [info.dir] })
+      factory = headersFactory(path.dirname(path.dirname(core)))
+    } catch (_) {}
+  }
+  c.wrap(proto, 'publish', wrapJsPublish(factory))
+}
+
+/** jetstream consumer.js: PullConsumerImpl / OrderedPullConsumerImpl consume, fetch, next. */
+function patchConsumer (exports) {
+  for (const name of ['PullConsumerImpl', 'OrderedPullConsumerImpl']) {
+    const K = exports && exports[name]
+    const proto = K && K.prototype
+    if (!proto) continue
+    if (typeof proto.consume === 'function') c.wrap(proto, 'consume', wrapConsume)
+    if (typeof proto.internalHandler === 'function') c.wrap(proto, 'internalHandler', wrapInternalHandler)
+    if (typeof proto.fetch === 'function') c.wrap(proto, 'fetch', wrapFetch)
+    if (typeof proto.next === 'function') c.wrap(proto, 'next', wrapNext)
+  }
+}
+
+module.exports = { patchConnection, patchIterator, patchJsClient, patchConsumer }