Kotlin refactoring O(n) algorithm

How to optimize the code search by keys?

Let’s imagine we have a stable working solution, no matter whether it's on Java or Kotlin. We search for the partner key in a vault. But is it working efficiently? Or could it be better?

private fun getAllPartnerPublicKeys(): Map<String, String> =
    secretsService.getSecret(secretName)

fun getPartnerPublicKey(partnerId: String): String {
    val keys = getAllPartnerPublicKeys()

    // safer than plain "contains"
    val matches = keys.filterKeys {
        key -> key.contains(partnerId, ignoreCase = true) }

    return when {
        matches.isEmpty() ->
            throw IllegalArgumentException("Unknown id: $partnerId")
        matches.size > 1 ->
            throw IllegalStateException("Ambiguous id: $partnerId")
        else ->
            matches.values.first()
    }
}

Will it be reasonable to refactor? Because keys.filterKeys is a linear scan key lookup O(n). We have a string matching per request. Logically thinking, we can use ConcurrentHashMap to store data.

private val partnerKeyCache = ConcurrentHashMap<String, RSAPublicKey>()

fun getPartnerPublicKey(partnerId: String): RSAPublicKey =
    partnerKeyCache.computeIfAbsent(partnerId) {
        val key = jwtKeyService.getPartnerPublicKey(partnerId)
        parsePublicKey(key)
    }

Pros:

• First call for a partnerId is slow (does the real lookup + parsing).
• Subsequent calls for the same partnerId are O(1) and very fast.
• Thread-safe without extra synchronization.
• Simple and idiomatic in Kotlin.

Cons need to watch:

• The cache never evicts or refreshes. If partner JWT keys rotate, the cached RSAPublicKey becomes stale until you restart the app.
• If jwtKeyService.getPartnerPublicKey() itself is expensive or not available, better load from cache with TTL/refresh.

We can optimize and refactor it. Let’s observe alternatives.

1. Use Spring Cache Abstraction

@Service
class PartnerKeyService(private val jwtKeyService: JwtKeyService) {

    @Cacheable(value = ["partnerPublicKeys"], key = "#partnerId")
    fun getPartnerPublicKey(partnerId: String): RSAPublicKey {
        val key = jwtKeyService.getPartnerPublicKey(partnerId)
        return parsePublicKey(key)
    }
}

Create config class “configuration\CacheConfig.kt”

@Configuration
@EnableCaching
class CacheConfig {
    @Bean
    fun cacheManager(): CacheManager =
        ConcurrentMapCacheManager("partnerPublicKeys")
}

Advantages:

• Easy to add TTL, size limits, or switch to Caffeine/Redis later.
• Add @CacheEvict or @CachePut for key rotation.
• Spring handles the caching logic cleanly.

2. Pre-load all keys at startup (if partners don’t change)

This turns it into pure O(1) lookup with no per-request cost.

@Component
class PartnerKeyCache(
    private val jwtKeyService: JwtKeyService
) {
    private val cache: MutableMap<String, RSAPublicKey> = ConcurrentHashMap()

    @PostConstruct
    fun loadAllKeys() { //init
        val all = jwtKeyService.getAllPartnerPublicKeys()
        all.forEach { (id, keyStr) ->
            cache[id.lowercase()] = parsePublicKey(keyStr) // some get
        }
    }
    fun getPartnerPublicKey(partnerId: String): RSAPublicKey =
        cache[partnerId.lowercase()] 
           ?: throw IllegalArgumentException("Unknown id: $partnerId")
}

3. Hybrid: Loading cache with refresh

ConcurrentMapCacheManager does not support TTL/expiration. So we can use Caffeine — it’s fast, lightweight, and supports expiration out of the box.

Add to build.gradle.kts

implementation(“org.springframework.boot:spring-boot-starter-cache”) implementation(“com.github.ben-manes.caffeine:caffeine:3.2.3”)

@Configuration
@EnableCaching
class CacheConfig {
    @Bean
    fun cacheManager(): CacheManager {
        val caffeineCacheManager= CaffeineCacheManager("partnerPublicKeys")

        caffeineCacheManager.setCaffeine(
            Caffeine.newBuilder()
                .maximumSize(1000)                   
                .expireAfterWrite(10, TimeUnit.MINUTES)
        )
        return caffeineCacheManager
    }
}

