1 files changed, 26 insertions, 0 deletions
diff --git a/src/main/scala/xyz/driver/core/trace/ExceptionLoggingSizedBufferingTraceConsumer.scala b/src/main/scala/xyz/driver/core/trace/ExceptionLoggingSizedBufferingTraceConsumer.scala
new file mode 100644
index 0000000..534df8c
--- /dev/null
+++ b/src/main/scala/xyz/driver/core/trace/ExceptionLoggingSizedBufferingTraceConsumer.scala
@@ -0,0 +1,26 @@
+package xyz.driver.core.trace
+
+import com.google.cloud.trace.v1.consumer.{FlushableTraceConsumer}
+import com.google.devtools.cloudtrace.v1.Traces
+import com.typesafe.scalalogging.Logger
+import scala.util.Try
+
+/**
+  * ExceptionLoggingFlushableTraceConsumer simply wraps a flushable trace consumer and catches/logs any exceptions
+  * @param traceConsumer the flusable trace consumer to wrap
+  * @param log where to log any exceptions
+  */
+class ExceptionLoggingFlushableTraceConsumer(traceConsumer: FlushableTraceConsumer, log: Logger)
+    extends FlushableTraceConsumer {
+
+  private val flushableTraceConsumer = traceConsumer
+
+  private def exceptionLogger(exception: Throwable): Unit =
+    log.trace(s"Encountered exception logging to google $exception")
+
+  override def receive(trace: Traces): Unit =
+    Try(flushableTraceConsumer.receive(trace)).fold(exceptionLogger, identity)
+
+  override def flush(): Unit =
+    Try(flushableTraceConsumer.flush()).fold(exceptionLogger, identity)
+}