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

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"

lapis2/src/main/kotlin/org/genspectrum/lapis/LapisApplication.kt

@@ -3,8 +3,10 @@ package org.genspectrum.lapis
 import org.genspectrum.lapis.config.ReferenceGenomeSchema
 import org.springframework.boot.autoconfigure.SpringBootApplication
 import org.springframework.boot.runApplication
+import org.springframework.cache.annotation.EnableCaching
 
 @SpringBootApplication
+@EnableCaching
 class Lapisv2Application
 
 fun main(args: Array<String>) {

lapis2/src/main/kotlin/org/genspectrum/lapis/silo/SiloClient.kt

@@ -7,6 +7,7 @@ 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
@@ -30,6 +31,7 @@ class SiloClient(
 ) {
     private val httpClient = HttpClient.newHttpClient()
 
+    @Cacheable("siloQueryCache", condition = "#query.action.cacheable")
     fun <ResponseType> sendQuery(query: SiloQuery<ResponseType>): ResponseType {
         val queryJson = objectMapper.writeValueAsString(query)
 

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)

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,expireAfterAccess=1m
+
+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

lapis2/src/test/kotlin/org/genspectrum/lapis/silo/SiloClientTest.kt

@@ -19,8 +19,12 @@ import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.assertThrows
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.Arguments
+import org.junit.jupiter.params.provider.MethodSource
 import org.mockserver.client.MockServerClient
 import org.mockserver.integration.ClientAndServer
+import org.mockserver.matchers.Times
 import org.mockserver.model.HttpRequest.request
 import org.mockserver.model.HttpResponse
 import org.mockserver.model.HttpResponse.response
@@ -39,12 +43,20 @@ class SiloClientTest(
 ) {
     private lateinit var mockServer: ClientAndServer
 
-    private val someQuery = SiloQuery(SiloAction.aggregated(), StringEquals("theColumn", "theValue"))
+    private lateinit var someQuery: SiloQuery<*>
+
+    private var counter = 0
 
     @BeforeEach
-    fun setupMockServer() {
+    fun setup() {
         mockServer = ClientAndServer.startClientAndServer(MOCK_SERVER_PORT)
         requestIdContext.requestId = REQUEST_ID_VALUE
+
+        someQuery = SiloQuery(
+            SiloAction.aggregated(),
+            StringEquals("theColumn", "a value that is difference for each test method: $counter"),
+        )
+        counter++
     }
 
     @AfterEach
@@ -335,15 +347,77 @@ class SiloClientTest(
         assertThat(exception.retryAfter, `is`(nullValue()))
     }
 
-    private fun expectQueryRequestAndRespondWith(httpResponse: HttpResponse?) {
+    @ParameterizedTest
+    @MethodSource("getQueriesThatShouldNotBeCached")
+    fun `GIVEN an action that should not be cached WHEN I send the same request twice THEN server is called twice`(
+        query: SiloQuery<*>,
+    ) {
+        val errorMessage = "make this fail so that we see a difference on the second call"
+
+        expectQueryRequestAndRespondWith(
+            response()
+                .withStatusCode(200)
+                .withBody("""{"queryResult": []}"""),
+            Times.exactly(1),
+        )
+        expectQueryRequestAndRespondWith(
+            response()
+                .withStatusCode(500)
+                .withBody(errorMessage),
+            Times.exactly(1),
+        )
+
+        underTest.sendQuery(query)
+
+        val exception = assertThrows<SiloException> { underTest.sendQuery(query) }
+        assertThat(exception.message, containsString(errorMessage))
+    }
+
+    @ParameterizedTest
+    @MethodSource("getQueriesThatShouldBeCached")
+    fun `GIVEN an action that should be cached WHEN I send the same request twice THEN second time is cached`(
+        query: SiloQuery<*>,
+    ) {
+        expectQueryRequestAndRespondWith(
+            response()
+                .withStatusCode(200)
+                .withBody("""{"queryResult": []}"""),
+            Times.once(),
+        )
+
+        val result1 = underTest.sendQuery(query)
+        val result2 = underTest.sendQuery(query)
+
+        assertThat(result1, `is`(result2))
+    }
+
+    private fun expectQueryRequestAndRespondWith(httpResponse: HttpResponse, times: Times = Times.unlimited()) =
         MockServerClient("localhost", MOCK_SERVER_PORT)
             .`when`(
                 request()
                     .withMethod("POST")
                     .withPath("/query")
                     .withContentType(MediaType.APPLICATION_JSON)
                     .withHeader("X-Request-Id", REQUEST_ID_VALUE),
+                times,
             )
             .respond(httpResponse)
+
+    companion object {
+        @JvmStatic
+        val queriesThatShouldNotBeCached = listOf(
+            Arguments.of(SiloQuery(SiloAction.aggregated(), True)),
+            Arguments.of(SiloQuery(SiloAction.details(), True)),
+            Arguments.of(SiloQuery(SiloAction.genomicSequence(SequenceType.ALIGNED, "sequenceName"), True)),
+            Arguments.of(SiloQuery(SiloAction.genomicSequence(SequenceType.UNALIGNED, "sequenceName"), True)),
+        )
+
+        @JvmStatic
+        val queriesThatShouldBeCached = listOf(
+            Arguments.of(SiloQuery(SiloAction.mutations(), True)),
+            Arguments.of(SiloQuery(SiloAction.aminoAcidMutations(), True)),
+            Arguments.of(SiloQuery(SiloAction.nucleotideInsertions(), True)),
+            Arguments.of(SiloQuery(SiloAction.aminoAcidInsertions(), True)),
+        )
     }
 }