test: test hapi tails

test/test-trace-hapi-tails.ts

@@ -0,0 +1,115 @@
+/**
+ * Copyright 2019 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as assert from 'assert';
+import axiosModule from 'axios';
+import * as semver from 'semver';
+
+import * as testTraceModule from './trace';
+import {assertSpanDuration, wait} from './utils';
+import {Hapi17} from './web-frameworks/hapi17';
+import {Hapi12, Hapi15, Hapi16, Hapi8} from './web-frameworks/hapi8_16';
+
+// The list of web frameworks to test.
+const FRAMEWORKS = [Hapi12, Hapi15, Hapi16, Hapi8, Hapi17];
+
+describe('Web framework tracing', () => {
+  let axios: typeof axiosModule;
+  before(() => {
+    testTraceModule.setCLSForTest();
+    testTraceModule.setPluginLoaderForTest();
+    testTraceModule.start();
+    axios = require('axios');
+  });
+
+  after(() => {
+    testTraceModule.setCLSForTest(testTraceModule.TestCLS);
+    testTraceModule.setPluginLoaderForTest(testTraceModule.TestPluginLoader);
+  });
+
+  FRAMEWORKS.forEach((webFrameworkConstructor) => {
+    const commonName = webFrameworkConstructor.commonName;
+    const versionRange = webFrameworkConstructor.versionRange;
+
+    // Skip this set for incompatible versions of Node
+    const skip = !semver.satisfies(process.version, versionRange);
+
+    (skip ? describe.skip : describe)(`Tracing ${commonName}`, () => {
+      // How this test works:
+      // On some WebFramework implementations (currently just Hapi), we can
+      // add "tail work" that is allowed to finish after the request ends.
+      // Hapi 8-16 provides built-in support to keep track of these types of
+      // calls, while Hapi 17 provides a more general mechanism for managing
+      // request lifecycle (both before and after the response has been sent.)
+      // Although the Trace Agent itself doesn't have any special behavior
+      // to account for the Hapi APIs, we would expect that a user using the
+      // tails API to observe a child span with the correct tail duration.
+      it('Traces tail calls correctly', async () => {
+        const framework = new webFrameworkConstructor();
+        try {
+          // "tail work" which will complete independent of the server response.
+          framework.addHandler({
+            path: '/tail',
+            hasResponse: false,
+            blocking: false,
+            fn: async () => {
+              const child =
+                  testTraceModule.get().createChildSpan({name: 'my-tail-work'});
+              await wait(100);
+              child.endSpan();
+            }
+          });
+          framework.addHandler({
+            path: '/tail',
+            hasResponse: true,
+            fn: async () => (
+                {statusCode: 200, message: 'there is still work to be done'})
+          });
+          // A Promise that resolves when the tail call is finished.
+          const tailCallMade =
+              new Promise((resolve) => framework.once('tail', resolve));
+          // Start listening.
+          const port = await framework.listen(0);
+          // Hit the server.
+          await testTraceModule.get().runInRootSpan(
+              {name: 'outer'}, async (span) => {
+                await axios.get(`http://localhost:${port}/tail`);
+                span.endSpan();
+              });
+          // A child span should have been observed by the Trace Writer.
+          const childSpanBeforeEnd =
+              testTraceModule.getOneSpan(span => span.name === 'my-tail-work');
+          assert.ok(!childSpanBeforeEnd.endTime);
+          // Simulate a "flush". The Trace Writer itself will not publish a
+          // span that doesn't have an end time.
+          testTraceModule.clearTraceData();
+          await tailCallMade;
+          // The same child span should have been observed again by the
+          // Trace Writer.
+          const childSpanAfterEnd =
+              testTraceModule.getOneSpan(span => span.name === 'my-tail-work');
+          assert.strictEqual(
+              childSpanAfterEnd.spanId, childSpanBeforeEnd.spanId);
+          // The child span only needs to be at least 100ms.
+          assertSpanDuration(childSpanAfterEnd, [100, Infinity]);
+        } finally {
+          framework.shutdown();
+          testTraceModule.clearTraceData();
+        }
+      });
+    });
+  });
+});

test/test-trace-web-frameworks.ts

@@ -92,6 +92,7 @@ describe('Web framework tracing', () => {
         webFramework.addHandler({
           path: '/two-handlers',
           hasResponse: false,
+          blocking: true,
           fn: async () => {
             await wait(DEFAULT_SPAN_DURATION / 2);
           }

test/web-frameworks/base.ts

@@ -29,12 +29,19 @@ export interface WebFrameworkResponse {
  * The underlying type of objects passed to WebFramework#addHandler.
  */
 export type WebFrameworkAddHandlerOptions = {
+  // The path which will invoke the handler.
   path: string;
 }&({
+  // This handler doesn't provide the response.
   hasResponse: false;
+  // Whether or not this handler should block the next handler.
+  blocking: boolean;
+  // The handler function.
   fn: (incomingHeaders: IncomingHttpHeaders) => Promise<void>;
 }|{
+  // This handler provides a response.
   hasResponse: true;
+  // The handler function.
   fn: (incomingHeaders: IncomingHttpHeaders) => Promise<WebFrameworkResponse>;
 });
 

