From 4b90c33722531082f688ca8e601014b21f70198e Mon Sep 17 00:00:00 2001
From: bgpo <bgpomjayma@gmail.com>
Date: Mon, 6 Dec 2021 13:38:00 +0800
Subject: [PATCH] fix: clientMethodTrace missing original properties

---
 .../src/grpc-js/index.ts                      |  6 ++--
 .../src/grpc/index.ts                         |  6 ++--
 .../test/helper.ts                            | 35 +++++++++++++++++++
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc-js/index.ts b/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc-js/index.ts
index ce9777c68b..7e854d7b77 100644
--- a/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc-js/index.ts
+++ b/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc-js/index.ts
@@ -274,7 +274,7 @@ export class GrpcJsInstrumentation extends InstrumentationBase {
     const instrumentation = this;
     return (original: GrpcClientFunc) => {
       instrumentation._diag.debug('patch all client methods');
-      return function clientMethodTrace(this: grpcJs.Client) {
+      function clientMethodTrace(this: grpcJs.Client) {
         const name = `grpc.${original.path.replace('/', '')}`;
         const args = [...arguments];
         const metadata = getMetadata.call(
@@ -289,7 +289,9 @@ export class GrpcJsInstrumentation extends InstrumentationBase {
         return context.with(trace.setSpan(context.active(), span), () =>
           makeGrpcClientRemoteCall(original, args, metadata, this)(span)
         );
-      };
+      }
+      Object.assign(clientMethodTrace, original);
+      return clientMethodTrace;
     };
   }
 
diff --git a/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc/index.ts b/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc/index.ts
index bf488028f0..cbc6bb7560 100644
--- a/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc/index.ts
+++ b/experimental/packages/opentelemetry-instrumentation-grpc/src/grpc/index.ts
@@ -285,7 +285,7 @@ export class GrpcNativeInstrumentation extends InstrumentationBase<
     const instrumentation = this;
     return (original: GrpcClientFunc) => {
       instrumentation._diag.debug('patch all client methods');
-      return function clientMethodTrace(this: grpcTypes.Client) {
+      function clientMethodTrace(this: grpcTypes.Client) {
         const name = `grpc.${(original.path as string | undefined)?.replace(
           '/',
           ''
@@ -304,7 +304,9 @@ export class GrpcNativeInstrumentation extends InstrumentationBase<
             this
           )(span)
         );
-      };
+      }
+      Object.assign(clientMethodTrace, original);
+      return clientMethodTrace;
     };
   }
 }
diff --git a/experimental/packages/opentelemetry-instrumentation-grpc/test/helper.ts b/experimental/packages/opentelemetry-instrumentation-grpc/test/helper.ts
index 5a39f3fd68..5620f233dc 100644
--- a/experimental/packages/opentelemetry-instrumentation-grpc/test/helper.ts
+++ b/experimental/packages/opentelemetry-instrumentation-grpc/test/helper.ts
@@ -646,6 +646,16 @@ export const runTests = (
       });
     };
 
+    const runClientMethodTest = (
+      method: typeof methodList[0]
+    ) => {
+      it(`should assign original properties for grpc remote method ${method.methodName}`, async () => {
+        const patchedClientMethod = (client as any)[method.methodName];
+        const properties = Object.keys(patchedClientMethod);
+        assert.ok(properties.length);
+      });
+    };
+
     describe('enable()', () => {
       const provider = new NodeTracerProvider();
       provider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter));
@@ -801,5 +811,30 @@ export const runTests = (
         });
       });
     });
+
+    describe('Test assigning properties from original client method to patched client method', () => {
+      before(async () => {
+        plugin.disable();
+        plugin.setConfig({});
+        plugin.enable();
+
+        const packageDefinition = await protoLoader.load(PROTO_PATH, options);
+        const proto = grpc.loadPackageDefinition(packageDefinition).pkg_test;
+
+        client = createClient(grpc, proto);
+      });
+
+      after(done => {
+        client.close();
+        server.tryShutdown(() => {
+          plugin.disable();
+          done();
+        });
+      });
+
+      methodList.map(method => {
+        runClientMethodTest(method);
+      });
+    });
   });
 };