feat(backend): add x-request-id header to trace requests across artif…

…acts in logs resolves #1072
loculus-project · Oct 15, 2024 · 08cb3c2 · 08cb3c2
1 parent f4d6e13
commit 08cb3c2
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 6 deletions.
diff --git a/backend/src/main/kotlin/org/loculus/backend/config/BackendSpringConfig.kt b/backend/src/main/kotlin/org/loculus/backend/config/BackendSpringConfig.kt
@@ -2,10 +2,16 @@ package org.loculus.backend.config
 
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.module.kotlin.readValue
+import io.swagger.v3.oas.models.headers.Header
+import io.swagger.v3.oas.models.media.StringSchema
+import io.swagger.v3.oas.models.parameters.HeaderParameter
 import org.flywaydb.core.Flyway
 import org.jetbrains.exposed.spring.autoconfigure.ExposedAutoConfiguration
 import org.jetbrains.exposed.sql.DatabaseConfig
 import org.jetbrains.exposed.sql.Slf4jSqlDebugLogger
+import org.loculus.backend.controller.LoculusCustomHeaders
+import org.loculus.backend.log.REQUEST_ID_HEADER_DESCRIPTION
+import org.springdoc.core.customizers.OperationCustomizer
 import org.springframework.beans.factory.annotation.Value
 import org.springframework.boot.autoconfigure.ImportAutoConfiguration
 import org.springframework.boot.autoconfigure.jdbc.DataSourceTransactionManagerAutoConfiguration
