fix gtfs tripupdates

1 files changed, 92 insertions, 83 deletions

diff --git a/lib/Server/GTFS_RT.hs b/lib/Server/GTFS_RT.hs
index 1641131..a697295 100644
--- a/lib/Server/GTFS_RT.hs
+++ b/lib/Server/GTFS_RT.hs
@@ -8,80 +8,89 @@
 
 module Server.GTFS_RT (gtfsRealtimeServer) where
 
-import           Prelude                                    hiding (id)
+import           Prelude                                                      hiding
+                                                                              (id)
 
-import           API                                        (GtfsRealtimeAPI)
-import           Control.Monad                              (forM)
-import           Control.Monad.IO.Class                     (MonadIO (..))
-import           Data.ByteString.Lazy                       (fromStrict)
-import           Data.Functor                               ((<&>))
-import qualified Data.Map                                   as M
-import           Data.Maybe                                 (catMaybes)
-import           Data.Pool                                  (Pool)
-import           Data.Sequence                              (Seq)
-import qualified Data.Sequence                              as Seq
-import           Data.Text                                  (Text)
-import qualified Data.Text                                  as T
-import           Data.Text.Encoding                         (encodeUtf8)
-import           Data.Time.Calendar                         (Day, toGregorian)
-import           Data.Time.Clock                            (UTCTime (utctDay),
-                                                             getCurrentTime)
-import           Data.Time.Clock.POSIX                      (getPOSIXTime)
-import           Data.Time.Clock.System                     (SystemTime (systemSeconds),
-                                                             getSystemTime,
-                                                             utcToSystemTime)
-import qualified Data.UUID                                  as UUID
-import           Data.Word                                  (Word64)
-import           Database.Persist                           (Entity (..),
-                                                             PersistQueryRead (selectFirst),
-                                                             selectList, (==.))
-import           Database.Persist.Postgresql                (SqlBackend)
-import           GHC.Float                                  (double2Float,
-                                                             int2Double)
-import           GTFS                                       (GTFS (..),
-                                                             Seconds (..),
-                                                             Stop (..),
-                                                             Trip (..), TripID)
-import           GTFS.Realtime.Alert                        as AL (Alert (..))
-import           GTFS.Realtime.Alert.SeverityLevel          (SeverityLevel (WARNING))
-import           GTFS.Realtime.EntitySelector               as ES (EntitySelector (..))
-import           GTFS.Realtime.FeedEntity                   as FE (FeedEntity (..))
-import           GTFS.Realtime.FeedHeader                   (FeedHeader (FeedHeader))
-import           GTFS.Realtime.FeedHeader.Incrementality    (Incrementality (FULL_DATASET))
-import           GTFS.Realtime.FeedMessage                  as FM (FeedMessage (..))
-import           GTFS.Realtime.Position                     as POS (Position (..))
-import           GTFS.Realtime.TimeRange                    (TimeRange (TimeRange))
-import           GTFS.Realtime.TranslatedString             (TranslatedString (TranslatedString))
-import           GTFS.Realtime.TranslatedString.Translation (Translation (Translation))
-import           GTFS.Realtime.TripDescriptor               as TD (TripDescriptor (..))
-import           GTFS.Realtime.TripUpdate                   as TU (TripUpdate (..))
-import           GTFS.Realtime.TripUpdate.StopTimeEvent     as STE (StopTimeEvent (..))
-import           GTFS.Realtime.TripUpdate.StopTimeUpdate    as STU (StopTimeUpdate (..))
-import           GTFS.Realtime.VehicleDescriptor            as VD (VehicleDescriptor (..))
-import           GTFS.Realtime.VehiclePosition              as VP (VehiclePosition (..))
-import           Persist                                    (Announcement (..),
-                                                             EntityField (..),
-                                                             Key (..),
-                                                             Running (..),
-                                                             Token (..),
-                                                             TrainAnchor (..),
-                                                             TrainPing (..),
-                                                             runSql)
-import           Servant.API                                ((:<|>) (..))
-import           Text.ProtocolBuffers                       (Utf8 (Utf8),
-                                                             defaultValue)
+import           API                                                          (GtfsRealtimeAPI)
+import           Control.Monad                                                (forM)
+import           Control.Monad.IO.Class                                       (MonadIO (..))
+import qualified Data.ByteString.Char8                                        as C8
+import           Data.ByteString.Lazy                                         (fromStrict)
+import           Data.Functor                                                 ((<&>))
+import qualified Data.Map                                                     as M
+import           Data.Maybe                                                   (catMaybes)
+import           Data.Pool                                                    (Pool)
+import           Data.Sequence                                                (Seq)
+import qualified Data.Sequence                                                as Seq
+import           Data.Text                                                    (Text)
+import qualified Data.Text                                                    as T
+import           Data.Text.Encoding                                           (encodeUtf8)
+import           Data.Time.Calendar                                           (Day,
+                                                                               toGregorian)
+import           Data.Time.Clock                                              (UTCTime (utctDay),
+                                                                               getCurrentTime)
+import           Data.Time.Clock.POSIX                                        (getPOSIXTime)
+import           Data.Time.Clock.System                                       (SystemTime (systemSeconds),
+                                                                               getSystemTime,
+                                                                               utcToSystemTime)
+import qualified Data.UUID                                                    as UUID
+import           Data.Word                                                    (Word64)
+import           Database.Persist                                             (Entity (..),
+                                                                               PersistQueryRead (selectFirst),
+                                                                               selectList,
+                                                                               (==.))
+import           Database.Persist.Postgresql                                  (SqlBackend)
+import           GHC.Float                                                    (double2Float,
+                                                                               int2Double)
+import           GTFS                                                         (GTFS (..),
+                                                                               Seconds (..),
+                                                                               Stop (..),
+                                                                               Trip (..),
+                                                                               TripID,
+                                                                               tripsOnDay)
+import           GTFS.Realtime.Alert                                          as AL (Alert (..))
+import           GTFS.Realtime.Alert.SeverityLevel                            (SeverityLevel (WARNING))
+import           GTFS.Realtime.EntitySelector                                 as ES (EntitySelector (..))
+import           GTFS.Realtime.FeedEntity                                     as FE (FeedEntity (..))
+import           GTFS.Realtime.FeedHeader                                     (FeedHeader (FeedHeader))
+import           GTFS.Realtime.FeedHeader.Incrementality                      (Incrementality (FULL_DATASET))
+import           GTFS.Realtime.FeedMessage                                    as FM (FeedMessage (..))
+import           GTFS.Realtime.Position                                       as POS (Position (..))
+import           GTFS.Realtime.TimeRange                                      (TimeRange (TimeRange))
+import           GTFS.Realtime.TranslatedString                               (TranslatedString (TranslatedString))
+import           GTFS.Realtime.TranslatedString.Translation                   (Translation (Translation))
+import           GTFS.Realtime.TripDescriptor                                 as TD (TripDescriptor (..))
+import           GTFS.Realtime.TripUpdate                                     as TU (TripUpdate (..))
+import           GTFS.Realtime.TripUpdate.StopTimeEvent                       as STE (StopTimeEvent (..))
+import           GTFS.Realtime.TripUpdate.StopTimeUpdate                      as STU (StopTimeUpdate (..))
+import qualified GTFS.Realtime.TripUpdate.StopTimeUpdate.ScheduleRelationship as SR
+import           GTFS.Realtime.VehicleDescriptor                              as VD (VehicleDescriptor (..))
+import           GTFS.Realtime.VehiclePosition                                as VP (VehiclePosition (..))
+import           Persist                                                      (Announcement (..),
+                                                                               EntityField (..),
+                                                                               Key (..),
+                                                                               Running (..),
+                                                                               Token (..),
+                                                                               TrainAnchor (..),
+                                                                               TrainPing (..),
+                                                                               runSql)
+import           Servant.API                                                  ((:<|>) (..))
+import           Text.ProtocolBuffers                                         (Utf8 (Utf8),
+                                                                               defaultValue)
 
