WIP

src/reporter.ts

@@ -12,7 +12,13 @@ import {
 import * as logger from './logger';
 import { HttpSpansAgent } from './httpSpansAgent';
 import { payloadStringify } from './utils/payloadStringify';
-import { decodeHttpBody } from './spans/awsSpan';
+import {
+  decodeHttpBody,
+  FUNCTION_SPAN,
+  getSpanInfo,
+  getSpanMetadata,
+  spansPrioritySorter,
+} from './spans/awsSpan';
 import untruncateJson from './tools/untrancateJson';
 export const NUMBER_OF_SPANS_IN_REPORT_OPTIMIZATION = 200;
 
@@ -115,6 +121,40 @@ export function scrubSpans(resultSpans: any[]) {
   return resultSpans.filter((span) => safeExecute(scrubSpan, 'Failed to scrub span')(span));
 }
 
+export function getPrioritizedSpans(spans: any[], maxSendBytes: number): any[] {
+  logger.debug('Using smart spans prioritization');
+  spans.sort(spansPrioritySorter);
+  let currentSize = 0;
+  const spansToSendSizes = {};
+  const spansToSend = {};
+
+  // First we try to take only the spans metadata
+  for (let index = 0; index < spans.length; index++) {
+    const spanMetadata = getSpanMetadata(spans[index]);
+    if (spanMetadata == undefined) continue;
+    const spanMetadataSize = getJSONBase64Size(spanMetadata);
+
+    if (currentSize + spanMetadata < maxSendBytes) {
+      spansToSendSizes[index] = spanMetadataSize;
+      spansToSend[index] = spanMetadata;
+      currentSize += spanMetadataSize;
+    }
+  }
+
+  // Replace metadata span with full spans
+  for (let index = 0; index < spans.length; index++) {
+    const spanSize = getJSONBase64Size(spans[index]);
+    const spanMetadataSize = spansToSendSizes[index];
+
+    if (currentSize + spanSize - spanMetadataSize < maxSendBytes) {
+      spansToSend[index] = spans[index];
+      currentSize += spanSize - spanMetadataSize;
+    }
+  }
+
+  return Object.values(spansToSend);
+}
+
 // We muted the spans itself to keep the memory footprint of the tracer to a minimum
 export const forgeAndScrubRequestBody = (spans, maxSendBytes): string | undefined => {
   const start = new Date().getTime();
@@ -124,13 +164,9 @@ export const forgeAndScrubRequestBody = (spans, maxSendBytes): string | undefine
     logger.debug(
       `Starting trim spans [${spans.length}] bigger than: [${maxSendBytes}] before send`
     );
-
-    const functionEndSpan = spans.pop();
-    spans.sort((a, b) => (spanHasErrors(a) ? -1 : spanHasErrors(b) ? 1 : 0));
-    let totalSize = getJSONBase64Size(functionEndSpan) + getJSONBase64Size(spans);
-    while (totalSize > maxSendBytes && spans.length > 0)
-      totalSize -= getJSONBase64Size(spans.pop());
-    spans.push(functionEndSpan);
+    if (getJSONBase64Size(spans) > maxSendBytes && spans.length > 0) {
+      spans = getPrioritizedSpans(spans, maxSendBytes);
+    }
   }
   spans = scrubSpans(spans);
   if (originalSize - spans.length > 0) {

src/spans/awsSpan.ts

@@ -34,11 +34,13 @@ import {
   safeExecute,
   SENDING_TIME_ID_KEY,
   setWarm,
+  spanHasErrors,
   TRANSACTION_ID_KEY,
 } from '../utils';
 import { payloadStringify, shallowMask, truncate } from '../utils/payloadStringify';
 import { Utf8Utils } from '../utils/utf8Utils';
 import { getW3CMessageId } from '../utils/w3cUtils';
+import { int } from 'aws-sdk/clients/datapipeline';
 
 export const HTTP_SPAN = 'http';
 export const FUNCTION_SPAN = 'function';
@@ -66,6 +68,78 @@ export const getSpanInfo = (): SpanInfo => {
   return { traceId, tracer, logGroupName, logStreamName };
 };
 
+export const getSpanPriority = (span): number => {
+  if (span.type === FUNCTION_SPAN) {
+    return 0;
+  }
+  if (span.type === ENRICHMENT_SPAN) {
+    return 1;
+  }
+  if (spanHasErrors(span)) {
+    return 2;
+  }
+  return 3;
+};
+
+export const spansPrioritySorter = (span1: any, span2: any): number => {
+  const span1Priority = getSpanPriority(span1);
+  const span2Priority = getSpanPriority(span2);
+
+  if (span1Priority < span2Priority) {
+    return -1;
+  } else if (span1Priority > span2Priority) {
+    return 1;
+  } else {
+    return 0;
+  }
+};
+
+export const getSpanMetadata = (span: any): any => {
+  const spanCopy = span;
+  spanCopy['isMetadata'] = true;
+
+  if (spanCopy.type === FUNCTION_SPAN) {
+    return {};
+  } else if (spanCopy.type === ENRICHMENT_SPAN) {
+    return {};
+  } else if (spanCopy.type === HTTP_SPAN) {
+    spanCopy?.request?.remove('headers');
+    spanCopy?.request?.remove('body');
+    spanCopy?.response?.remove('headers');
+    spanCopy?.response?.remove('body');
+    return {};
+  } else if (spanCopy.type === MONGO_SPAN) {
+    spanCopy?.remove('request');
+    spanCopy?.remove('response');
+    return spanCopy;
+  } else if (spanCopy.type === REDIS_SPAN) {
+    spanCopy?.remove('requestArgs');
+    spanCopy?.remove('response');
+    return spanCopy;
+  } else if (spanCopy.type === NEO4J_SPAN) {
+    spanCopy?.remove('summary');
+    spanCopy?.remove('query');
+    spanCopy?.remove('values');
+    spanCopy?.remove('response');
+    return spanCopy;
+  } else if (
+    spanCopy.type === PG_SPAN ||
+    spanCopy.type === MYSQL_SPAN ||
+    spanCopy.type === MSSQL_SPAN
+  ) {
+    spanCopy?.remove('query');
+    spanCopy?.remove('values');
+    spanCopy?.remove('response');
+    return spanCopy;
+  } else if (spanCopy.type === PRISMA_SPAN) {
+    spanCopy?.remove('queryArgs');
+    spanCopy?.remove('result');
+    return spanCopy;
+  }
+  logger.warn(`Got unknown span type: ${spanCopy.type}`);
+  return {};
+};
+
 export const getCurrentTransactionId = (): string => {
   return getSpanInfo().traceId.transactionId;
 };