Service solution using the Caffeine library:

@EnableScheduling
@Configuration
class SchedulingConfig {
}

@Service
class JwtKeyService(
    private val secretsService: SecretsService,
    @Value("\${jwt.secret-name}")
    private val secretName: String
) {
    companion object {
        private const val TTL: Long = 5 * 60 * 1000L
    }

    @Volatile
    private var cachedKeys: Map<String, String> = emptyMap()

    @Volatile
    private var lastFetchTime: Long = 0L

    private val lock = Any()

    private fun getKeys(): Map<String, String> {
        val now = System.currentTimeMillis()
        if (cachedKeys.isNotEmpty() && now - lastFetchTime < TTL) {
            return cachedKeys
        }
        synchronized(lock) {
            val recheckNow = System.currentTimeMillis()
            if (cachedKeys.isEmpty() || recheckNow-lastFetchTime >= TTL) {
                refreshKeysInternal()
            }
        }
        return cachedKeys
    }

    @PostConstruct
    fun init() {
        refreshKeys()
    }

    @Scheduled(fixedDelay = TTL)
    fun refreshKeys() {
        try {
            refreshKeysInternal()
        } catch (e: Exception) {\  
            println("Failed to refresh JWT keys: ${e.message}")
        }
    }

    private fun refreshKeysInternal() {
        val freshKeys = secretsService.getSecret(secretName)
        if (freshKeys.isNotEmpty()) {
            cachedKeys = freshKeys
            lastFetchTime = System.currentTimeMillis() // extend new period
        }
    }

    private fun getAllPartnerPublicKeys(): Map<String, String> =
        getKeys().filterKeys { it.endsWith("Public", ignoreCase = true) }


    @Cacheable(value=["partnerPublicKeys"], key="#partnerId.toLowerCase()")
    fun getPartnerPublicKey(partnerId: String): String {
        if (partnerId.isBlank()) {
            throw IllegalArgumentException("partner_id cannot be blank")
        }
        // Get fresh keys, cache already handles TTL
        val allPartnerKeys = getAllPartnerPublicKeys()

        val matches = allPartnerKeys.filterKeys { key ->
            val lowerKey = key.lowercase()
            lowerKey.contains(partnerId.lowercase())
        }
        return when {
            matches.isEmpty() ->
                throw IllegalArgumentException("Unknown id: $partnerId")
            matches.size > 1 ->
                throw IllegalStateException("Ambiguous id: $partnerId") 
            else ->
                matches.values.first()
        }
    }
}

Conclusion

Why the original code is problematic

• getAllPartnerPublicKeys() is called on every request.
• filterKeys { key.contains(partnerId, ignoreCase = true) } — this is a linear scan (O(n)) over all partners + string contains() (which is relatively expensive).
• No matter how many partners log in, it scales poorly and adds unnecessary latency + CPU usage under load.

After optimization

1. Cache expiration — The Caffeine cache will automatically evict entries after 10 minutes.
2. Performance:

• Hidden refresh happens in the background every N minutes.
• getPartnerPublicKey() is now O(1) most of the time thanks to @Cacheable.
• The inner getAllPartnerPublicKeys() is cheap because the map is already in memory.

Set up part of the unit test :

@SpringBootTest
@AutoConfigureMockMvc
@Transactional
@ActiveProfiles("test")
class ApiIntegrationTest @Autowired constructor(
    private val mockMvc: MockMvc,
    private val objectMapper: ObjectMapper,
    private val jwtKeyService: JwtKeyService,
) {
    @MockitoBean
    lateinit var secretsService: SecretsService
    @Autowired
    private lateinit var cacheManager: CacheManager   // inject

    private val partnerId = "myPartner"

    @BeforeEach
    fun setup() {
        whenever(secretsService.getSecret(any())).thenAnswer {
            mapOf(
                "jwt-my-key.pub" to TestKeys.Keys.Public
            )
        }
        // Clear cache so that getPartnerPublicKey re-executes with mock
        cacheManager.getCache("partnerPublicKeys")?.clear()

        // Optional: force refresh of internal keys too
        jwtKeyService.refreshKeys()
    }
}