-import           Control.Monad.Extra                        (mapMaybeM)
-import           Data.List.NonEmpty                         (nonEmpty)
-import           Data.UUID                                  (toASCIIBytes,
-                                                             toLazyASCIIBytes)
-import qualified Data.Vector                                as V
-import           Extrapolation                              (Extrapolator (extrapolateAtPosition, extrapolateAtSeconds),
-                                                             LinearExtrapolator (..))
-import           GTFS                                       (Depth (..))
-import           GTFS.Realtime.TripUpdate                   (TripUpdate (TripUpdate))
-import           Server.Util                                (Service,
-                                                             secondsNow)
+import           Control.Monad.Extra                                          (mapMaybeM)
+import           Data.List.NonEmpty                                           (nonEmpty)
+import           Data.Time.Format.ISO8601                                     (ISO8601 (iso8601Format),
+                                                                               iso8601Show)
+import           Data.UUID                                                    (toASCIIBytes,
+                                                                               toLazyASCIIBytes)
+import qualified Data.Vector                                                  as V
+import           Extrapolation                                                (Extrapolator (extrapolateAtPosition, extrapolateAtSeconds),
+                                                                               LinearExtrapolator (..))
+import           GTFS                                                         (Depth (..))
+import           GTFS.Realtime.TripUpdate                                     (TripUpdate (TripUpdate))
+import           Server.Util                                                  (Service,
+                                                                               secondsNow)
 
 uuidUtf8 :: UUID.UUID -> Utf8
 uuidUtf8 = Utf8 . fromStrict . UUID.toASCIIBytes
