/
QueryableFinatraWindowStore.scala
160 lines (139 loc) · 5.86 KB
/
QueryableFinatraWindowStore.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
package com.twitter.finatra.kafkastreams.query
import com.twitter.conversions.DurationOps._
import com.twitter.finatra.kafkastreams.transformer.FinatraTransformer.DateTimeMillis
import com.twitter.finatra.kafkastreams.transformer.FinatraTransformer.WindowStartTime
import com.twitter.finatra.kafkastreams.transformer.aggregation.TimeWindowed
import com.twitter.finatra.kafkastreams.transformer.domain.Time
import com.twitter.finatra.kafkastreams.transformer.stores.internal.FinatraStoresGlobalManager
import com.twitter.finatra.kafkastreams.utils.time._
import com.twitter.finatra.streams.queryable.thrift.domain.ServiceShardId
import com.twitter.finatra.streams.queryable.thrift.partitioning.KafkaPartitioner
import com.twitter.finatra.streams.queryable.thrift.partitioning.StaticServiceShardPartitioner
import com.twitter.util.Duration
import com.twitter.util.logging.Logging
import java.io.File
import org.apache.kafka.common.serialization.Serde
import org.apache.kafka.streams.errors.InvalidStateStoreException
import org.joda.time.DateTimeUtils
import scala.collection.JavaConverters._
object QueryableFinatraWindowStore {
/**
* Calculate the default range of a window query which can be used
* when a query doesn't specify an explicit start and end time
*
* @param windowSize Window size configuration
* @param allowedLateness Window allowed lateness configuration
* @param queryableAfterClose Window queryable after close configuration
*
* @return Duration in millis of how long the default query range should be
*/
def defaultQueryRange(
windowSize: Duration,
allowedLateness: Duration,
queryableAfterClose: Duration
): Duration = {
val windowLifetime = windowSize + allowedLateness + queryableAfterClose
(math
.ceil(
windowLifetime.inMillis.toDouble / windowSize.inMillis.toDouble).toLong * windowSize.inMillis).millis
}
/**
* Calculate the default range of a window query which can be used
* when the query doesn't specify an explicit start and end time
*
* @param windowSize Window's size configuration
* @param defaultQueryRange Duration in millis of the query range (end - start)
* @param queryParamStartTime Query start time
* @param queryParamEndTime Query end time
*
* @return Start and end times to be queried snapped to window boundaries
*/
def queryStartAndEndTime(
windowSize: Duration,
defaultQueryRange: Duration,
queryParamStartTime: Option[DateTimeMillis],
queryParamEndTime: Option[DateTimeMillis]
): (DateTimeMillis, DateTimeMillis) = {
val endTime = queryParamEndTime.getOrElse(DateTimeUtils.currentTimeMillis)
val windowSnappedEndTime =
TimeWindowed.windowEnd(messageTime = Time(endTime), size = windowSize)
val startTimeMillis =
queryParamStartTime.getOrElse(windowSnappedEndTime.millis - defaultQueryRange.inMillis)
val windowSnappedStartTime =
TimeWindowed.windowStart(messageTime = Time(startTimeMillis), size = windowSize)
(windowSnappedStartTime.millis, windowSnappedEndTime.millis)
}
}
/**
* A queryable Finatra window store for use by endpoints exposing queryable state
*/
class QueryableFinatraWindowStore[K, V](
stateDir: File,
storeName: String,
windowSize: Duration,
allowedLateness: Duration,
queryableAfterClose: Duration,
keySerde: Serde[K],
numShards: Int,
numQueryablePartitions: Int,
currentShardId: Int)
extends Logging {
private val keySerializer = keySerde.serializer()
private val currentServiceShardId = ServiceShardId(currentShardId)
private val windowSizeMillis = windowSize.inMillis
private val partitioner = new KafkaPartitioner(
StaticServiceShardPartitioner(numShards = numShards),
numPartitions = numQueryablePartitions)
private val defaultQueryRange = QueryableFinatraWindowStore
.defaultQueryRange(windowSize, allowedLateness, queryableAfterClose)
/**
* Get the windowed values associated with this key and return them combined in a map. If the
* key is non-local to this Kafka Streams instance, return an exception indicating
* which instance is hosting this key
*
* @param key The key to fetch
* @param startTime Optional start time of the range of windows to query
* @param endTime Optional end time of the range of windows to query
*
* @return A map of windowed values for the key being queried
*
* @throws NullPointerException If null is used for key.
* @throws InvalidStateStoreException if the store is not initialized
*/
@throws[InvalidStateStoreException]
def get(
key: K,
startTime: Option[Long] = None,
endTime: Option[Long] = None
): Map[WindowStartTime, V] = {
val primaryKeyBytes = FinatraStoresGlobalManager.primaryKeyBytesIfLocalKey(
partitioner,
currentServiceShardId,
key,
keySerializer)
val (startWindowRange, endWindowRange) = QueryableFinatraWindowStore.queryStartAndEndTime(
windowSize,
defaultQueryRange,
startTime,
endTime)
val windowedMap = new java.util.TreeMap[DateTimeMillis, V]
var currentWindowStart = startWindowRange
while (currentWindowStart <= endWindowRange) {
val windowedKey = TimeWindowed.forSize(Time(currentWindowStart), windowSize, key)
debug(s"Query for $storeName $windowedKey")
val result = FinatraStoresGlobalManager
.getWindowedStore[K, V](stateDir, storeName, numQueryablePartitions, primaryKeyBytes)
.get(windowedKey)
if (result != null) {
debug(s"Found $storeName $windowedKey = $result")
windowedMap.put(currentWindowStart, result)
} else {
debug(s"NotFound $storeName $windowedKey = $result")
}
currentWindowStart = currentWindowStart + windowSizeMillis
}
info(
s"Query $key ${startWindowRange.iso8601Millis} to ${endWindowRange.iso8601Millis} = $windowedMap")
windowedMap.asScala.toMap
}
}