1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 
10 package org.locationtech.geomesa.convert.redis
11 
12 import com.github.benmanes.caffeine.cache.{CacheLoader, Caffeine, LoadingCache}
13 import com.typesafe.config.Config
14 import org.locationtech.geomesa.convert.{EnrichmentCache, EnrichmentCacheFactory}
15 import redis.clients.jedis.util.JedisURIHelper
16 import redis.clients.jedis.{Jedis, JedisPool}
17 
18 import java.io.Closeable
19 import java.net.URI
20 import java.util.concurrent.TimeUnit
21 import scala.util.Try
22 
23 trait RedisConnectionBuilder extends Closeable {
24   def getConn: Jedis
25 }
26 
27 class RedisEnrichmentCache(jedisPool: RedisConnectionBuilder,
28                            expiration: Long = -1,
29                            localCache: Boolean) extends EnrichmentCache {
30   type KV = java.util.Map[String, String]
31 
32   private val builder =
33     if (expiration > 0) {
34       Caffeine.newBuilder().expireAfterWrite(expiration, TimeUnit.MILLISECONDS)
35     } else {
36       if (!localCache) {
37         Caffeine.newBuilder().expireAfterWrite(0, TimeUnit.MILLISECONDS).maximumSize(0)
38       } else {
39         Caffeine.newBuilder()
40       }
41     }
42 
43   private val cache: LoadingCache[String, KV] =
44     builder
45       .build(new CacheLoader[String, KV] {
46         override def load(k: String): KV = {
47           val conn = jedisPool.getConn
48           try {
49             conn.hgetAll(k)
50           } finally {
51             // Note: for a JedisPool this only returns it to the pool instead of actionally
52             // closing the connection (so it's safe to call close() on the conn)
53             conn.close()
54           }
55         }
56       })
57 
58   override def get(args: Array[String]): Any = cache.get(args(0)).get(args(1))
59   override def put(args: Array[String], value: Any): Unit = ???
60   override def clear(): Unit = ???
61   override def close(): Unit = jedisPool.close()
62 }
63 
64 class RedisEnrichmentCacheFactory extends EnrichmentCacheFactory {
65   override def canProcess(conf: Config): Boolean = conf.hasPath("type") && conf.getString("type").equals("redis")
66 
67   override def build(conf: Config): EnrichmentCache = {
68     val redisUrl = {
69       val url = conf.getString("redis-url")
70       Some(url).filter(u => Try(new URI(u)).toOption.exists(JedisURIHelper.isValid)).getOrElse {
71         if (url.indexOf(":") == -1) { url } else { s"redis://$url" }
72       }
73     }
74     val timeout = if (conf.hasPath("expiration")) conf.getLong("expiration") else -1
75     val connBuilder: RedisConnectionBuilder = new RedisConnectionBuilder {
76       private val pool = new JedisPool(redisUrl)
77       override def getConn: Jedis = pool.getResource
78       override def close(): Unit = pool.close()
79     }
80 
81     val localCache = Try(conf.getBoolean("local-cache")).getOrElse(true)
82     new RedisEnrichmentCache(connBuilder, timeout, localCache)
83   }
84 }