feat: implement a request cache for mutation and insertion queries #137

Also make sure to invalidate the cache when SILO has a new data version
GenSpectrum · Mar 4, 2024 · 17ff89c · 17ff89c
1 parent 8313996
commit 17ff89c
Show file tree

Hide file tree

Showing 9 changed files with 304 additions and 28 deletions.
diff --git a/lapis2/build.gradle b/lapis2/build.gradle
@@ -29,6 +29,8 @@ repositories {
 
 dependencies {
     implementation 'org.springframework.boot:spring-boot-starter-web'
+    implementation 'org.springframework.boot:spring-boot-starter-actuator'
+    implementation 'org.springframework.boot:spring-boot-starter-cache'
     implementation 'com.fasterxml.jackson.module:jackson-module-kotlin'
     implementation 'org.jetbrains.kotlin:kotlin-reflect'
     implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml'
@@ -39,6 +41,7 @@ dependencies {
     implementation 'org.antlr:antlr4-runtime:4.13.1'
     implementation 'org.apache.commons:commons-csv:1.10.0'
     implementation 'com.github.luben:zstd-jni:1.5.5-11'
+    implementation 'com.github.ben-manes.caffeine:caffeine:3.1.8'
 
     testImplementation('org.springframework.boot:spring-boot-starter-test') {
         exclude group: "org.mockito"

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/LapisSpringConfig.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/LapisSpringConfig.kt
@@ -20,12 +20,16 @@ import org.genspectrum.lapis.openApi.buildOpenApiSchema
 import org.genspectrum.lapis.util.TimeFactory
 import org.genspectrum.lapis.util.YamlObjectMapper
 import org.springframework.beans.factory.annotation.Value
+import org.springframework.cache.annotation.EnableCaching
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
+import org.springframework.scheduling.annotation.EnableScheduling
 import org.springframework.web.filter.CommonsRequestLoggingFilter
 import java.io.File
 
 @Configuration
+@EnableScheduling
+@EnableCaching
 class LapisSpringConfig {
     @Bean
     fun openAPI(

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/auth/DataOpennessAuthorizationFilter.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/auth/DataOpennessAuthorizationFilter.kt
@@ -111,6 +111,7 @@ private class ProtectedDataAuthorizationFilter(
         private val WHITELISTED_PATH_PREFIXES = listOf(
             "/swagger-ui",
             "/api-docs",
+            "/actuator",
             "/sample$DATABASE_CONFIG_ROUTE",
             "/sample$REFERENCE_GENOME_ROUTE",
         )

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/scheduler/DataVersionCacheInvalidator.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/scheduler/DataVersionCacheInvalidator.kt
@@ -0,0 +1,43 @@
+package org.genspectrum.lapis.scheduler
+
+import mu.KotlinLogging
+import org.genspectrum.lapis.silo.CachedSiloClient
+import org.genspectrum.lapis.silo.SILO_QUERY_CACHE_NAME
+import org.springframework.cache.annotation.CacheEvict
+import org.springframework.scheduling.annotation.Scheduled
+import org.springframework.stereotype.Component
+import java.util.concurrent.TimeUnit
+
+private val log = KotlinLogging.logger {}
+
+@Component
+class DataVersionCacheInvalidator(
+    private val cachedSiloClient: CachedSiloClient,
+    private val cacheClearer: CacheClearer,
+) {
+    private var currentlyCachedDataVersion = "uninitialized"
+
+    @Scheduled(fixedRate = 1, timeUnit = TimeUnit.SECONDS)
+    @Synchronized
+    fun invalidateSiloCache() {
+        log.debug { "checking for data version change" }
+
+        val info = cachedSiloClient.callInfo()
+        if (info.dataVersion != currentlyCachedDataVersion) {
+            log.info {
+                "Invalidating cache, old data version: $currentlyCachedDataVersion, " +
+                    "new data version: ${info.dataVersion}"
+            }
+            cacheClearer.clearCache()
+            currentlyCachedDataVersion = info.dataVersion
+        }
+    }
+}
+
+@Component
+class CacheClearer {
+    @CacheEvict(SILO_QUERY_CACHE_NAME, allEntries = true)
+    fun clearCache() {
+        log.info { "Clearing cache $SILO_QUERY_CACHE_NAME" }
+    }
+}
diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/silo/SiloClient.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/silo/SiloClient.kt
@@ -7,10 +7,12 @@ import org.genspectrum.lapis.controller.LapisHeaders.REQUEST_ID
 import org.genspectrum.lapis.logging.RequestIdContext
 import org.genspectrum.lapis.response.InfoData
 import org.springframework.beans.factory.annotation.Value
+import org.springframework.cache.annotation.Cacheable
 import org.springframework.http.HttpHeaders
 import org.springframework.http.HttpStatus
 import org.springframework.http.MediaType
 import org.springframework.stereotype.Component
+import org.springframework.web.context.request.RequestContextHolder
 import java.net.URI
 import java.net.http.HttpClient
 import java.net.http.HttpRequest
@@ -23,14 +25,34 @@ const val SILO_RESPONSE_MAX_LOG_LENGTH = 10_000
 
 @Component
 class SiloClient(
+    private val cachedSiloClient: CachedSiloClient,
+    private val dataVersion: DataVersion,
+) {
+    fun <ResponseType> sendQuery(query: SiloQuery<ResponseType>): ResponseType {
+        val result = cachedSiloClient.sendQuery(query)
+        dataVersion.dataVersion = result.dataVersion
+        return result.queryResult
+    }
+
+    fun callInfo(): InfoData {
+        val info = cachedSiloClient.callInfo()
+        dataVersion.dataVersion = info.dataVersion
+        return info
+    }
+}
+
+const val SILO_QUERY_CACHE_NAME = "siloQueryCache"
+
+@Component
+class CachedSiloClient(
     @Value("\${silo.url}") private val siloUrl: String,
     private val objectMapper: ObjectMapper,
-    private val dataVersion: DataVersion,
     private val requestIdContext: RequestIdContext,
 ) {
     private val httpClient = HttpClient.newHttpClient()
 
-    fun <ResponseType> sendQuery(query: SiloQuery<ResponseType>): ResponseType {
+    @Cacheable(SILO_QUERY_CACHE_NAME, condition = "#query.action.cacheable")
+    fun <ResponseType> sendQuery(query: SiloQuery<ResponseType>): WithDataVersion<ResponseType> {
         val queryJson = objectMapper.writeValueAsString(query)
 
         log.info { "Calling SILO: $queryJson" }
@@ -41,7 +63,10 @@ class SiloClient(
         }
 
         try {
-            return objectMapper.readValue(response.body(), query.action.typeReference).queryResult
+            return WithDataVersion(
+                queryResult = objectMapper.readValue(response.body(), query.action.typeReference).queryResult,
+                dataVersion = getDataVersion(response),
+            )
         } catch (exception: Exception) {
             val message = "Could not parse response from silo: " + exception::class.toString() + " " + exception.message
             throw RuntimeException(message, exception)
@@ -63,7 +88,7 @@ class SiloClient(
         val request = HttpRequest.newBuilder(uri)
             .apply(buildRequest)
             .apply {
-                if (requestIdContext.requestId != null) {
+                if (RequestContextHolder.getRequestAttributes() != null && requestIdContext.requestId != null) {
                     header(REQUEST_ID, requestIdContext.requestId)
                 }
             }
@@ -112,7 +137,6 @@ class SiloClient(
             )
         }
 
-        addDataVersionToRequestScope(response)
         return response
     }
 
@@ -124,10 +148,6 @@ class SiloClient(
             null
         }
 
-    private fun addDataVersionToRequestScope(response: HttpResponse<String>) {
-        dataVersion.dataVersion = getDataVersion(response)
-    }
-
     private fun getDataVersion(response: HttpResponse<String>): String {
         return response.headers().firstValue("data-version").orElse("")
     }
@@ -141,4 +161,9 @@ data class SiloQueryResponse<ResponseType>(
     val queryResult: ResponseType,
 )
 
+data class WithDataVersion<ResponseType>(
+    val dataVersion: String,
+    val queryResult: ResponseType,
+)
+
 data class SiloErrorResponse(val error: String, val message: String)
diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/silo/SiloQuery.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/silo/SiloQuery.kt
@@ -34,6 +34,7 @@ interface CommonActionFields {
 
 sealed class SiloAction<ResponseType>(
     @JsonIgnore val typeReference: TypeReference<SiloQueryResponse<ResponseType>>,
+    @JsonIgnore val cacheable: Boolean,
 ) : CommonActionFields {
     companion object {
         fun aggregated(
@@ -104,7 +105,7 @@ sealed class SiloAction<ResponseType>(
         override val limit: Int? = null,
         override val offset: Int? = null,
         val type: String = "Aggregated",
-    ) : SiloAction<List<AggregationData>>(AggregationDataTypeReference())
+    ) : SiloAction<List<AggregationData>>(AggregationDataTypeReference(), cacheable = false)
 
     @JsonInclude(JsonInclude.Include.NON_EMPTY)
     private data class MutationsAction(
@@ -113,7 +114,7 @@ sealed class SiloAction<ResponseType>(
         override val limit: Int? = null,
         override val offset: Int? = null,
         val type: String = "Mutations",
-    ) : SiloAction<List<MutationData>>(MutationDataTypeReference())
+    ) : SiloAction<List<MutationData>>(MutationDataTypeReference(), cacheable = true)
 
     @JsonInclude(JsonInclude.Include.NON_EMPTY)
     private data class AminoAcidMutationsAction(
@@ -122,7 +123,7 @@ sealed class SiloAction<ResponseType>(
         override val limit: Int? = null,
         override val offset: Int? = null,
         val type: String = "AminoAcidMutations",
-    ) : SiloAction<List<MutationData>>(AminoAcidMutationDataTypeReference())
+    ) : SiloAction<List<MutationData>>(AminoAcidMutationDataTypeReference(), cacheable = true)
 
     @JsonInclude(JsonInclude.Include.NON_EMPTY)
     private data class DetailsAction(
@@ -131,23 +132,23 @@ sealed class SiloAction<ResponseType>(
         override val limit: Int? = null,
         override val offset: Int? = null,
         val type: String = "Details",
-    ) : SiloAction<List<DetailsData>>(DetailsDataTypeReference())
+    ) : SiloAction<List<DetailsData>>(DetailsDataTypeReference(), cacheable = false)
 
     @JsonInclude(JsonInclude.Include.NON_EMPTY)
     private data class NucleotideInsertionsAction(
         override val orderByFields: List<OrderByField> = emptyList(),
         override val limit: Int? = null,
         override val offset: Int? = null,
         val type: String = "Insertions",
-    ) : SiloAction<List<InsertionData>>(InsertionDataTypeReference())
+    ) : SiloAction<List<InsertionData>>(InsertionDataTypeReference(), cacheable = true)
 
     @JsonInclude(JsonInclude.Include.NON_EMPTY)
     private data class AminoAcidInsertionsAction(
         override val orderByFields: List<OrderByField> = emptyList(),
         override val limit: Int? = null,
         override val offset: Int? = null,
         val type: String = "AminoAcidInsertions",
-    ) : SiloAction<List<InsertionData>>(InsertionDataTypeReference())
+    ) : SiloAction<List<InsertionData>>(InsertionDataTypeReference(), cacheable = true)
 
     @JsonInclude(JsonInclude.Include.NON_EMPTY)
     private data class SequenceAction(
@@ -156,7 +157,7 @@ sealed class SiloAction<ResponseType>(
         override val offset: Int? = null,
         val type: SequenceType,
         val sequenceName: String,
-    ) : SiloAction<List<SequenceData>>(SequenceDataTypeReference())
+    ) : SiloAction<List<SequenceData>>(SequenceDataTypeReference(), cacheable = false)
 }
 
 sealed class SiloFilterExpression(val type: String)

diff --git a/lapis2/src/main/resources/application.properties b/lapis2/src/main/resources/application.properties
@@ -2,3 +2,12 @@ springdoc.default-produces-media-type=application/json
 springdoc.api-docs.path=/api-docs
 springdoc.swagger-ui.operationsSorter=alpha
 server.forward-headers-strategy=framework
+
+spring.cache.cache-names=siloQueryCache
+spring.cache.caffeine.spec=maximumSize=50000
+
+management.endpoints.enabled-by-default=false
+management.endpoint.health.enabled=true
+management.endpoint.caches.enabled=true
+management.endpoint.metrics.enabled=true
+management.endpoints.web.exposure.include=health,caches,metrics