From 9efc9728bb043d33aee9ac170040b750c342e579 Mon Sep 17 00:00:00 2001
From: Ivan Topolnjak <ivantopo@kamon.io>
Date: Tue, 7 Jul 2020 09:10:21 +0200
Subject: [PATCH] ensure the bundle works when the Cassandra Driver isn't in
 the classpath, fixes #805

---
 .../com/datastax/driver/core/ConnectionPoolAdvices.scala | 9 +++++----
 .../com/datastax/driver/core/SessionInterceptor.scala    | 8 ++++----
 .../cassandra/driver/DriverInstrumentation.scala         | 4 ++--
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/ConnectionPoolAdvices.scala b/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/ConnectionPoolAdvices.scala
index 93c80d22c..5568cd021 100644
--- a/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/ConnectionPoolAdvices.scala
+++ b/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/ConnectionPoolAdvices.scala
@@ -58,6 +58,7 @@ object PoolCloseAdvice {
   * Measure time spent waiting for a connection
   * Record number of in-flight queries on just-acquired connection
   */
+class BorrowAdvice
 object BorrowAdvice {
 
   @Advice.OnMethodEnter
@@ -65,11 +66,11 @@ object BorrowAdvice {
     Kamon.clock().nanos()
   }
 
-  @Advice.OnMethodExit(suppress = classOf[Throwable])
+  @Advice.OnMethodExit(suppress = classOf[Throwable], inline = false)
   def onBorrowed(
-      @Advice.Return(readOnly = false) connection: ListenableFuture[Connection],
-      @Advice.Enter start:                               Long,
-      @Advice.This poolMetrics:                          HasPoolMetrics,
+      @Advice.Return connection:  ListenableFuture[Connection],
+      @Advice.Enter start:        Long,
+      @Advice.This poolMetrics:   HasPoolMetrics,
       @Advice.FieldValue("totalInFlight") totalInflight: AtomicInteger
   ): Unit = {
 
diff --git a/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/SessionInterceptor.scala b/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/SessionInterceptor.scala
index 26a11c278..1b35c7255 100644
--- a/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/SessionInterceptor.scala
+++ b/instrumentation/kamon-cassandra/src/main/scala/com/datastax/driver/core/SessionInterceptor.scala
@@ -19,13 +19,13 @@ package com.datastax.driver.core
 import java.util.concurrent.Callable
 
 import kamon.instrumentation.cassandra.driver.InstrumentedSession
-import kanela.agent.libs.net.bytebuddy.asm.Advice
-import kanela.agent.libs.net.bytebuddy.implementation.bind.annotation.SuperCall
+import kanela.agent.libs.net.bytebuddy.implementation.bind.annotation.{RuntimeType, SuperCall}
 
+class SessionInterceptor
 object SessionInterceptor {
 
-  @Advice.OnMethodExit
-  def wrapSession(@SuperCall session: Callable[Session]): Session = {
+  @RuntimeType
+  def newSession(@SuperCall session: Callable[Session]): Session = {
     new InstrumentedSession(session.call())
   }
 }
diff --git a/instrumentation/kamon-cassandra/src/main/scala/kamon/instrumentation/cassandra/driver/DriverInstrumentation.scala b/instrumentation/kamon-cassandra/src/main/scala/kamon/instrumentation/cassandra/driver/DriverInstrumentation.scala
index f0be7dcce..d64a63916 100644
--- a/instrumentation/kamon-cassandra/src/main/scala/kamon/instrumentation/cassandra/driver/DriverInstrumentation.scala
+++ b/instrumentation/kamon-cassandra/src/main/scala/kamon/instrumentation/cassandra/driver/DriverInstrumentation.scala
@@ -32,7 +32,7 @@ class DriverInstrumentation extends InstrumentationBuilder {
     * Wraps the client session with an InstrumentedSession.
     */
   onType("com.datastax.driver.core.Cluster$Manager")
-    .intercept(method("newSession"), SessionInterceptor)
+    .intercept(method("newSession"), classOf[SessionInterceptor])
     .bridge(classOf[ClusterManagerBridge])
 
   /**
@@ -42,7 +42,7 @@ class DriverInstrumentation extends InstrumentationBuilder {
     * Pool metrics are mixed in the pool object itself
     */
   onType("com.datastax.driver.core.HostConnectionPool")
-    .advise(method("borrowConnection"), BorrowAdvice)
+    .advise(method("borrowConnection"), classOf[BorrowAdvice])
     .advise(method("trashConnection"), TrashConnectionAdvice)
     .advise(method("addConnectionIfUnderMaximum"), CreateConnectionAdvice)
     .advise(method("onConnectionDefunct"), ConnectionDefunctAdvice)