feat(opentracing): support server stream calls

edvardchen · Aug 10, 2019 · 2918b93 · 2918b93
1 parent c25b27a
commit 2918b93
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 22 deletions.
diff --git a/__tests__/opentracing.test.ts b/__tests__/opentracing.test.ts
@@ -8,17 +8,17 @@ import { Point } from './fixtures/static_codegen/route_guide_pb';
 describe('opentracing', () => {
   const tracer = initTracer({ serviceName: 'grpc-exp-server' }, {});
 
-  describe('start span', () => {
+  describe('start span witout parent', () => {
     runTest({
       testcase(getServer, client) {
-        it('without parent', done => {
+        it('report span on finished', done => {
           const server = getServer();
           server.use(opentracing());
-
-          server.use(async ({ call: { span } }, next) => {
+          server.use(async ({ call }, next) => {
+            const span = call.span as Span;
             expect(span).toBeInstanceOf(Span);
+            span.finish = done;
             await next();
-            done();
           });
 
           // @ts-ignore
@@ -28,10 +28,10 @@ describe('opentracing', () => {
     });
   });
 
-  describe('start span ', () => {
+  describe('start span with parent', () => {
     runTest({
       testcase(getServer, client) {
-        it('with parent', done => {
+        it('chain on parent', done => {
           const metadata = new Metadata();
           const prx = new Proxy(metadata, {
             set(target, key: string, value) {
@@ -43,9 +43,11 @@ describe('opentracing', () => {
           tracer.inject(clientSpan, FORMAT_HTTP_HEADERS, prx);
 
           const server = getServer();
-          server.use(opentracing({ tracer }));
 
-          server.use(async ({ call: { span } }, next) => {
+          server.use(async ({ call }, next) => {
+            await next();
+
+            const { span } = call;
             expect(span).not.toBeUndefined();
 
             const context = (span as Span).context();
@@ -57,9 +59,9 @@ describe('opentracing', () => {
             // not empty
             expect(parentIdStr).toEqual(spanIdStr);
 
-            await next();
             done();
           });
+          server.use(opentracing({ tracer }));
 
           // @ts-ignore
           client().getFeature(new Point(), metadata, () => {});

diff --git a/src/interceptors/opentracing.ts b/src/interceptors/opentracing.ts
@@ -1,11 +1,11 @@
-import { Tracer, globalTracer, FORMAT_HTTP_HEADERS, Span } from 'opentracing';
+import { Tracer, globalTracer, FORMAT_HTTP_HEADERS, Span, Tags } from 'opentracing';
 import { Interceptor } from '../server';
 
 declare module 'grpc' {
   interface ServerUnaryCall<RequestType> {
     span?: Span;
   }
-  interface ServerReadableStream<RequestType> {
+  interface ServerWriteableStream<RequestType> {
     span?: Span;
   }
   interface ServerReadableStream<RequestType> {
@@ -20,26 +20,32 @@ declare module 'grpc' {
  * extract parent span from metadata and chian on it
  * @param {Tracer} [tracer] If omitted, global tracer would be used
  */
-export default function opentracing(options?: {
-  tracer?: Tracer;
-}): Interceptor {
-  return async ({ call, definition: { path } }, next) => {
+export default function opentracing(options?: { tracer?: Tracer }): Interceptor {
+  return async (ctx, next) => {
+    const {
+      call,
+      definition: { path },
+    } = ctx;
     const tracer = (options && options.tracer) || globalTracer();
 
     const { metadata } = call;
 
     const parent = tracer.extract(FORMAT_HTTP_HEADERS, metadata.getMap());
 
-    const span = tracer.startSpan(
-      path,
-      parent ? { childOf: parent } : undefined
-    );
+    const span = tracer.startSpan(path, parent ? { childOf: parent } : undefined);
 
     // expose span
     call.span = span;
 
-    await next();
+    ctx.onFinished((error: Error | null) => {
+      if (error) {
+        span.addTags({
+          [Tags.ERROR]: error,
+        });
+      }
+      span.finish();
+    });
 
-    span.finish();
+    await next();
   };
 }