@@ -76,6 +82,36 @@ class BackendSpringConfig {
 
     @Bean
     fun openApi(backendConfig: BackendConfig) = buildOpenApiSchema(backendConfig)
+
+    @Bean
+    fun headerCustomizer() = OperationCustomizer { operation, _ ->
+        val foundRequestIdHeaderParameter = operation.parameters?.any { it.name == LoculusCustomHeaders.REQUEST_ID }
+        if (foundRequestIdHeaderParameter == false || foundRequestIdHeaderParameter == null) {
+            operation.addParametersItem(
+                HeaderParameter().apply {
+                    name = LoculusCustomHeaders.REQUEST_ID
+                    required = false
+                    example = "1747481c-816c-4b60-af20-a61717a35067"
+                    description = REQUEST_ID_HEADER_DESCRIPTION
+                    schema = StringSchema()
+                },
+            )
+        }
+        for ((_, response) in operation.responses) {
+            if (response.headers == null || !response.headers.containsKey(LoculusCustomHeaders.REQUEST_ID)) {
+                response.addHeaderObject(
+                    LoculusCustomHeaders.REQUEST_ID,
+                    Header().apply {
+                        description = REQUEST_ID_HEADER_DESCRIPTION
+                        required = false
+                        example = "1747481c-816c-4b60-af20-a61717a35067"
+                        schema = StringSchema()
+                    },
+                )
+            }
+        }
+        operation
+    }
 }
 
 fun readBackendConfig(objectMapper: ObjectMapper, configPath: String): BackendConfig {

diff --git a/backend/src/main/kotlin/org/loculus/backend/controller/LoculusCustomHeaders.kt b/backend/src/main/kotlin/org/loculus/backend/controller/LoculusCustomHeaders.kt
@@ -0,0 +1,6 @@
+package org.loculus.backend.controller
+
+object LoculusCustomHeaders {
+    const val REQUEST_ID = "x-request-id"
+    const val X_TOTAL_RECORDS = "x-total-records"
+}
diff --git a/backend/src/main/kotlin/org/loculus/backend/controller/SubmissionController.kt b/backend/src/main/kotlin/org/loculus/backend/controller/SubmissionController.kt
@@ -33,13 +33,17 @@ import org.loculus.backend.api.UnprocessedData
 import org.loculus.backend.api.WarningsFilter
 import org.loculus.backend.auth.AuthenticatedUser
 import org.loculus.backend.auth.HiddenParam
+import org.loculus.backend.controller.LoculusCustomHeaders.X_TOTAL_RECORDS
+import org.loculus.backend.log.REQUEST_ID_MDC_KEY
+import org.loculus.backend.log.RequestIdContext
 import org.loculus.backend.model.RELEASED_DATA_RELATED_TABLES
 import org.loculus.backend.model.ReleasedDataModel
 import org.loculus.backend.model.SubmissionParams
 import org.loculus.backend.model.SubmitModel
 import org.loculus.backend.service.submission.SubmissionDatabaseService
 import org.loculus.backend.utils.Accession
 import org.loculus.backend.utils.IteratorStreamer
+import org.slf4j.MDC
 import org.springframework.http.HttpHeaders
 import org.springframework.http.HttpStatus
 import org.springframework.http.MediaType
@@ -71,6 +75,7 @@ open class SubmissionController(
     private val releasedDataModel: ReleasedDataModel,
     private val submissionDatabaseService: SubmissionDatabaseService,
     private val iteratorStreamer: IteratorStreamer,
+    private val requestIdContext: RequestIdContext,
 ) {
 
     @Operation(description = SUBMIT_DESCRIPTION)
@@ -259,7 +264,7 @@ open class SubmissionController(
         ],
         headers = [
             Header(
-                name = "x-total-records",
+                name = X_TOTAL_RECORDS,
                 description = "The total number of records sent in responseBody",
                 schema = Schema(type = "integer"),
             ),
@@ -297,7 +302,7 @@ open class SubmissionController(
         compression?.let { headers.add(HttpHeaders.CONTENT_ENCODING, it.compressionName) }
 
         val totalRecords = submissionDatabaseService.countReleasedSubmissions(organism)
-        headers.add("x-total-records", totalRecords.toString())
+        headers.add(X_TOTAL_RECORDS, totalRecords.toString())
         // TODO(https://github.com/loculus-project/loculus/issues/2778)
         // There's a possibility that the totalRecords change between the count and the actual query
         // this is not too bad, if the client ends up with a few more records than expected
@@ -369,7 +374,7 @@ open class SubmissionController(
         description = GET_ORIGINAL_METADATA_RESPONSE_DESCRIPTION,
         headers = [
             Header(
-                name = "x-total-records",
+                name = X_TOTAL_RECORDS,
                 description = "The total number of records sent in responseBody",
                 schema = Schema(type = "integer"),
             ),
@@ -411,7 +416,7 @@ open class SubmissionController(
             groupIdsFilter?.takeIf { it.isNotEmpty() },
             statusesFilter?.takeIf { it.isNotEmpty() },
         )
-        headers.add("x-total-records", totalRecords.toString())
+        headers.add(X_TOTAL_RECORDS, totalRecords.toString())
         // TODO(https://github.com/loculus-project/loculus/issues/2778)
         // There's a possibility that the totalRecords change between the count and the actual query
         // this is not too bad, if the client ends up with a few more records than expected
@@ -476,6 +481,8 @@ open class SubmissionController(
         compressionFormat: CompressionFormat? = null,
         sequenceProvider: () -> Sequence<T>,
     ) = StreamingResponseBody { responseBodyStream ->
+        MDC.put(REQUEST_ID_MDC_KEY, requestIdContext.requestId)
+
         val outputStream = when (compressionFormat) {
             CompressionFormat.ZSTD -> ZstdCompressorOutputStream(responseBodyStream)
             null -> responseBodyStream
@@ -493,5 +500,6 @@ open class SubmissionController(
                 }
             }
         }
+        MDC.remove(REQUEST_ID_MDC_KEY)
     }
 }
diff --git a/backend/src/main/kotlin/org/loculus/backend/log/RequestId.kt b/backend/src/main/kotlin/org/loculus/backend/log/RequestId.kt
@@ -0,0 +1,47 @@
+package org.loculus.backend.log
+
+import jakarta.servlet.FilterChain
+import jakarta.servlet.http.HttpServletRequest
+import jakarta.servlet.http.HttpServletResponse
+import org.loculus.backend.controller.LoculusCustomHeaders.REQUEST_ID
+import org.slf4j.MDC
+import org.springframework.core.annotation.Order
+import org.springframework.stereotype.Component
+import org.springframework.web.context.annotation.RequestScope
+import org.springframework.web.filter.OncePerRequestFilter
+import java.util.UUID
+
+@Component
+@RequestScope
+class RequestIdContext {
+    var requestId: String? = null
+}
+
+const val REQUEST_ID_MDC_KEY = "RequestId"
+private const val HIGH_PRECEDENCE_BUT_LOW_ENOUGH_TO_HAVE_REQUEST_SCOPE_AVAILABLE = -100
+const val REQUEST_ID_HEADER_DESCRIPTION = """
+A UUID that uniquely identifies the request for tracing purposes.
+If none is provided in the request, the backend will generate one.
+"""
+
+@Component
+@Order(HIGH_PRECEDENCE_BUT_LOW_ENOUGH_TO_HAVE_REQUEST_SCOPE_AVAILABLE)
+class RequestIdFilter(private val requestIdContext: RequestIdContext) : OncePerRequestFilter() {
+    override fun doFilterInternal(
+        request: HttpServletRequest,
+        response: HttpServletResponse,
+        filterChain: FilterChain,
+    ) {
+        val requestId = request.getHeader(REQUEST_ID) ?: UUID.randomUUID().toString()
+
+        MDC.put(REQUEST_ID_MDC_KEY, requestId)
+        requestIdContext.requestId = requestId
+        response.addHeader(REQUEST_ID, requestId)
+
+        try {
+            filterChain.doFilter(request, response)
+        } finally {
+            MDC.remove(REQUEST_ID_MDC_KEY)
+        }
+    }
+}
diff --git a/backend/src/main/resources/logback.xml b/backend/src/main/resources/logback.xml
@@ -1,6 +1,6 @@
 <configuration>
 
-    <property name="PATTERN" value="%date %level [%thread] [%X{organism}] - %class: %message%n"/>
+    <property name="PATTERN" value="%date %level [%thread] [%X{RequestId}] [%X{organism}] - %class: %message%n"/>
 
     <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
         <layout class="ch.qos.logback.classic.PatternLayout">

diff --git a/backend/src/test/resources/logback-test.xml b/backend/src/test/resources/logback-test.xml
@@ -1,7 +1,7 @@
 <configuration>
     <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
         <layout class="ch.qos.logback.classic.PatternLayout">
-            <Pattern>%date %level [%thread] [%X{organism}] - %class: %message%n</Pattern>
+            <Pattern>%date %level [%thread] [%X{RequestId}] [%X{organism}] - %class: %message%n</Pattern>
         </layout>
     </appender>