@@ -138,23 +147,23 @@ gtfsRealtimeServer gtfs@GTFS{..} dbpool = handleServiceAlerts :<|> handleTripUpd
         handleTripUpdates = runSql dbpool $ do
           today <- liftIO $ getCurrentTime <&> utctDay
           nowSeconds <- secondsNow today
-          running <- selectList [RunningDay ==. today] []
-          anchors <- flip mapMaybeM running $ \r@(Entity key Running{..}) -> do
-            entities <- selectList [TrainAnchorTrip ==. runningTrip, TrainAnchorDay ==. today] []
+          let running = M.toList $ tripsOnDay gtfs today
+          anchors <- flip mapMaybeM running $ \(tripId, trip@Trip{..}) -> do
+            entities <- selectList [TrainAnchorTrip ==. tripId, TrainAnchorDay ==. today] []
             case nonEmpty (fmap entityVal entities) of
               Nothing -> pure Nothing
-              Just anchors -> case M.lookup runningTrip trips of
-                Nothing   -> pure Nothing
-                Just trip -> pure (Just (r, trip, anchors))
+              Just anchors ->
+                  pure $ Just (tripId, trip, anchors)
 
-          dFeedMessage $ Seq.fromList $ mkTripUpdate nowSeconds <$> anchors
-          where mkTripUpdate nowSeconds (Entity (RunningKey (Token uuid)) Running{..}, Trip{..} :: Trip Deep Deep, anchors) =
+          dFeedMessage $ Seq.fromList $ fmap (mkTripUpdate today nowSeconds) anchors
+          where mkTripUpdate today nowSeconds (tripId :: Text, Trip{..} :: Trip Deep Deep, anchors) =
                   let lastCall = extrapolateAtSeconds LinearExtrapolator anchors nowSeconds
                       stations = tripStops
                         <&> (\stop@Stop{..} -> fmap (, stop) $ extrapolateAtPosition LinearExtrapolator anchors (int2Double stopSequence))
-                  in (dFeedEntity (Utf8 $ toLazyASCIIBytes uuid))
+                  -- note: these IDs should be stable across iterations, so just do tripId + runningday. TODO: breaks in case of cross-midnight?
+                  in (dFeedEntity (Utf8 $ fromStrict $ (encodeUtf8 tripId <> "-" <> (C8.pack $ iso8601Show today))))
                   { FE.trip_update = Just $ TripUpdate
-                    { TU.trip = dTripDescriptor runningTrip (Just runningDay)
+                    { TU.trip = dTripDescriptor tripId (Just today) -- TODO will break if cross-midnight train
                     , TU.vehicle = Nothing
                     , TU.stop_time_update = Seq.fromList
                       $ fmap (\(TrainAnchor{..}, Stop{..}) -> StopTimeUpdate
@@ -164,7 +173,7 @@ gtfsRealtimeServer gtfs@GTFS{..} dbpool = handleServiceAlerts :<|> handleTripUpd
                                  defaultValue { STE.delay = Just $ fromIntegral $ unSeconds $ trainAnchorDelay })
                              , STU.departure = Nothing
                              , STU.departure_occupancy_status = Nothing
-                             , STU.schedule_relationship = Nothing
+                             , STU.schedule_relationship = Just SR.SCHEDULED
                              , STU.stop_time_properties = Nothing
                              , STU.ext'field = defaultValue
                              })