refactor: externalize singleton accessors from trace writer (#694)

PR-URL: #694

src/index.ts

@@ -33,7 +33,7 @@ import * as PluginTypes from './plugin-types';
 import {PluginLoaderConfig} from './trace-plugin-loader';
 import * as pluginLoader from './trace-plugin-loader';
 import {TraceAgent} from './trace-api';
-import {traceWriter, TraceWriterSingletonConfig} from './trace-writer';
+import {traceWriter, TraceWriterConfig} from './trace-writer';
 import * as traceUtil from './util';
 
 export {Config, PluginTypes};
@@ -57,10 +57,9 @@ interface TopLevelConfig {
   forceNewAgent_: boolean;
 }
 
-// TraceWriterSingletonConfig = TraceWriterConfig & { forceNewAgent_: boolean }
 // PluginLoaderConfig extends TraceAgentConfig
-type NormalizedConfig =
-    TraceWriterSingletonConfig&PluginLoaderConfig&TopLevelConfig;
+type NormalizedConfig = TraceWriterConfig&PluginLoaderConfig&TopLevelConfig&
+    {forceNewAgent_: boolean};
 
 /**
  * Normalizes the user-provided configuration object by adding default values
@@ -162,7 +161,7 @@ export function start(projectConfig?: Config): PluginTypes.TraceAgent {
   }
   // CLS namespace for context propagation
   cls.createNamespace();
-  traceWriter.create(logger, config, (err) => {
+  traceWriter.create(logger, config).initialize((err) => {
     if (err) {
       stop();
     }

src/trace-writer.ts

@@ -22,6 +22,7 @@ import * as util from 'util';
 import {Constants} from './constants';
 import {SpanKind, Trace} from './trace';
 import {TraceLabels} from './trace-labels';
+import {Singleton} from './util';
 
 const pjson = require('../../package.json');
 
@@ -340,32 +341,4 @@ export class TraceWriter extends common.Service {
   }
 }
 
-export type TraceWriterSingletonConfig = TraceWriterConfig&{
-  forceNewAgent_: boolean;
-};
-
-// Singleton
-let singleton: TraceWriter;
-
-export const traceWriter = {
-  create(
-      logger: common.Logger, config: TraceWriterSingletonConfig,
-      cb?: (err?: Error) => void): TraceWriter {
-    if (!cb) {
-      // tslint:disable-next-line:no-empty
-      cb = () => {};
-    }
-    if (!singleton || config.forceNewAgent_) {
-      singleton = new TraceWriter(logger, config);
-      singleton.initialize(cb);
-    }
-    return singleton;
-  },
-
-  get(): TraceWriter {
-    if (!singleton) {
-      throw new Error('TraceWriter singleton was not initialized.');
-    }
-    return singleton;
-  }
-};
+export const traceWriter = new Singleton(TraceWriter);

src/util.ts

@@ -17,6 +17,11 @@
 import * as fs from 'fs';
 import Module = require('module');
 import * as path from 'path';
+import {Logger} from '@google-cloud/common';  // for types only.
+
+// This symbol must be exported (for now).
+// See: https://github.com/Microsoft/TypeScript/issues/20080
+export const kSingleton = Symbol();
 
 /**
  * Trace API expects stack frames to be a JSON string with the following
@@ -41,6 +46,41 @@ interface PackageJson {
   version: string;
 }
 
+export interface Constructor<T, Config> {
+  new(logger: Logger, config: Config): T;
+  prototype: T;
+  name: string;
+}
+
+/**
+ * A class that provides access to a singleton.
+ * We assume that any such singleton is always constructed with two arguments:
+ * A logger and an arbitrary configuration object.
+ * Instances of this type should only be constructed in module scope.
+ */
+export class Singleton<T, Config> {
+  // Note: private[symbol] is enforced by clang-format.
+  private[kSingleton]: T|null = null;
+
+  constructor(private implementation: Constructor<T, Config>) {}
+
+  create(logger: Logger, config: Config&{forceNewAgent_?: boolean}): T {
+    if (!this[kSingleton] || config.forceNewAgent_) {
+      this[kSingleton] = new this.implementation(logger, config);
+      return this[kSingleton]!;
+    } else {
+      throw new Error(`${this.implementation.name} has already been created.`);
+    }
+  }
+
+  get(): T {
+    if (!this[kSingleton]) {
+      throw new Error(`${this.implementation.name} has not yet been created.`);
+    }
+    return this[kSingleton]!;
+  }
+}
+
 /**
  * Truncates the provided `string` to be at most `length` bytes
  * after utf8 encoding and the appending of '...'.

test/test-config-credentials.ts

@@ -47,12 +47,13 @@ describe('Credentials Configuration', () => {
   before(() => {
     savedProject = process.env.GCLOUD_PROJECT;
     process.env.GCLOUD_PROJECT = '0';
-    trace.enableTraceWriter();
+    trace.setTraceWriterEnabled(true);
     disableNetConnect();
   });
 
   after(() => {
     process.env.GCLOUD_PROJECT = savedProject;
+    trace.setTraceWriterEnabled(false);
     enableNetConnect();
   });
 

test/test-trace-api.ts

@@ -77,7 +77,7 @@ describe('Trace Interface', function() {
       Object.assign(defaultConfig, {
         projectId: '0',
         forceNewAgent_: false
-      }), function(err) {
+      })).initialize(function(err) {
         assert.ok(!err);
         done();
       });

test/test-trace-writer.ts

@@ -107,7 +107,8 @@ describe('TraceWriter', function() {
         serviceContext: {},
         onUncaughtException: 'ignore',
         forceNewAgent_: true
-      } as createTraceWriterOptions, function() {
+      } as createTraceWriterOptions);
+      writer.initialize(function() {
         var spanData = createFakeTrace('fake span');
         writer.defaultLabels = {
           fakeKey: 'value'
@@ -328,7 +329,7 @@ describe('TraceWriter', function() {
           forceNewAgent_: true,
           onUncaughtException: 'ignore',
           serviceContext: {}
-        }, testCase.config), function(err) {
+        }, testCase.config)).initialize(function(err) {
           testCase.assertResults(err, traceWriter.get());
           done();
         });

test/test-util.ts

@@ -18,6 +18,7 @@
 
 import { Constants } from '../src/constants';
 import * as util from '../src/util';
+import { Logger } from '@google-cloud/common';
 
 var assert = require('assert');
 var common = require('./plugins/common'/*.js*/);
