From fa137d5a258fe40dd2e487a34958b1679eccb8e8 Mon Sep 17 00:00:00 2001
From: Bartlomiej Obecny <bobecny@gmail.com>
Date: Wed, 28 Oct 2020 15:37:48 +0100
Subject: [PATCH] pg spans disconnected from parent (#225)

* chore: gitignore

* chore: fixing context propagation for koa middleware layer
---
 .../src/koa.ts                                |  9 +++++---
 .../test/koa.test.ts                          | 21 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/plugins/node/opentelemetry-koa-instrumentation/src/koa.ts b/plugins/node/opentelemetry-koa-instrumentation/src/koa.ts
index 73fe921f765..f7be8855f0d 100644
--- a/plugins/node/opentelemetry-koa-instrumentation/src/koa.ts
+++ b/plugins/node/opentelemetry-koa-instrumentation/src/koa.ts
@@ -129,9 +129,12 @@ export class KoaInstrumentation extends BasePlugin<typeof koa> {
       const span = this._tracer.startSpan(metadata.name, {
         attributes: metadata.attributes,
       });
-      const result = await middlewareLayer(context, next);
-      span.end();
-      return result;
+
+      return this._tracer.withSpan(span, async () => {
+        const result = await middlewareLayer(context, next);
+        span.end();
+        return result;
+      });
     };
   }
 }
diff --git a/plugins/node/opentelemetry-koa-instrumentation/test/koa.test.ts b/plugins/node/opentelemetry-koa-instrumentation/test/koa.test.ts
index df90119227e..cefcb26a6cf 100644
--- a/plugins/node/opentelemetry-koa-instrumentation/test/koa.test.ts
+++ b/plugins/node/opentelemetry-koa-instrumentation/test/koa.test.ts
@@ -93,6 +93,12 @@ describe('Koa Instrumentation - Core Tests', () => {
     await next();
   };
 
+  const spanCreateMiddleware: koa.Middleware = async (ctx, next) => {
+    const span = tracer.startSpan('foo');
+    span.end();
+    await next();
+  };
+
   const asyncMiddleware: koa.Middleware = async (ctx, next) => {
     const start = Date.now();
     await next();
@@ -106,11 +112,12 @@ describe('Koa Instrumentation - Core Tests', () => {
       app.use((ctx, next) => tracer.withSpan(rootSpan, next));
       app.use(customMiddleware);
       app.use(simpleResponse);
+      app.use(spanCreateMiddleware);
 
       await tracer.withSpan(rootSpan, async () => {
         await httpRequest.get(`http://localhost:${port}`);
         rootSpan.end();
-        assert.deepStrictEqual(memoryExporter.getFinishedSpans().length, 5);
+        assert.deepStrictEqual(memoryExporter.getFinishedSpans().length, 8);
 
         assert.notStrictEqual(
           memoryExporter
@@ -119,6 +126,18 @@ describe('Koa Instrumentation - Core Tests', () => {
           undefined
         );
 
+        const fooParentSpan = memoryExporter
+          .getFinishedSpans()
+          .find(span => span.name.includes('spanCreateMiddleware'));
+        assert.notStrictEqual(fooParentSpan, undefined);
+
+        const fooSpan = memoryExporter.getFinishedSpans().find(span => 'foo');
+        assert.notStrictEqual(fooSpan, undefined);
+        assert.strictEqual(
+          fooSpan!.parentSpanId,
+          fooParentSpan!.spanContext.spanId
+        );
+
         const simpleResponseSpan = memoryExporter
           .getFinishedSpans()
           .find(span => span.name.includes('simpleResponse'));