test/web-frameworks/connect.ts

@@ -34,6 +34,11 @@ export class Connect3 implements WebFramework {
   }
 
   addHandler(options: WebFrameworkAddHandlerOptions): void {
+    if (!options.hasResponse && !options.blocking) {
+      throw new Error(`${
+          this.constructor
+              .name} wrapper for testing doesn't support non-blocking handlers.`);
+    }
     this.app.use(
         options.path,
         async (

test/web-frameworks/express.ts

@@ -34,6 +34,11 @@ export class Express4 implements WebFramework {
   }
 
   addHandler(options: WebFrameworkAddHandlerOptions): void {
+    if (!options.hasResponse && !options.blocking) {
+      throw new Error(`${
+          this.constructor
+              .name} wrapper for testing doesn't support non-blocking handlers.`);
+    }
     this.app.get(options.path, async (req, res, next) => {
       let response: WebFrameworkResponse|void;
       try {

test/web-frameworks/hapi17.ts

@@ -14,11 +14,19 @@
  * limitations under the License.
  */
 
+import {EventEmitter} from 'events';
+
 import {hapi_17} from '../../src/plugins/types';
 
 import {WebFramework, WebFrameworkAddHandlerOptions, WebFrameworkHandlerFunction, WebFrameworkResponse} from './base';
 
-export class Hapi17 implements WebFramework {
+const TAIL_WORK = Symbol('tail work for hapi');
+
+type AppState = {
+  [TAIL_WORK]?: Array<Promise<void>>;
+};
+
+export class Hapi17 extends EventEmitter implements WebFramework {
   static commonName = `hapi@17`;
   static expectedTopStackFrame = '_executeWrap';
   static versionRange = '>=7.5';
@@ -28,21 +36,27 @@ export class Hapi17 implements WebFramework {
   // So instead of registering a new Hapi plugin per path,
   // register only the first time -- passing a function that will iterate
   // through a list of routes keyed under the path.
-  private routes = new Map<string, WebFrameworkHandlerFunction[]>();
+  private routes = new Map<string, WebFrameworkAddHandlerOptions[]>();
   private registering = Promise.resolve();
 
   constructor() {
+    super();
     const hapi = require('../plugins/fixtures/hapi17') as typeof hapi_17;
     this.server = new hapi.Server();
+    this.server.events.on('response', (request: hapi_17.Request) => {
+      Promise.all((request.app as AppState)[TAIL_WORK] || [])
+          .then(
+              () => this.emit('tail'), (err: Error) => this.emit('tail', err));
+    });
   }
 
   addHandler(options: WebFrameworkAddHandlerOptions): void {
     let shouldRegister = false;
     if (!this.routes.has(options.path)) {
-      this.routes.set(options.path, [options.fn]);
+      this.routes.set(options.path, [options]);
       shouldRegister = true;
     } else {
-      this.routes.get(options.path)!.push(options.fn);
+      this.routes.get(options.path)!.push(options);
     }
 
     // Only register a new plugin for the first occurrence of this path.
@@ -56,10 +70,21 @@ export class Hapi17 implements WebFramework {
               path: options.path,
               handler: async (request, h) => {
                 let result;
-                for (const handler of this.routes.get(options.path)!) {
-                  result = await handler(request.raw.req.headers);
-                  if (result) {
-                    return result;
+                for (const localOptions of this.routes.get(options.path)!) {
+                  if (localOptions.hasResponse || localOptions.blocking) {
+                    result = await localOptions.fn(request.raw.req.headers);
+                    if (result) {
+                      return result;
+                    }
+                  } else {
+                    // Use Hapi 17's application state to keep track of
+                    // tail work.
+                    const appState: AppState = request.app;
+                    if (!appState[TAIL_WORK]) {
+                      appState[TAIL_WORK] = [];
+                    }
+                    appState[TAIL_WORK]!.push(
+                        localOptions.fn(request.raw.req.headers));
                   }
                 }
                 return h.continue;

test/web-frameworks/hapi8_16.ts

@@ -14,13 +14,13 @@
  * limitations under the License.
  */
 
-import * as http from 'http';
+import {EventEmitter} from 'events';
 
 import {hapi_16} from '../../src/plugins/types';
 
 import {WebFramework, WebFrameworkAddHandlerOptions, WebFrameworkResponse} from './base';
 
-export class Hapi implements WebFramework {
+export class Hapi extends EventEmitter implements WebFramework {
   server: hapi_16.Server;
   // In Hapi, handlers are added after a connection is specified.
   // Since a port number is required to initialize a connection,
@@ -29,8 +29,10 @@ export class Hapi implements WebFramework {
   queuedHandlers: Array<() => void> = [];
 
   constructor(path: string) {
+    super();
     const hapi = require(path) as typeof hapi_16;
     this.server = new hapi.Server();
+    this.server.on('tail', () => this.emit('tail'));
   }
 
   addHandler(options: WebFrameworkAddHandlerOptions): void {
@@ -51,15 +53,24 @@ export class Hapi implements WebFramework {
           }
         });
       } else {
-        this.server.ext('onPreHandler', async (request, reply) => {
-          try {
-            await options.fn(request.raw.req.headers);
-          } catch (e) {
-            reply(e);
-            return;
-          }
-          reply.continue();
-        });
+        if (options.blocking) {
+          this.server.ext('onPreHandler', async (request, reply) => {
+            try {
+              await options.fn(request.raw.req.headers);
+            } catch (e) {
+              reply(e);
+              return;
+            }
+            reply.continue();
+          });
+        } else {
+          // Use Hapi's request.tail to keep track of tail work.
+          this.server.ext('onPreHandler', (request, reply) => {
+            const tail = request.tail();
+            options.fn(request.raw.req.headers).then(tail, tail);
+            reply.continue();
+          });
+        }
       }
     });
   }

test/web-frameworks/koa1.ts

@@ -36,6 +36,11 @@ export class Koa1 implements WebFramework {
   }
 
   addHandler(options: WebFrameworkAddHandlerOptions): void {
+    if (!options.hasResponse && !options.blocking) {
+      throw new Error(`${
+          this.constructor
+              .name} wrapper for testing doesn't support non-blocking handlers.`);
+    }
     this.app.use(function*(next) {
       if (this.request.path === options.path) {
         // Context doesn't automatically get propagated to yielded functions.

test/web-frameworks/koa2.ts

@@ -35,6 +35,11 @@ export class Koa2 implements WebFramework {
   }
 
   addHandler(options: WebFrameworkAddHandlerOptions): void {
+    if (!options.hasResponse && !options.blocking) {
+      throw new Error(`${
+          this.constructor
+              .name} wrapper for testing doesn't support non-blocking handlers.`);
+    }
     this.app.use(async (ctx, next) => {
       if (ctx.request.path === options.path) {
         const response = await options.fn(ctx.req.headers);

test/web-frameworks/restify.ts

@@ -29,6 +29,11 @@ export class Restify implements WebFramework {
   }
 
   addHandler(options: WebFrameworkAddHandlerOptions): void {
+    if (!options.hasResponse && !options.blocking) {
+      throw new Error(`${
+          this.constructor
+              .name} wrapper for testing doesn't support non-blocking handlers.`);
+    }
     if (options.hasResponse) {
       this.server.get(options.path, async (req, res, next) => {
         let response: WebFrameworkResponse;

tsconfig.json

@@ -38,6 +38,7 @@
     "test/test-trace-api.ts",
     "test/test-trace-api-none-cls.ts",
     "test/test-trace-cluster.ts",
+    "test/test-trace-hapi-tails.ts",
     "test/test-trace-policy.ts",
     "test/test-trace-uncaught-exception.ts",
     "test/test-trace-web-frameworks.ts",