@@ -31,6 +32,64 @@ function notNull<T>(arg: T|null): T {
   return arg as T;
 }
 
+// TODO(kjin): Use TypeScript in the rest of this file. This is already done
+// in PR #686.
+describe('Singleton', () => {
+  // A real test logger class is also introduced as part of #686.
+  const logger = {} as any as Logger;
+  class MyClass {
+    constructor(public logger: Logger, public config: {}) {}
+  }
+
+  describe('create', () => {
+    it('creates an instance of the given class', () => {
+      const createResult = new util.Singleton(MyClass).create(logger, {});
+      assert.ok(createResult instanceof MyClass);
+    });
+
+    it('passes arguments to the underlying constructor', () => {
+      const config = {};
+      const createResult = new util.Singleton(MyClass).create(logger, config);
+      assert.strictEqual(createResult.logger, logger);
+      assert.strictEqual(createResult.config, config);
+    });
+
+    it('throws when used more than once, by default', () => {
+      const singleton = new util.Singleton(MyClass);
+      singleton.create(logger, {});
+      assert.throws(() => singleton.create(logger, {}));
+    });
+
+    it('creates a new instance when forceNewAgent_ is true in the config', () => {
+      const singleton = new util.Singleton(MyClass);
+      const createResult1 = singleton.create(logger, {});
+      const createResult2 = singleton.create(logger, { forceNewAgent_: true });
+      assert.notStrictEqual(createResult1, createResult2);
+    });
+  });
+
+  describe('get', () => {
+    it('throws if create was not called first', () => {
+      assert.throws(() => new util.Singleton(MyClass).get());
+    });
+
+    it('returns the same value returned by create function', () => {
+      const singleton = new util.Singleton(MyClass);
+      const createResult = singleton.create(logger, {});
+      const getResult = singleton.get();
+      assert.strictEqual(getResult, createResult);
+    });
+
+    it('does not return a stale value', () => {
+      const singleton = new util.Singleton(MyClass);
+      singleton.create(logger, {});
+      const createResult = singleton.create(logger, { forceNewAgent_: true });
+      const getResult = singleton.get();
+      assert.strictEqual(getResult, createResult);
+    });
+  });
+});
+
 describe('util.truncate', function() {
   it('should truncate objects larger than size', function() {
     assert.strictEqual(util.truncate('abcdefghijklmno', 5), 'ab...');

test/trace.ts

@@ -38,7 +38,7 @@ import * as trace from '../src';
 import {Config, PluginTypes} from '../src';
 import {RootSpanData} from '../src/span-data';
 import {Trace, TraceSpan} from '../src/trace';
-import {LabelObject, TraceWriter, traceWriter, TraceWriterConfig, TraceWriterSingletonConfig} from '../src/trace-writer';
+import {LabelObject, TraceWriter, traceWriter, TraceWriterConfig} from '../src/trace-writer';
 
 export {Config, PluginTypes};
 
@@ -61,9 +61,11 @@ class TestTraceWriter extends TraceWriter {
     });
   }
 }
+setTraceWriterEnabled(false);
 
-let singleton: TraceWriter|null = null;
-disableTraceWriter();
+export function setTraceWriterEnabled(enabled: boolean) {
+  traceWriter['implementation'] = enabled ? TraceWriter : TestTraceWriter;
+}
 
 export type Predicate<T> = (value: T) => boolean;
 
@@ -76,38 +78,6 @@ export function get(): PluginTypes.TraceAgent {
   return trace.get();
 }
 
-export function enableTraceWriter() {
-  if (traceWriter.get.__wrapped) {
-    assert.ok(!singleton);
-    shimmer.massUnwrap([traceWriter], ['create', 'get']);
-  }
-}
-
-export function disableTraceWriter() {
-  if (!traceWriter.get.__wrapped) {
-    assert.throws(traceWriter.get);
-    shimmer.wrap(
-        traceWriter, 'create',
-        () =>
-            (logger: common.Logger, config: TraceWriterSingletonConfig,
-             cb?: (err?: Error) => void): TraceWriter => {
-              if (singleton) {
-                throw new Error('Trace Writer already created.');
-              }
-              singleton = new TestTraceWriter(logger, config);
-              singleton.initialize(cb || (() => {}));
-              return singleton;
-            });
-
-    shimmer.wrap(traceWriter, 'get', () => (): TraceWriter => {
-      if (!singleton) {
-        throw new Error('Trace Writer not initialized.');
-      }
-      return singleton;
-    });
-  }
-}
-
 export function getTraces(predicate?: Predicate<TraceSpan[]>): string[] {
   if (!predicate) {
     predicate = () => true;