Skip to content

Commit

Permalink
feat: Removed references to segment.transaction and instead get
Browse files Browse the repository at this point in the history
transaction either from tracer or just within the function scope.
  • Loading branch information
bizob2828 committed Oct 11, 2024
1 parent 13edec5 commit e9af6d1
Show file tree
Hide file tree
Showing 71 changed files with 533 additions and 341 deletions.
5 changes: 3 additions & 2 deletions api.js
Original file line number Diff line number Diff line change
Expand Up @@ -1524,12 +1524,13 @@ API.prototype.getTraceMetadata = function getTraceMetadata() {
const metadata = {}

const segment = this.agent.tracer.getSegment()
if (!segment) {
const transaction = this.agent.tracer.getTransaction()
if (!(segment || transaction)) {
logger.debug('No transaction found when calling API#getTraceMetadata')
} else if (!this.agent.config.distributed_tracing.enabled) {
logger.debug('Distributed tracing disabled when calling API#getTraceMetadata')
} else {
metadata.traceId = segment.transaction.traceId
metadata.traceId = transaction.traceId

const spanId = segment.getSpanId()
if (spanId) {
Expand Down
3 changes: 2 additions & 1 deletion lib/agent.js
Original file line number Diff line number Diff line change
Expand Up @@ -834,6 +834,7 @@ Agent.prototype._listenForConfigChanges = function _listenForConfigChanges() {
*/
Agent.prototype.getLinkingMetadata = function getLinkingMetadata() {
const segment = this.tracer.getSegment()
const transaction = this.tracer.getTransaction()
const config = this.config

const linkingMetadata = {
Expand All @@ -843,7 +844,7 @@ Agent.prototype.getLinkingMetadata = function getLinkingMetadata() {
}

if (config.distributed_tracing.enabled && segment) {
linkingMetadata['trace.id'] = segment.transaction.traceId
linkingMetadata['trace.id'] = transaction.traceId
const spanId = segment.getSpanId()
if (spanId) {
linkingMetadata['span.id'] = spanId
Expand Down
4 changes: 2 additions & 2 deletions lib/attributes.js
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ class Attributes {
* @param {string} scope
* The scope of the attributes this will collect. Must be `transaction` or
* `segment`.
* @param {number} [limit=Infinity]
* @param {number} [limit]
* The maximum number of attributes to retrieve for each destination.
*/
constructor(scope, limit = Infinity) {
Expand Down Expand Up @@ -104,7 +104,7 @@ class Attributes {
* @param {DESTINATIONS} destinations - The default destinations for this key.
* @param {string} key - The attribute name.
* @param {string} value - The attribute value.
* @param {boolean} [truncateExempt=false] - Flag marking value exempt from truncation
* @param {boolean} [truncateExempt] - Flag marking value exempt from truncation
*/
addAttribute(destinations, key, value, truncateExempt = false) {
if (this.attributeCount + 1 > this.limit) {
Expand Down
2 changes: 1 addition & 1 deletion lib/collector/remote-method.js
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ RemoteMethod.prototype._reportDataUsage = function reportDataUsage(sent, receive
* you're doing it wrong.
*
* @param {object} payload Serializable payload.
* @param {object} [nrHeaders=null] NR request headers from connect response.
* @param {object} [nrHeaders] NR request headers from connect response.
* @param {Function} callback What to do next. Gets passed any error.
*/
RemoteMethod.prototype.invoke = function invoke(payload, nrHeaders, callback) {
Expand Down
11 changes: 8 additions & 3 deletions lib/db/parsed-statement.js
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,9 @@ function ParsedStatement(type, operation, collection, raw) {
}
}

ParsedStatement.prototype.recordMetrics = function recordMetrics(segment, scope) {
ParsedStatement.prototype.recordMetrics = function recordMetrics(segment, scope, transaction) {
const duration = segment.getDurationInMillis()
const exclusive = segment.getExclusiveDurationInMillis()
const transaction = segment.transaction
const type = transaction.isWeb() ? DB.WEB : DB.OTHER
const thisTypeSlash = this.type + '/'
const operation = DB.OPERATION + '/' + thisTypeSlash + this.operation
Expand Down Expand Up @@ -68,7 +67,13 @@ ParsedStatement.prototype.recordMetrics = function recordMetrics(segment, scope)
}

if (this.raw) {
transaction.agent.queries.add(segment, this.type.toLowerCase(), this.raw, this.trace)
transaction.agent.queries.add({
segment,
transaction,
type: this.type.toLowerCase(),
query: this.raw,
trace: this.trace
})
}
}

Expand Down
6 changes: 3 additions & 3 deletions lib/db/query-sample.js
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ QuerySample.prototype.merge = function merge(sample) {
}

QuerySample.prototype.prepareJSON = function prepareJSON(done) {
const transaction = this.trace.segment.transaction
const transaction = this.trace.transaction
const sample = this
const trace = sample.trace

Expand All @@ -56,7 +56,7 @@ QuerySample.prototype.prepareJSON = function prepareJSON(done) {
}

QuerySample.prototype.prepareJSONSync = function prepareJSONSync() {
const transaction = this.trace.segment.transaction
const transaction = this.trace.transaction
const sample = this
const trace = sample.trace

Expand Down Expand Up @@ -99,7 +99,7 @@ QuerySample.prototype.getParams = function getParams() {
}

if (this.tracer.config.distributed_tracing.enabled) {
this.trace.segment.transaction.addDistributedTraceIntrinsics(params)
this.trace.transaction.addDistributedTraceIntrinsics(params)
}

return params
Expand Down
8 changes: 4 additions & 4 deletions lib/db/query-trace-aggregator.js
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ class QueryTraceAggregator extends Aggregator {
}
}

add(segment, type, query, trace) {
add({ segment, transaction, type, query, trace }) {
const ttConfig = this.config.transaction_tracer

// If DT is enabled and the segment is part of a sampled transaction
Expand All @@ -57,12 +57,12 @@ class QueryTraceAggregator extends Aggregator {
let slowQuery
switch (ttConfig.record_sql) {
case 'raw':
slowQuery = new SlowQuery(segment, type, query, trace)
slowQuery = new SlowQuery({ segment, transaction, type, query, trace })
logger.trace('recording raw sql')
segment.addAttribute('sql', slowQuery.query, true)
break
case 'obfuscated':
slowQuery = new SlowQuery(segment, type, query, trace)
slowQuery = new SlowQuery({ transaction, segment, type, query, trace })
logger.trace('recording obfuscated sql')
segment.addAttribute('sql_obfuscated', slowQuery.obfuscated, true)
break
Expand All @@ -78,7 +78,7 @@ class QueryTraceAggregator extends Aggregator {
return
}

slowQuery = slowQuery || new SlowQuery(segment, type, query, trace)
slowQuery = slowQuery || new SlowQuery({ segment, transaction, type, query, trace })

segment.addAttribute('backtrace', slowQuery.trace)

Expand Down
3 changes: 2 additions & 1 deletion lib/db/slow-query.js
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,15 @@ const crypto = require('crypto')
const path = require('path')
const NR_ROOT = path.resolve(__dirname, '..')

function SlowQuery(segment, type, query, trace) {
function SlowQuery({ segment, transaction, type, query, trace }) {
this.obfuscated = obfuscate(query, type)
this.normalized = this.obfuscated.replace(/\?\s*,\s*|\s*/g, '')
this.id = normalizedHash(this.normalized)
this.segment = segment
this.query = query
this.metric = segment.name
this.trace = formatTrace(trace)
this.transaction = transaction
this.duration = segment.getDurationInMillis()
}

Expand Down
2 changes: 1 addition & 1 deletion lib/environment.js
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ function clearSetting(name) {
* the provided root.
*
* @param {string} root - Path to start listing packages from.
* @param {Array} [packages=[]] - Array to append found packages to.
* @param {Array} [packages] - Array to append found packages to.
*/
async function listPackages(root, packages = []) {
_log('Listing packages in %s', root)
Expand Down
2 changes: 1 addition & 1 deletion lib/grpc/connection.js
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ class GrpcConnection extends EventEmitter {
*
* @param {object} infiniteTracingConfig config item config.infinite_tracing
* @param {MetricAggregator} metrics metric aggregator, for supportability metrics
* @param {number} [reconnectDelayMs=15000] number of milliseconds to wait before reconnecting
* @param {number} [reconnectDelayMs] number of milliseconds to wait before reconnecting
* for error states that require a reconnect delay.
*/
constructor(infiniteTracingConfig, metrics, reconnectDelayMs = DEFAULT_RECONNECT_DELAY_MS) {
Expand Down
37 changes: 28 additions & 9 deletions lib/instrumentation/aws-sdk/v3/bedrock.js
Original file line number Diff line number Diff line change
Expand Up @@ -70,10 +70,11 @@ function recordEvent({ agent, type, msg }) {
* @param {object} params input params
* @param {Agent} params.agent NR agent instance
* @param {TraceSegment} params.segment active segment
* @param params.transaction
*/
function addLlmMeta({ agent, segment }) {
function addLlmMeta({ agent, segment, transaction }) {
agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount()
segment.transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
// end segment to get a consistent segment duration
// for both the LLM events and the segment
segment.end()
Expand All @@ -92,11 +93,13 @@ function addLlmMeta({ agent, segment }) {
* @param {Error|null} params.err error from request if exists
* @param params.bedrockResponse
* @param params.shim
* @param params.transaction
*/
function recordChatCompletionMessages({
agent,
shim,
segment,
transaction,
bedrockCommand,
bedrockResponse,
err
Expand All @@ -110,6 +113,7 @@ function recordChatCompletionMessages({
agent,
bedrockResponse,
bedrockCommand,
transaction,
segment,
isError: err !== null
})
Expand All @@ -119,6 +123,7 @@ function recordChatCompletionMessages({
segment,
bedrockCommand,
bedrockResponse,
transaction,
index: 0,
completionId: summary.id
})
Expand All @@ -128,6 +133,7 @@ function recordChatCompletionMessages({
const chatCompletionMessage = new LlmChatCompletionMessage({
agent,
segment,
transaction,
bedrockCommand,
bedrockResponse,
isResponse: true,
Expand All @@ -142,7 +148,7 @@ function recordChatCompletionMessages({

if (err) {
const llmError = new LlmError({ bedrockResponse, err, summary })
agent.errors.add(segment.transaction, err, llmError)
agent.errors.add(transaction, err, llmError)
}
}

Expand All @@ -157,8 +163,17 @@ function recordChatCompletionMessages({
* @param {BedrockCommand} params.bedrockCommand parsed input
* @param {Error|null} params.err error from request if exists
* @param params.bedrockResponse
* @param params.transaction
*/
function recordEmbeddingMessage({ agent, shim, segment, bedrockCommand, bedrockResponse, err }) {
function recordEmbeddingMessage({
agent,
shim,
segment,
transaction,
bedrockCommand,
bedrockResponse,
err
}) {
if (shouldSkipInstrumentation(agent.config) === true) {
shim.logger.debug('skipping sending of ai data')
return
Expand All @@ -167,6 +182,7 @@ function recordEmbeddingMessage({ agent, shim, segment, bedrockCommand, bedrockR
const embedding = new LlmEmbedding({
agent,
segment,
transaction,
bedrockCommand,
bedrockResponse,
isError: err !== null
Expand All @@ -175,7 +191,7 @@ function recordEmbeddingMessage({ agent, shim, segment, bedrockCommand, bedrockR
recordEvent({ agent, type: 'LlmEmbedding', msg: embedding })
if (err) {
const llmError = new LlmError({ bedrockResponse, err, embedding })
agent.errors.add(segment.transaction, err, llmError)
agent.errors.add(transaction, err, llmError)
}
}

Expand Down Expand Up @@ -222,12 +238,13 @@ function getBedrockSpec({ commandName }, shim, _original, _name, args) {
return new RecorderSpec({
promise: true,
name: `Llm/${modelType}/Bedrock/${commandName}`,
after: ({ shim, error: err, result: response, segment }) => {
after: ({ shim, error: err, result: response, segment, transaction }) => {
const passThroughParams = {
shim,
err,
response,
segment,
transaction,
bedrockCommand,
modelType
}
Expand All @@ -250,22 +267,23 @@ function getBedrockSpec({ commandName }, shim, _original, _name, args) {
'ai_monitoring.streaming.enabled is set to `false`, stream will not be instrumented.'
)
agent.metrics.getOrCreateMetric(AI.STREAMING_DISABLED).incrementCallCount()
addLlmMeta({ agent, segment })
addLlmMeta({ agent, segment, transaction })
}
}
})
}

function handleResponse({ shim, err, response, segment, bedrockCommand, modelType }) {
function handleResponse({ shim, err, response, segment, transaction, bedrockCommand, modelType }) {
const { agent } = shim
const bedrockResponse = createBedrockResponse({ bedrockCommand, response, err })

addLlmMeta({ agent, segment })
addLlmMeta({ agent, segment, transaction })
if (modelType === 'completion') {
recordChatCompletionMessages({
agent,
shim,
segment,
transaction,
bedrockCommand,
bedrockResponse,
err
Expand All @@ -275,6 +293,7 @@ function handleResponse({ shim, err, response, segment, bedrockCommand, modelTyp
agent,
shim,
segment,
transaction,
bedrockCommand,
bedrockResponse,
err
Expand Down
4 changes: 1 addition & 3 deletions lib/instrumentation/core/globals.js
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,7 @@ function initialize(agent, nodule, name, shim) {
!process.domain &&
process.listenerCount('unhandledRejection') === 0
) {
// If there are no unhandledRejection handlers report the error.
const segment = promise[symbols.context] && promise[symbols.context].getSegment()
const tx = segment && segment.transaction
const tx = promise[symbols.context] && promise[symbols.context].getTransaction()
shim.logger.trace('Captured unhandled rejection for transaction %s', tx && tx.id)
agent.errors.add(tx, error)
}
Expand Down
Loading

0 comments on commit e9af6d1

Please sign in to comment.