@@ -0,0 +1,408 @@
1+{-# LANGUAGE RecursiveDo #-}
2+{-# LANGUAGE GADTs #-}
3+{-# LANGUAGE FlexibleContexts #-}
4+{-# LANGUAGE UndecidableInstances #-}
5+{-# LANGUAGE DeriveGeneric #-}
6+{-# LANGUAGE ScopedTypeVariables #-}
7+{-# LANGUAGE Rank2Types #-}
8+module Obelisk.Api.Pipeline where
9+
10+import Prelude hiding (id, (.))
11+
12+import Obelisk.Request (Request, SomeRequest (..))
13+
14+import Control.Concurrent
15+import Control.Monad
16+import Control.Monad.State.Strict
17+import Data.Typeable
18+import GHC.Generics
19+import Control.Lens (imapM_)
20+import Data.Aeson
21+import Reflex.Query.Base
22+import Reflex.Query.Class
23+import Reflex.Patch
24+import Reflex.Class
25+import Control.Category
26+import Snap.Core (Snap, MonadSnap)
27+import Data.Text (Text)
28+import Data.Map.Monoidal (MonoidalMap)
29+import qualified Data.AppendMap as Map
30+import Data.MonoidMap (MonoidMap (..), monoidMap)
31+import Data.Foldable
32+import Data.IORef
33+import Debug.Trace
34+import qualified Network.WebSockets as WS
35+import qualified Network.WebSockets.Connection as WS
36+import qualified Network.WebSockets.Stream as WS
37+import Network.WebSockets.Snap
38+import Control.Exception
39+
40+-- | A way for a pipeline to retrieve data
41+newtype QueryHandler q m = QueryHandler
42+ { runQueryHandler :: q -> m (QueryResult q) }
43+
44+-- | A way to push data into a pipeline
45+newtype Recipient q m = Recipient
46+ { tellRecipient :: QueryResult q -> m () }
47+
48+-- | A way of attaching to (and later detaching from) a pipeline
49+newtype Registrar q = Registrar { unRegistrar :: Recipient q IO -> IO (QueryHandler q IO, IO ()) }
50+
51+-- | A way of connecting a source of data to a consumer of data. The consumer
52+-- can pull data from the datasource and the datasource can push data to the
53+-- consumer.
54+--
55+-- q is the consumer side
56+-- q' is the datasource side
57+newtype Pipeline m q q' = Pipeline { unPipeline :: QueryHandler q' m -> Recipient q m -> IO (QueryHandler q m, Recipient q' m) }
58+
59+instance Category (Pipeline m) where
60+ id = Pipeline $ \qh r -> return (qh, r)
61+ Pipeline f . Pipeline g = Pipeline $ \qh r -> do
62+ rec (qh'', r') <- g qh' r
63+ (qh', r'') <- f qh r'
64+ return (qh'', r'')
65+
66+-- | A key used to track particular consumers
67+newtype ClientKey = ClientKey { unClientKey :: Integer }
68+ deriving (Eq, Ord, Show)
69+
70+-- | Produce a multi-client 'Recipient' and a single client 'QueryHandler'
71+fanQuery
72+ :: forall k q. (Ord k, Monoid (QueryResult q))
73+ => (k -> IO (Recipient q IO))
74+ -- ^ Look up a recipient
75+ -> QueryHandler (MonoidalMap k q) IO
76+ -- ^ A 'QueryHandler' for multiple clients
77+ -> ( Recipient (MonoidalMap k q) IO
78+ -- Used to notify recipients of new 'QueryResult' data
79+ , k -> QueryHandler q IO
80+ -- Used by 'multiplexQuery' to lookup the 'QueryHandler' for a given client
81+ )
82+fanQuery lookupRecipient qh = (multiRecipient lookupRecipient, fanQueryHandler qh)
83+ where
84+ -- | Given a recipient lookup function, produces a 'Recipient' that can
85+ -- handle distribution of a map of 'QueryResult's to many recipients
86+ multiRecipient lookupR = Recipient $ imapM_ $ \k qr -> do
87+ s <- lookupR k
88+ tellRecipient s qr
89+ -- | Given a multi-client 'QueryHandler', produces a 'QueryHandler' for
90+ -- a single client
91+ fanQueryHandler qh' = \k -> mapQueryHandler (singletonQuery k) qh'
92+
93+-- | Maintains a Map from connected clients to a 'Recipient' that can be used to transmit data to the clients
94+--
95+-- Takes a function that looks up a the 'QueryHandler' for a given client
96+-- Returns: 1. A lookup function into the Map of clients
97+-- 2. A way to register a new client
98+-- a. A 'QueryHandler' for the newly registered client
99+-- b. A removal callback to de-register a particular client
100+multiplexQuery
101+ :: (MonadIO m, Group q)
102+ => (ClientKey -> QueryHandler q m)
103+ -> IO ( ClientKey -> IO (Recipient q m)
104+ , Recipient q m -> IO (QueryHandler q m, m ())
105+ )
106+multiplexQuery lookupQueryHandler = do
107+ clients <- newIORef (ClientKey 0, Map.empty :: MonoidalMap ClientKey (Recipient q m, q))
108+ let
109+ lookupRecipient k = do
110+ (_, cs) <- readIORef clients
111+ case Map.lookup k cs of
112+ Nothing -> do
113+ putStrLn $ mconcat
114+ [ "Rhyolite.Backend.App.multiplexQuery: Failed to find sender for key"
115+ , show k
116+ , " in known keys "
117+ , show $ Map.keys cs
118+ ]
119+ return $ Recipient $ const $ return ()
120+ Just s -> return $ fst s
121+
122+ registerRecipient s = do
123+ cid <- atomicModifyIORef' clients $ \(cid, recipients) ->
124+ ((ClientKey (unClientKey cid + 1), Map.insert cid (s, mempty) recipients), cid)
125+ let
126+ queryHandler = QueryHandler $ \q -> do
127+ liftIO $ atomicModifyIORef' clients $ \(nextCid, recipients) ->
128+ ((nextCid, Map.update (\(r, oldQ) -> Just (r, oldQ <> q)) cid recipients), ())
129+ runQueryHandler (lookupQueryHandler cid) q
130+
131+ unregisterRecipient = do
132+ antiQ <- liftIO $ atomicModifyIORef' clients $ \(nextCid, recipients) ->
133+ case Map.updateLookupWithKey (\_ _ -> Nothing) cid recipients of
134+ (Nothing, _) -> trace
135+ ("Rhyolite.Backend.App.multiplexQuery: Tried to unregister a client key that is not registered " <> show cid)
136+ ((nextCid, recipients), mempty)
137+ (Just (_, removedQuery), newRecipients) -> ((nextCid, newRecipients), negateG removedQuery)
138+
139+ -- TODO: Should we have a way of ensuring that this doesn't actually cause a query to be run?
140+ -- It shouldn't cause the query to be run again but it depends on if the callee will notice
141+ -- that the new query is strictly smaller than the old one.
142+ _ <- runQueryHandler (lookupQueryHandler cid) antiQ
143+ return ()
144+
145+ return (queryHandler, unregisterRecipient)
146+
147+ return (lookupRecipient, registerRecipient)
148+
149+-- | Represents a WebSocket message from one of two channels: ViewSelector declarations or API requests
150+data WebSocketRequest q r
151+ = WebSocketRequest_ViewSelector q
152+ | WebSocketRequest_Api (TaggedRequest r)
153+ deriving (Typeable, Generic)
154+
155+instance (Request r, FromJSON q) => FromJSON (WebSocketRequest q r)
156+instance (Request r, ToJSON q) => ToJSON (WebSocketRequest q r)
157+
158+-- | Represents a WebSocket response from one of three channels: incoming 'View's, API responses, or version info
159+data WebSocketResponse q
160+ = WebSocketResponse_View (QueryResult q)
161+ | WebSocketResponse_Api TaggedResponse
162+ | WebSocketResponse_Version Text
163+ deriving (Typeable, Generic)
164+
165+instance FromJSON (QueryResult q) => FromJSON (WebSocketResponse q)
166+instance ToJSON (QueryResult q) => ToJSON (WebSocketResponse q)
167+
168+-- | A request tagged with an identifier
169+data TaggedRequest r = TaggedRequest Value (SomeRequest r)
170+ deriving (Typeable, Generic)
171+
172+instance Request r => FromJSON (TaggedRequest r)
173+instance Request r => ToJSON (TaggedRequest r)
174+
175+-- | A response tagged with an identifier matching the one in the 'TaggedRequest'. The identifier is the first argument.
176+data TaggedResponse = TaggedResponse Value Value
177+ deriving (Typeable, Generic)
178+
179+instance FromJSON TaggedResponse
180+instance ToJSON TaggedResponse
181+
182+-- | Handles a websocket connection
183+handleWebsocket
184+ :: forall q qr r.
185+ ( Eq (q SelectedCount)
186+ , Eq (qr SelectedCount)
187+ , QueryResult (q SelectedCount) ~ qr SelectedCount
188+ , QueryResult (q ()) ~ qr ()
189+ , ToJSON (qr ())
190+ , Functor qr
191+ , Group (q SelectedCount)
192+ , Monoid (qr SelectedCount)
193+ , FromJSON (q ())
194+ , Request r
195+ , Functor q
196+ )
197+ => Text -- ^ Version
198+ -> (forall a. r a -> IO a)
199+ -> Registrar (q SelectedCount)
200+ -> Snap ()
201+handleWebsocket v rh register = withWebsocketsConnection $ \conn -> do
202+ let sender = Recipient $ sendEncodedDataMessage conn . (\a -> WebSocketResponse_View (void a) :: WebSocketResponse (q ()))
203+ sendEncodedDataMessage conn (WebSocketResponse_Version v :: WebSocketResponse (q ()))
204+ bracket (unRegistrar register sender) snd $ \(vsHandler, _) -> flip evalStateT mempty $ forever $ do
205+ (wsr :: WebSocketRequest (q ()) r) <- liftIO $ getDataMessage conn
206+ case wsr of
207+ WebSocketRequest_Api (TaggedRequest reqId (SomeRequest req)) -> lift $ do
208+ a <- rh req
209+ sendEncodedDataMessage conn
210+ (WebSocketResponse_Api $ TaggedResponse reqId (toJSON a) :: WebSocketResponse (q ()))
211+ WebSocketRequest_ViewSelector new -> do
212+ old <- get
213+ let new' = SelectedCount 1 <$ new
214+ vsDiff = new' ~~ old
215+ when (vsDiff /= mempty) $ do
216+ qr <- lift $ runQueryHandler vsHandler vsDiff
217+ put new'
218+ when (qr /= mempty) $ lift $
219+ sendEncodedDataMessage conn (WebSocketResponse_View (void qr) :: WebSocketResponse (q ()))
220+
221+-------------------------------------------------------------------------------
222+
223+-- | Connect a datasource (e.g., a database) to a pipeline
224+--
225+-- Data taken from 'getNextNotification' is pushed into the pipeline and
226+-- when the pipeline pulls data, it is retrieved using 'qh'
227+feedPipeline
228+ :: Monoid q
229+ => IO (q -> IO (QueryResult q))
230+ -- ^ Get the next notification to be sent to the pipeline. If no notification
231+ -- is available, this should block until one is available
232+ -> QueryHandler q IO
233+ -- ^ Retrieve data when requested by pipeline
234+ -> Recipient q IO
235+ -- ^ A way to push data into the pipeline
236+ -> IO (QueryHandler q IO, IO ())
237+ -- ^ A way for the pipeline to request data
238+feedPipeline getNextNotification qh r = do
239+ currentQuery <- newIORef mempty
240+ let qhSaveQuery = QueryHandler $ \q -> do
241+ atomicModifyIORef' currentQuery $ \old -> (q <> old, ())
242+ runQueryHandler qh q
243+ tid <- forkIO . forever $ do
244+ nm <- getNextNotification
245+ q <- readIORef currentQuery
246+ qr <- nm q
247+ tellRecipient r qr
248+ return (qhSaveQuery, killThread tid)
249+
250+-- | Connects the pipeline to websockets consumers
251+connectPipelineToWebsockets
252+ :: ( Eq (q SelectedCount)
253+ , QueryResult (q SelectedCount) ~ qr SelectedCount
254+ , QueryResult (q ()) ~ qr ()
255+ , Group (q SelectedCount)
256+ , Monoid (qr SelectedCount)
257+ , Eq (qr SelectedCount)
258+ , ToJSON (qr ())
259+ , FromJSON (q ())
260+ , Request r
261+ , Functor qr
262+ , Functor q
263+ )
264+ => Text
265+ -> (forall a. r a -> IO a)
266+ -- ^ API handler
267+ -> QueryHandler (MonoidalMap ClientKey (q SelectedCount)) IO
268+ -- ^ A way to retrieve more data for each consumer
269+ -> IO (Recipient (MonoidalMap ClientKey (q SelectedCount)) IO, Snap ())
270+ -- ^ A way to send data to many consumers and a handler for websockets connections
271+connectPipelineToWebsockets ver rh qh = do
272+ (allRecipients, registerRecipient) <- connectPipelineToWebsockets' qh
273+ return (allRecipients, handleWebsocket ver rh registerRecipient)
274+
275+-- | Like 'connectPipelineToWebsockets' but returns a Registrar that can
276+-- be used to construct a handler for a particular client
277+connectPipelineToWebsockets'
278+ :: (Monoid (QueryResult q), Group q)
279+ => QueryHandler (MonoidalMap ClientKey q) IO
280+ -> IO (Recipient (MonoidalMap ClientKey q) IO, Registrar q)
281+ -- ^ A way to send data to many consumers, and a way to register new consumers
282+connectPipelineToWebsockets' qh = do
283+ rec (lookupRecipient, registerRecipient) <- multiplexQuery clientQueryHandler
284+ let (allRecipients, clientQueryHandler) = fanQuery lookupRecipient qh
285+ return (allRecipients, Registrar registerRecipient)
286+
287+-- | Extends a 'Registrar' with a 'Pipeline'
288+extendRegistrar :: Pipeline IO q q' -> Registrar q' -> Registrar q
289+extendRegistrar (Pipeline p) (Registrar r) = Registrar $ \recipient -> do
290+ rec (qh', close) <- r recipient'
291+ (qh, recipient') <- p qh' recipient
292+ return (qh, close)
293+
294+-------------------------------------------------------------------------------
295+
296+singletonQuery
297+ :: ( Monoid (QueryResult q)
298+ , Ord k
299+ )
300+ => k
301+ -> QueryMorphism q (MonoidalMap k q)
302+singletonQuery k = QueryMorphism { _queryMorphism_mapQuery = Map.singleton k
303+ , _queryMorphism_mapQueryResult = Map.findWithDefault mempty k
304+ }
305+
306+monoidMapQueryMorphism :: (Eq q, Monoid q) => QueryMorphism (MonoidalMap k q) (MonoidMap k q)
307+monoidMapQueryMorphism = QueryMorphism
308+ { _queryMorphism_mapQuery = monoidMap
309+ , _queryMorphism_mapQueryResult = unMonoidMap
310+ }
311+
312+transposeMonoidMap
313+ :: forall a a' k q qr.
314+ ( Ord k
315+ , Eq a
316+ , Monoid a
317+ , Monoid a'
318+ , Eq (q (MonoidMap k a))
319+ , Eq (QueryResult (q a))
320+ , Foldable qr
321+ , Functor q
322+ , Functor qr
323+ , FunctorMaybe qr
324+ , Monoid (q (MonoidMap k a))
325+ , Monoid (QueryResult (q a))
326+ , QueryResult (q (MonoidMap k a)) ~ qr (MonoidMap k a')
327+ , QueryResult (q a) ~ qr a'
328+ )
329+ => QueryMorphism (MonoidMap k (q a)) (q (MonoidMap k a))
330+transposeMonoidMap = QueryMorphism
331+ { _queryMorphism_mapQuery = aggregateQueries
332+ , _queryMorphism_mapQueryResult = distributeResults
333+ }
334+ where
335+ aggregateQueries :: MonoidMap k (q a) -> q (MonoidMap k a)
336+ aggregateQueries = fold . monoidMap . Map.mapWithKey (\k q -> fmap (monoidMap . Map.singleton k) q) . unMonoidMap
337+ distributeResults :: qr (MonoidMap k a') -> MonoidMap k (qr a')
338+ distributeResults v = monoidMap $ Map.mapWithKey (\k _ -> fmapMaybe (Map.lookup k . unMonoidMap) v) $ fold $ fmap unMonoidMap v
339+
340+mapQueryHandlerAndRecipient
341+ :: Functor f
342+ => QueryMorphism q q'
343+ -> QueryHandler q' f
344+ -> Recipient q f
345+ -> (QueryHandler q f, Recipient q' f)
346+mapQueryHandlerAndRecipient qm qh s = (mapQueryHandler qm qh, mapRecipient qm s)
347+
348+mapQueryHandler :: Functor f => QueryMorphism q q' -> QueryHandler q' f -> QueryHandler q f
349+mapQueryHandler qm qh = QueryHandler $ \qs -> mapQueryResult qm <$> runQueryHandler qh (mapQuery qm qs)
350+
351+mapRecipient :: QueryMorphism q q' -> Recipient q m -> Recipient q' m
352+mapRecipient qm s = Recipient $ \qr -> tellRecipient s $ mapQueryResult qm qr
353+
354+fmapQueryMorphism
355+ :: ( Functor f, QueryResult (f q) ~ f (QueryResult q)
356+ , QueryResult (f q') ~ f (QueryResult q') )
357+ => QueryMorphism q q'
358+ -> QueryMorphism (f q) (f q')
359+fmapQueryMorphism qm = QueryMorphism
360+ { _queryMorphism_mapQuery = fmap $ _queryMorphism_mapQuery qm
361+ , _queryMorphism_mapQueryResult = fmap $ _queryMorphism_mapQueryResult qm
362+ }
363+
364+queryMorphismPipeline :: Functor m => QueryMorphism q q' -> Pipeline m q q'
365+queryMorphismPipeline qm = Pipeline $ \qh r -> pure $ mapQueryHandlerAndRecipient qm qh r
366+
367+-- | Accepts a websockets connection and runs the supplied action with it
368+withWebsocketsConnection :: MonadSnap m => (WS.Connection -> IO ()) -> m ()
369+withWebsocketsConnection f = runWebSocketsSnap $ \pc -> do
370+ conn <- WS.acceptRequest pc
371+ handleSomeException $ handleConnectionException pc $ f conn
372+ where
373+ handleSomeException = handle $ \(SomeException e) -> putStrLn $ "withWebsocketsConnection: " <> displayException e
374+ handleConnectionException pc = handle $ \e -> case e of
375+ WS.ConnectionClosed -> return ()
376+ WS.CloseRequest _ _ -> print e >> WS.close (WS.pendingStream pc) >> throwIO e
377+ _ -> do putStr $ "withWebsocketsConnection: Exception: " <> displayException e
378+ throwIO e
379+
380+-- | Attempts to json decode a websockets data message
381+decodeWebsocketsDataMessage :: FromJSON a => WS.DataMessage -> Either String a
382+decodeWebsocketsDataMessage dm = eitherDecode' $ case dm of
383+ WS.Text r' _ -> r'
384+ WS.Binary r' -> r'
385+
386+-- | Parse and process a single websocket data message
387+getDataMessage
388+ :: FromJSON a
389+ => WS.Connection
390+ -> IO a
391+getDataMessage conn = do
392+ dm <- WS.receiveDataMessage conn
393+ case decodeWebsocketsDataMessage dm of
394+ Left err -> liftIO $ throwIO $ AssertionFailed $ mconcat
395+ [ "getDataMessage: error: "
396+ , err
397+ , "; received: "
398+ , show dm
399+ ]
400+ Right a -> return a
401+
402+-- | Send a json encoded data message over the websocket connection
403+sendEncodedDataMessage
404+ :: ToJSON a
405+ => WS.Connection
406+ -> a
407+ -> IO ()
408+sendEncodedDataMessage conn = WS.sendDataMessage conn . (\x -> WS.Text x Nothing) . encode