From 473c152eae989a35a30570626bdbb92ad17bbfe5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Carlos=20Ch=C3=A1vez?= <jcchavezs@gmail.com>
Date: Mon, 29 Jul 2019 11:20:36 +0200
Subject: [PATCH] chore: adds response reader.

---
 .../zipkin-instrumentation-fetch/index.d.ts     | 17 +++++++++++++++++
 .../src/wrapFetch.js                            |  7 ++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100644 packages/zipkin-instrumentation-fetch/index.d.ts

diff --git a/packages/zipkin-instrumentation-fetch/index.d.ts b/packages/zipkin-instrumentation-fetch/index.d.ts
new file mode 100644
index 00000000..dc7fb76e
--- /dev/null
+++ b/packages/zipkin-instrumentation-fetch/index.d.ts
@@ -0,0 +1,17 @@
+import {Tracer} from 'zipkin';
+
+// ResponseReader allows the user to customize the response based on
+// the response. In order to not to leak context, user should
+// bind this to `instrumentation`. An example use case for such
+// feature is read response header to know if the response comes
+// was cached or not.
+type ResponseReader = (traceId: string, res: any) => void
+
+declare function wrapFetch(got: any, {
+  tracer: Tracer,
+  serviceName: string,
+  remoteServiceName: string,
+  responseReader: ResponseReader
+}): any;
+
+export default wrapFetch;
diff --git a/packages/zipkin-instrumentation-fetch/src/wrapFetch.js b/packages/zipkin-instrumentation-fetch/src/wrapFetch.js
index 1c8c0791..f4c6ea8f 100644
--- a/packages/zipkin-instrumentation-fetch/src/wrapFetch.js
+++ b/packages/zipkin-instrumentation-fetch/src/wrapFetch.js
@@ -2,7 +2,7 @@ const {
   Instrumentation
 } = require('zipkin');
 
-function wrapFetch(fetch, {tracer, serviceName, remoteServiceName}) {
+function wrapFetch(fetch, {tracer, serviceName, remoteServiceName, responseReader}) {
   const instrumentation = new Instrumentation.HttpClient({tracer, serviceName, remoteServiceName});
   return function zipkinfetch(input, opts = {}) {
     return new Promise((resolve, reject) => {
@@ -17,6 +17,11 @@ function wrapFetch(fetch, {tracer, serviceName, remoteServiceName}) {
             instrumentation.recordResponse(traceId, res.status);
           });
           resolve(res);
+          if (responseReader) {
+            tracer.scoped(() => {
+              responseReader.bind(instrumentation)(traceId, res);
+            });
+          }
         }).catch((err) => {
           tracer.scoped(() => {
             instrumentation.recordError(traceId, err);