Adds distributed tracing but with bug in coroutines context: micrometer-metrics/tracing#174

Author: omprakashsridharan

build.gradle.kts

@@ -43,6 +43,15 @@ dependencies {
     // Avro
     implementation("io.confluent:kafka-avro-serializer:7.4.0")
 
+    // Tracing
+    implementation("org.springframework.boot:spring-boot-starter-actuator")
+    implementation("io.micrometer:micrometer-tracing")
+    implementation("io.micrometer:micrometer-tracing-bridge-brave")
+    implementation("io.zipkin.reporter2:zipkin-reporter-brave")
+    implementation("net.logstash.logback:logstash-logback-encoder:7.3")
+    implementation("io.zipkin.brave:brave-context-slf4j")
+    implementation("io.zipkin.brave:brave-instrumentation-kafka-clients")
+
 }
 
 tasks.withType<KotlinCompile> {

docker-compose.yaml

@@ -58,5 +58,10 @@ services:
       SCHEMA_REGISTRY_KAFKASTORE_BOOTSTRAP_SERVERS: 'kafka:29092'
       SCHEMA_REGISTRY_LISTENERS: http://0.0.0.0:8081
 
+  zipkin-server:
+    image: openzipkin/zipkin
+    ports:
+      - "9411:9411"
+
 volumes:
   postgres_dev:

src/main/kotlin/com/omprakash/springbootplayground/SpringBootPlaygroundApplication.kt

@@ -1,11 +1,14 @@
 package com.omprakash.springbootplayground
 
+import com.omprakash.springbootplayground.config.TracingConfig
 import org.springframework.boot.autoconfigure.SpringBootApplication
 import org.springframework.boot.runApplication
+import org.springframework.context.annotation.Import
 import org.springframework.kafka.annotation.EnableKafka
 
 @SpringBootApplication
 @EnableKafka
+@Import(TracingConfig::class)
 class SpringBootPlaygroundApplication
 
 fun main(args: Array<String>) {

src/main/kotlin/com/omprakash/springbootplayground/config/TracingConfig.kt

@@ -0,0 +1,20 @@
+package com.omprakash.springbootplayground.config
+
+import brave.baggage.BaggageField
+import brave.baggage.CorrelationScopeConfig
+import brave.context.slf4j.MDCScopeDecorator
+import brave.propagation.CurrentTraceContext
+import org.springframework.beans.factory.annotation.Qualifier
+import org.springframework.context.annotation.Bean
+
+class TracingConfig {
+    @Bean(name = ["bookId"])
+    fun bookId(): BaggageField = BaggageField.create("bookId")
+
+    @Bean
+    fun mdcScopeDecorator(@Qualifier("bookId") bookId: BaggageField): CurrentTraceContext.ScopeDecorator {
+        return MDCScopeDecorator.newBuilder().clear()
+            .add(CorrelationScopeConfig.SingleCorrelationField.newBuilder(bookId).flushOnUpdate().build())
+            .build()
+    }
+}

src/main/kotlin/com/omprakash/springbootplayground/controllers/BookController.kt

@@ -1,29 +1,51 @@
 package com.omprakash.springbootplayground.controllers
 
+import brave.baggage.BaggageField
 import com.omprakash.springbootplayground.kafka.BookCreatedKafkaProducer
 import com.omprakash.springbootplayground.models.Book
+import com.omprakash.springbootplayground.models.request.CreateBook
 import com.omprakash.springbootplayground.services.BookService
+import io.micrometer.core.instrument.kotlin.asContextElement
+import io.micrometer.observation.ObservationRegistry
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.reactor.mono
+import org.slf4j.LoggerFactory
+import org.springframework.beans.factory.annotation.Qualifier
 import org.springframework.web.bind.annotation.GetMapping
+import org.springframework.web.bind.annotation.PostMapping
+import org.springframework.web.bind.annotation.RequestBody
 import org.springframework.web.bind.annotation.RequestMapping
 import org.springframework.web.bind.annotation.RestController
+import reactor.core.publisher.Mono
 
 @RestController
 @RequestMapping("/books")
-class BookController(private val bookService: BookService, private val bookKafkaProducer: BookCreatedKafkaProducer) {
+class BookController(
+    private val bookService: BookService,
+    private val bookKafkaProducer: BookCreatedKafkaProducer,
+    @Qualifier("bookId") var bookIdBaggageField: BaggageField,
+    private val observationRegistry: ObservationRegistry
+) {
+
+    var logger = LoggerFactory.getLogger(this::class.java)
 
     @GetMapping
     fun findAll(): Flow<Book> {
         return bookService.findAll()
     }
 
-    @GetMapping("/publish")
-    suspend fun publishBook(): String {
-        return try {
-            bookKafkaProducer.publishBook(1L, "Test", "1234").toString()
-        } catch (e: Exception) {
-            println("Exception while publishing books ${e.message}")
-            e.message ?: "Exception"
+    @PostMapping("/publish")
+    fun publishBook(@RequestBody createBook: CreateBook): Mono<String> {
+        return mono(observationRegistry.asContextElement()) {
+            try {
+                bookIdBaggageField.updateValue(createBook.id.toString())
+                bookKafkaProducer.publishBook(createBook.id, createBook.title, createBook.isbn).toString()
+                logger.info("Book published")
+            } catch (e: Exception) {
+                println("Exception while publishing books ${e.message}")
+                e.message ?: "Exception"
+            }
+            "Completed"
         }
     }
 }

src/main/kotlin/com/omprakash/springbootplayground/kafka/BookCreatedKafkaConsumer.kt

@@ -1,14 +1,17 @@
 package com.omprakash.springbootplayground.kafka
 
 import org.apache.avro.generic.GenericRecord
+import org.slf4j.LoggerFactory
 import org.springframework.kafka.annotation.KafkaListener
 import org.springframework.stereotype.Service
 
 @Service
 class BookCreatedKafkaConsumer {
 
+    var logger = LoggerFactory.getLogger(this::class.java)
+
     @KafkaListener(topics = [TOPICS.BOOKS_CREATED], groupId = GROUP_ID.BOOKS_CREATED_CONSUMER, containerFactory = "bookCreatedKafkaListenerContainerFactory")
     fun onBookCreated(bookCreated: GenericRecord) {
-        println("Consumed message $bookCreated")
+        logger.info("Consumed message $bookCreated")
     }
 }

src/main/kotlin/com/omprakash/springbootplayground/kafka/ConsumerConfig.kt

@@ -1,5 +1,7 @@
 package com.omprakash.springbootplayground.kafka
 
+import brave.Tracing
+import brave.kafka.clients.KafkaTracing
 import io.confluent.kafka.serializers.KafkaAvroDeserializer
 import org.apache.avro.generic.GenericRecord
 import org.apache.kafka.clients.consumer.ConsumerConfig.*
@@ -31,9 +33,17 @@ class ConsumerConfig(val kafkaProperties: KafkaProperties) {
     }
 
     @Bean
-    fun bookCreatedKafkaListenerContainerFactory(): ConcurrentKafkaListenerContainerFactory<String, GenericRecord> {
+    fun kafkaTracing(tracing: Tracing?): KafkaTracing? {
+        return KafkaTracing.create(tracing)
+    }
+
+    @Bean
+    fun bookCreatedKafkaListenerContainerFactory(kafkaTracing: KafkaTracing): ConcurrentKafkaListenerContainerFactory<String, GenericRecord> {
         return ConcurrentKafkaListenerContainerFactory<String, GenericRecord>().apply {
-            consumerFactory = consumerFactory()
+            val cf = consumerFactory()
+            cf.addPostProcessor(kafkaTracing::consumer)
+            consumerFactory = cf
+            this.containerProperties.isObservationEnabled = true
         }
     }
 }

src/main/kotlin/com/omprakash/springbootplayground/kafka/Extensions.kt

@@ -12,6 +12,7 @@ import kotlin.coroutines.resumeWithException
 suspend inline fun <reified K : Any, reified V : Any> KafkaTemplate<K, V>.sendAsync(record: ProducerRecord<K, V>): SendResult<K, V> {
     return suspendCancellableCoroutine { cancellableContinuation ->
         val future = this.send(record)
+        this.setObservationEnabled(true)
         future.whenComplete { metadata, exception ->
             if (metadata != null) {
                 cancellableContinuation.resume(metadata)

src/main/kotlin/com/omprakash/springbootplayground/kafka/ProducerConfig.kt

@@ -25,6 +25,10 @@ class ProducerConfig(val kafkaProperties: KafkaProperties) {
     }
 
     @Bean
-    fun booksCreatedKafkaTemplate() = KafkaTemplate(producerFactory())
+    fun booksCreatedKafkaTemplate(): KafkaTemplate<String, GenericRecord> {
+        val kafkaTemplate: KafkaTemplate<String, GenericRecord> = KafkaTemplate(producerFactory())
+        kafkaTemplate.setObservationEnabled(true)
+        return kafkaTemplate
+    }
 
 }

src/main/kotlin/com/omprakash/springbootplayground/models/request/CreateBook.kt

@@ -0,0 +1,3 @@
+package com.omprakash.springbootplayground.models.request
+
+data class CreateBook(val id: Long, val title: String, val isbn: String)

src/main/resources/application.yaml

@@ -1,5 +1,25 @@
 spring:
+  application:
+    name: spring-boot-playground
   sql:
     init:
       schema-locations: classpath:schema.sql
       mode: always
+
+logging:
+  config: classpath:logback-spring.xml
+  pattern:
+    level: "%5p [${spring.application.name:},%X{traceId:-},%X{spanId:-}]"
+
+management:
+  tracing:
+    enabled: true
+    propagation:
+      type: b3
+    sampling:
+      probability: 1.0
+    baggage:
+      correlation:
+        enabled: true
+        fields: bookId
+      remote-fields: bookId

src/main/resources/logback-spring.xml

@@ -0,0 +1,16 @@
+<configuration>
+    <springProperty scope="context" name="service" source="spring.application.name"/>
+    <property scope="context" name="hostname" value="${HOSTNAME}"/>
+    <appender name="console" class="ch.qos.logback.core.ConsoleAppender">
+        <encoder class="net.logstash.logback.encoder.LogstashEncoder">
+            <timeZone>UTC</timeZone>
+            <includeMdcKeyName>traceId</includeMdcKeyName>
+            <includeMdcKeyName>bookId</includeMdcKeyName>
+            <includeContext>false</includeContext>
+            <customFields>{"hostname":"${HOSTNAME}","service":"${service}"}</customFields>
+        </encoder>
+    </appender>
+    <root level="info">
+        <appender-ref ref="console"/>
+    </root>
+</configuration>