more tests

DataDog · Sep 23, 2024 · 7f56ce0 · 7f56ce0
1 parent 8186fbd
commit 7f56ce0
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 20 deletions.
diff --git a/packages/datadog-plugin-protobufjs/test/index.spec.js b/packages/datadog-plugin-protobufjs/test/index.spec.js
@@ -14,6 +14,7 @@ const {
 } = require('../../dd-trace/src/constants')
 const sinon = require('sinon')
 const { loadMessage } = require('./helpers')
+const { SchemaBuilder } = require('../../dd-trace/src/datastreams/schemas/schema_builder')
 
 const schemas = JSON.parse(fs.readFileSync(path.join(__dirname, 'schemas/expected_schemas.json'), 'utf8'))
 const MESSAGE_SCHEMA_DEF = schemas.MESSAGE_SCHEMA_DEF
@@ -271,5 +272,60 @@ describe('Plugin', () => {
         expect(span.context()._tags).to.have.property(SCHEMA_WEIGHT, 1)
       })
     })
+
+    describe('during schema sampling', function () {
+      let cacheSetSpy
+      let cacheGetSpy
+
+      beforeEach(() => {
+        const cache = SchemaBuilder.getCache()
+        cache.clear()
+        cacheSetSpy = sinon.spy(cache, 'set')
+        cacheGetSpy = sinon.spy(cache, 'get')
+      })
+
+      afterEach(() => {
+        cacheSetSpy.restore()
+        cacheGetSpy.restore()
+      })
+
+      it('should use the schema cache and not re-extract an already sampled schema', async () => {
+        const loadedMessages = await loadMessage(protobuf, 'MyMessage')
+
+        tracer.trace('message_pb2.serialize', span => {
+          loadedMessages.MyMessage.type.encode(loadedMessages.MyMessage.instance).finish()
+
+          expect(span._name).to.equal('message_pb2.serialize')
+
+          expect(compareJson(MESSAGE_SCHEMA_DEF, span)).to.equal(true)
+          expect(span.context()._tags).to.have.property(SCHEMA_TYPE, 'protobuf')
+          expect(span.context()._tags).to.have.property(SCHEMA_NAME, 'MyMessage')
+          expect(span.context()._tags).to.have.property(SCHEMA_OPERATION, 'serialization')
+          expect(span.context()._tags).to.have.property(SCHEMA_ID, MESSAGE_SCHEMA_ID)
+          expect(span.context()._tags).to.have.property(SCHEMA_WEIGHT, 1)
+
+          // we sampled 1 schema with 1 subschema, so the constructor should've only been called twice
+          expect(cacheSetSpy.callCount).to.equal(2)
+          expect(cacheGetSpy.callCount).to.equal(2)
+        })
+
+        tracer.trace('message_pb2.serialize', span => {
+          loadedMessages.MyMessage.type.encode(loadedMessages.MyMessage.instance).finish()
+
+          expect(span._name).to.equal('message_pb2.serialize')
+
+          expect(compareJson(MESSAGE_SCHEMA_DEF, span)).to.equal(true)
+          expect(span.context()._tags).to.have.property(SCHEMA_TYPE, 'protobuf')
+          expect(span.context()._tags).to.have.property(SCHEMA_NAME, 'MyMessage')
+          expect(span.context()._tags).to.have.property(SCHEMA_OPERATION, 'serialization')
+          expect(span.context()._tags).to.have.property(SCHEMA_ID, MESSAGE_SCHEMA_ID)
+          expect(span.context()._tags).to.have.property(SCHEMA_WEIGHT, 1)
+
+          // ensure schema was sampled and returned via the cache, so no extra cache set calls were needed, only gets
+          expect(cacheSetSpy.callCount).to.equal(2)
+          expect(cacheGetSpy.callCount).to.equal(3)
+        })
+      })
+    })
   })
 })
diff --git a/packages/dd-trace/src/datastreams/schemas/schema_builder.js b/packages/dd-trace/src/datastreams/schemas/schema_builder.js
@@ -10,7 +10,30 @@ class SchemaBuilder {
   constructor (iterator) {
     this.schema = new OpenApiSchema()
     this.iterator = iterator
-    this.proerties = 0
+    this.properties = 0
+  }
+
+  static getCache () {
+    return CACHE
+  }
+
+  static getSchemaDefinition (schema) {
+    const noNones = convertToJsonCompatible(schema)
+    const definition = jsonStringify(noNones)
+    const id = fnv64(Buffer.from(definition, 'utf-8')).toString()
+    return new Schema(definition, id)
+  }
+
+  static getSchema (schemaName, iterator, builder) {
+    if (!CACHE.has(schemaName)) {
+      CACHE.set(schemaName, (builder ?? new SchemaBuilder(iterator)).build())
+    }
+    return CACHE.get(schemaName)
+  }
+
+  build () {
+    this.iterator.iterateOverSchema(this)
+    return this.schema
   }
 
   addProperty (schemaName, fieldName, isArray, type, description, ref, format, enumValues) {
@@ -26,18 +49,6 @@ class SchemaBuilder {
     return true
   }
 
-  build () {
-    this.iterator.iterateOverSchema(this)
-    return this.schema
-  }
-
-  static getSchemaDefinition (schema) {
-    const noNones = convertToJsonCompatible(schema)
-    const definition = jsonStringify(noNones)
-    const id = fnv64(Buffer.from(definition, 'utf-8')).toString()
-    return new Schema(definition, id)
-  }
-
   shouldExtractSchema (schemaName, depth) {
     if (depth > maxDepth) {
       return false
@@ -48,13 +59,6 @@ class SchemaBuilder {
     this.schema.components.schemas[schemaName] = new OpenApiSchema.SCHEMA()
     return true
   }
-
-  static getSchema (schemaName, iterator, builder) {
-    if (!CACHE.has(schemaName)) {
-      CACHE.set(schemaName, (builder ?? new SchemaBuilder(iterator)).build())
-    }
-    return CACHE.get(schemaName)
-  }
 }
 
 class OpenApiSchema {