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 package org.locationtech.geomesa.utils.iterators
10 
11 import com.typesafe.scalalogging.StrictLogging
12 import org.geotools.api.data.{DataStore, Query, Transaction}
13 import org.geotools.api.feature.simple.SimpleFeature
14 import org.geotools.api.filter.Filter
15 import org.geotools.api.filter.sort.SortOrder
16 import org.geotools.factory.CommonFactoryFinder
17 import org.geotools.feature.simple.SimpleFeatureImpl
18 import org.geotools.temporal.`object`.{DefaultInstant, DefaultPeriod, DefaultPosition}
19 import org.geotools.util.factory.Hints
20 import org.locationtech.geomesa.utils.collection.{CloseableIterator, SelfClosingIterator}
21 
22 import java.util.Date
23 import java.util.concurrent.{Executors, LinkedBlockingQueue}
24 import scala.concurrent.duration.Duration
25 import scala.util.control.NonFatal
26 
27 /**
28  * Query over a time frame and return the features in sorted order, delayed based on the date of each feature
29  * to simulate the original ingestion stream
30  *
31  * @param ds data store
32  * @param typeName simple feature type name
33  * @param interval interval to query
34  * @param dtg date attribute to sort by
35  * @param filter additional filter predicate, if any
36  * @param transforms query transforms, if any
37  * @param window length of a single query window, used to chunk up the total features
38  * @param rate multiplier for the rate of returning features, applied to the original delay between features
39  * @param live project dates to current time
40  * @param readAhead size of the read-ahead queue used for holding features before returning them
41  */
42 class PlaybackIterator(
43     ds: DataStore,
44     typeName: String,
45     interval: (Date, Date),
46     dtg: Option[String] = None,
47     filter: Option[Filter] = None,
48     transforms: Array[String] = null,
49     window: Option[Duration] = None,
50     rate: Float = 10f,
51     live: Boolean = false,
52     readAhead: Int = 10000
53   ) extends CloseableIterator[SimpleFeature] with StrictLogging {
54 
55   import PlaybackIterator.ff
56   import org.locationtech.geomesa.utils.geotools.RichSimpleFeatureType.RichSimpleFeatureType
57 
58   require(interval._2.after(interval._1), s"Interval is not ordered correctly: ${interval._1}/${interval._2}")
59 
60   private val sft = ds.getSchema(typeName)
61   private val dtgName = dtg.orElse(sft.getDtgField).getOrElse {
62     throw new IllegalArgumentException("Schema does not have a default date field")
63   }
64   private val tdefs = transforms match {
65     case null => null
66     case t if t.indexOf(dtgName) == -1 => t :+ dtgName
67     case t => t
68   }
69   private val dtgIndex = tdefs match {
70     case null => sft.indexOf(dtgName)
71     case t => t.indexOf(dtgName)
72   }
73   require(dtgIndex != -1, "Invalid date field")
74   private val dtgProp = ff.property(dtgName)
75   private val sort = ff.sort(dtgName, SortOrder.ASCENDING)
76 
77   private val windowMillis = window.map(_.toMillis).getOrElse(interval._2.getTime - interval._1.getTime + 1)
78 
79   private var start: Long = -1
80   private var eventStart: Long = -1
81 
82   private val features = new LinkedBlockingQueue[SimpleFeature](readAhead)
83   private var staged: SimpleFeature = _
84 
85   private val executor = Executors.newSingleThreadExecutor()
86   executor.submit(new QueryRunnable())
87 
88   override def hasNext: Boolean = {
89     if (staged != null) {
90       true
91     } else {
92       staged = features.take()
93       if (!PlaybackIterator.terminal.eq(staged)) {
94         true
95       } else {
96         features.put(staged) // re-queue the terminal value to keep this method idempotent
97         staged = null
98         false
99       }
100     }
101   }
102 
103   override def next(): SimpleFeature = {
104     val feature = staged
105     staged = null
106     val featureTime = feature.getAttribute(dtgIndex).asInstanceOf[Date].getTime
107     if (start == -1L) {
108       // emit the first feature as soon as it's available, and set the clock to start timing from here
109       logger.debug("Starting replay clock")
110       start = System.currentTimeMillis()
111       eventStart = featureTime
112     }
113     val featureRelativeTime = start + ((featureTime - eventStart) / rate).toLong
114     val sleep = featureRelativeTime - System.currentTimeMillis()
115     if (sleep > 0) {
116       Thread.sleep(sleep)
117     }
118     if (live) {
119       feature.setAttribute(dtgIndex, new Date(featureRelativeTime))
120     }
121     feature
122   }
123 
124   override def close(): Unit = executor.shutdownNow()
125 
126   private class QueryRunnable extends Runnable {
127     override def run(): Unit = {
128       try {
129         var from = interval._1
130         var to = new Date(from.getTime + windowMillis)
131         var loop = true
132 
133         while (loop && !Thread.currentThread().isInterrupted) {
134           if (interval._2.before(to)) {
135             // this query will finish the last window
136             to = interval._2
137             loop = false
138           }
139 
140           logger.debug(s"Running query window $from to $to")
141 
142           val during = {
143             val period = new DefaultPeriod(
144               new DefaultInstant(new DefaultPosition(from)),
145               new DefaultInstant(new DefaultPosition(to))
146             )
147             ff.during(dtgProp, ff.literal(period))
148           }
149           val query = new Query(typeName, filter.map(ff.and(_, during)).getOrElse(during), tdefs: _*)
150           query.setSortBy(sort)
151           // prevent ContentDataStore from sorting on disk
152           query.getHints.put(Hints.MAX_MEMORY_SORT, java.lang.Integer.MAX_VALUE)
153 
154           var count = 0L
155 
156           // populate the queue - this will block if we get too far ahead
157           SelfClosingIterator(ds.getFeatureReader(query, Transaction.AUTO_COMMIT)).foreach { f =>
158             features.put(f)
159             count += 1
160           }
161 
162           logger.debug(s"Returned $count features from query window $from to $to")
163 
164           // increment time window
165           from = to
166           to = new Date(from.getTime + windowMillis)
167         }
168       } catch {
169         case NonFatal(e) => logger.error("Error querying playback:", e)
170       } finally {
171         features.put(PlaybackIterator.terminal)
172       }
173     }
174   }
175 }
176 
177 object PlaybackIterator {
178   private val ff = CommonFactoryFinder.getFilterFactory
179   private val terminal = new SimpleFeatureImpl(null, null, null, false, null)
180 }