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