diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,5 @@ +# Revision history for greskell-websocket + +## 0.1.0.0 -- 2018-06-21 + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2018, Toshio Ito + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Toshio Ito nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,6 @@ +# greskell-websocket + + +## Author + +Toshio Ito diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/greskell-websocket.cabal b/greskell-websocket.cabal new file mode 100644 --- /dev/null +++ b/greskell-websocket.cabal @@ -0,0 +1,114 @@ +name: greskell-websocket +version: 0.1.0.0 +author: Toshio Ito +maintainer: Toshio Ito +license: BSD3 +license-file: LICENSE +synopsis: Haskell client for Gremlin Server using WebSocket serializer +description: Haskell client for [Gremlin Server](http://tinkerpop.apache.org/docs/current/reference/#gremlin-server) using WebSocket serializer. + See [README.md](https://github.com/debug-ito/greskell/blob/master/README.md) for detail. + . + This package is based on [greskell-core](http://hackage.haskell.org/package/greskell-core), + and best used with [greskell](http://hackage.haskell.org/package/greskell) package. +category: Network +cabal-version: >= 1.10 +build-type: Simple +extra-source-files: README.md, ChangeLog.md, + test/samples/*.json +homepage: https://github.com/debug-ito/greskell/ +bug-reports: https://github.com/debug-ito/greskell/issues/ + +library + default-language: Haskell2010 + hs-source-dirs: src + ghc-options: -Wall -fno-warn-unused-imports + -- default-extensions: + other-extensions: OverloadedStrings, DuplicateRecordFields, + DeriveGeneric, PartialTypeSignatures, FlexibleContexts, + CPP + exposed-modules: Network.Greskell.WebSocket, + Network.Greskell.WebSocket.Connection, + Network.Greskell.WebSocket.Connection.Settings, + Network.Greskell.WebSocket.Request, + Network.Greskell.WebSocket.Request.Common, + Network.Greskell.WebSocket.Request.Standard, + Network.Greskell.WebSocket.Request.Session, + Network.Greskell.WebSocket.Response, + Network.Greskell.WebSocket.Codec, + Network.Greskell.WebSocket.Codec.JSON + Network.Greskell.WebSocket.Client, + Network.Greskell.WebSocket.Client.Options + other-modules: Network.Greskell.WebSocket.Request.Aeson, + Network.Greskell.WebSocket.Connection.Impl, + Network.Greskell.WebSocket.Connection.Type, + Network.Greskell.WebSocket.Client.Impl, + Network.Greskell.WebSocket.Util + build-depends: base >=4.9.1.0 && <4.12, + greskell-core >=0.1.2.0 && <0.2, + bytestring >=0.10.8.1 && <0.11, + base64-bytestring >=1.0.0.1 && <1.1, + text >=1.2.2.2 && <1.3, + aeson >=1.1.2.0 && <1.5, + unordered-containers >=0.2.8 && <0.3, + uuid >=1.3.13 && <1.4, + async >=2.1.1.1 && <2.3, + stm >=2.4.4.1 && <2.5, + websockets >=0.10 && <0.13, + hashtables >=1.2.2.1 && <1.3, + safe-exceptions >=0.1.6 && <0.2, + vector >=0.12.0.1 && <0.13 + +-- executable greskell-websocket +-- default-language: Haskell2010 +-- hs-source-dirs: app +-- main-is: Main.hs +-- ghc-options: -Wall -fno-warn-unused-imports +-- -- other-modules: +-- -- default-extensions: +-- -- other-extensions: +-- build-depends: base + +test-suite spec + type: exitcode-stdio-1.0 + default-language: Haskell2010 + hs-source-dirs: test + ghc-options: -Wall -fno-warn-unused-imports "-with-rtsopts=-M512m" + main-is: Spec.hs + -- default-extensions: + other-extensions: OverloadedStrings, DuplicateRecordFields, + NoMonomorphismRestriction + other-modules: Network.Greskell.WebSocket.Codec.JSONSpec + build-depends: base, greskell-websocket, aeson, uuid, bytestring, unordered-containers, vector, greskell-core, + hspec >=2.4.4 + + +flag server-test + description: Do test with a Gremlin Server + default: False + +test-suite server-test + type: exitcode-stdio-1.0 + default-language: Haskell2010 + hs-source-dirs: test + ghc-options: -Wall -fno-warn-unused-imports "-with-rtsopts=-M512m" + -threaded + main-is: ServerTest.hs + -- default-extensions: + -- other-extensions: + other-modules: TestUtil.TCounter, + TestUtil.Env, + TestUtil.MockServer, + ServerTest.Connection, + ServerTest.Client + if flag(server-test) + build-depends: base, greskell-websocket, greskell-core, + aeson, uuid, unordered-containers, text, async, safe-exceptions, + websockets, bytestring, stm, vector, + hspec + + else + buildable: False + +source-repository head + type: git + location: https://github.com/debug-ito/greskell.git diff --git a/src/Network/Greskell/WebSocket.hs b/src/Network/Greskell/WebSocket.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket.hs @@ -0,0 +1,38 @@ +-- | +-- Module: Network.Greskell.WebSocket +-- Description: Client of Gremlin Server using WebSocket +-- Maintainer: Toshio Ito +-- +module Network.Greskell.WebSocket + ( module Network.Greskell.WebSocket.Client + -- $doc + ) where + +import Network.Greskell.WebSocket.Client + +-- $doc +-- +-- Client for Gremlin Server using the WebSocket serializer. For +-- example, see the project +-- [README.md](https://github.com/debug-ito/greskell#submit-to-the-gremlin-server) +-- +-- End-users usually only have to use +-- "Network.Greskell.WebSocket.Client", so this module re-exports only +-- that module. +-- +-- Other modules are low-level implementation and for advanced uses. +-- +-- - "Network.Greskell.WebSocket.Connection": Connection to the +-- Gremlin Server implementing the Driver protocol described in +-- . +-- - "Network.Greskell.WebSocket.Codec": Encoder and decoder of +-- RequestMessage and ResponseMessage. +-- - "Network.Greskell.WebSocket.Request": RequestMessage object sent +-- to Gremlin Server. +-- - "Network.Greskell.WebSocket.Request.Standard": Request objects +-- for Standard OpProcessor. +-- - "Network.Greskell.WebSocket.Request.Session": Request objects for +-- Session OpProcessor. +-- - "Network.Greskell.WebSocket.Response": ResponseMessage object +-- returned from Gremlin Server. + diff --git a/src/Network/Greskell/WebSocket/Client.hs b/src/Network/Greskell/WebSocket/Client.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Client.hs @@ -0,0 +1,30 @@ +-- | +-- Module: Network.Greskell.WebSocket.Client +-- Description: High-level interface to Gremlin Server +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Client + ( -- * Make a Client + connect, + connectWith, + close, + Client, + Host, + Port, + -- ** Options for Client + module Network.Greskell.WebSocket.Client.Options, + -- * Submit evaluation requests + submit, + submitRaw, + ResultHandle, + nextResult, + nextResultSTM, + slurpResults, + -- * Exceptions + SubmitException(..) + ) where + +import Network.Greskell.WebSocket.Client.Impl +import Network.Greskell.WebSocket.Client.Options +import Network.Greskell.WebSocket.Connection (Host, Port) diff --git a/src/Network/Greskell/WebSocket/Client/Impl.hs b/src/Network/Greskell/WebSocket/Client/Impl.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Client/Impl.hs @@ -0,0 +1,215 @@ +{-# LANGUAGE DuplicateRecordFields, TypeFamilies #-} +-- | +-- Module: Network.Greskell.WebSocket.Client.Impl +-- Description: +-- Maintainer: Toshio Ito +-- +-- __Internal module__. It's like +-- "Network.Greskell.WebSocket.Connection.Impl". +module Network.Greskell.WebSocket.Client.Impl + where + +import Control.Applicative ((<$>), (<*>)) +import Control.Concurrent.STM + ( STM, atomically, + TVar, newTVarIO, readTVar, writeTVar + ) +import Control.Exception.Safe + ( throw, Typeable, Exception, SomeException, catch + ) +import Data.Aeson (Object) +import qualified Data.Aeson as Aeson +import qualified Data.Aeson.Types as Aeson (Parser) +import Data.Greskell.Greskell (ToGreskell(GreskellReturn), toGremlin) +import Data.Greskell.GraphSON (GraphSON, gsonValue, GValue, FromGraphSON(..), parseEither) +import Data.Greskell.AsIterator (AsIterator(IteratorItem)) +import Data.Monoid (mempty) +import Data.Vector (Vector, (!)) +import Data.Text (Text) +import Data.Traversable (traverse) +import Data.Vector (Vector) + +import Network.Greskell.WebSocket.Client.Options (Options) +import qualified Network.Greskell.WebSocket.Client.Options as Opt +import Network.Greskell.WebSocket.Connection + ( Host, Port, Connection, ResponseHandle + ) +import qualified Network.Greskell.WebSocket.Connection as Conn +import qualified Network.Greskell.WebSocket.Request.Standard as ReqStd +import Network.Greskell.WebSocket.Response (ResponseCode, ResponseMessage) +import qualified Network.Greskell.WebSocket.Response as Res +import Network.Greskell.WebSocket.Util (slurp) + + +-- | A client that establishes a connection to the Gremlin Server. You +-- can send Gremlin expression for evaluation by 'submit' function. +data Client = + Client + { clientOpts :: Options, + clientConn :: Connection GValue + } + +-- | Create a 'Client' to a Gremlin Server, with the default 'Options'. +connect :: Host -> Port -> IO Client +connect = connectWith Opt.defOptions + +-- | Create a 'Client' to a Gremlin Server. +connectWith :: Options -> Host -> Port -> IO Client +connectWith opts host port = do + conn <- Conn.connect (Opt.connectionSettings opts) host port + return $ Client { clientOpts = opts, + clientConn = conn + } + +-- | Close the connection to the server and release other resources of +-- the 'Client'. +close :: Client -> IO () +close c = Conn.close $ clientConn c + +data HandleState = + HandleOpen + | HandleClose + | HandleError SomeException + deriving (Show) + +-- | A handle to receive the result of evaluation of a Gremlin script +-- from the server. +data ResultHandle v = + ResultHandle + { rhResHandle :: ResponseHandle GValue, + rhParseGValue :: GValue -> Either String (Vector v), + rhResultCache :: TVar (Vector v), + rhNextResultIndex :: TVar Int, + rhState :: TVar HandleState + } + +submitBase :: FromGraphSON r => Client -> Text -> Maybe Object -> IO (ResultHandle r) +submitBase client script bindings = do + rh <- Conn.sendRequest conn op + (cache, index, state) <- (,,) <$> newTVarIO mempty <*> newTVarIO 0 <*> newTVarIO HandleOpen + return $ ResultHandle { rhResHandle = rh, + rhParseGValue = parseEither, + rhResultCache = cache, + rhNextResultIndex = index, + rhState = state + } + where + conn = clientConn client + opts = clientOpts client + op = ReqStd.OpEval { ReqStd.batchSize = Opt.batchSize opts, + ReqStd.gremlin = script, + ReqStd.bindings = bindings, + ReqStd.language = Opt.language opts, + ReqStd.aliases = Opt.aliases opts, + ReqStd.scriptEvaluationTimeout = Opt.scriptEvaluationTimeout opts + } + +-- | Submit a Gremlin script to the server. You can get its results by +-- 'ResultHandle'. The result type @v@ is determined by the script +-- type @g@. +-- +-- Usually this function does not throw any exception. Exceptions +-- about sending requests are reported when you operate on +-- 'ResultHandle'. +submit :: (ToGreskell g, r ~ GreskellReturn g, AsIterator r, v ~ IteratorItem r, FromGraphSON v) + => Client + -> g -- ^ Gremlin script + -> Maybe Object -- ^ bindings + -> IO (ResultHandle v) +submit client greskell bindings = submitBase client (toGremlin greskell) bindings + +-- | Less type-safe version of 'submit'. +submitRaw :: Client + -> Text -- ^ Gremlin script + -> Maybe Object -- ^ bindings + -> IO (ResultHandle GValue) +submitRaw = submitBase + +-- | Exception about 'submit' operation and getting its result. +data SubmitException = + ResponseError (ResponseMessage GValue) + -- ^ The server returns a 'ResponseMessage' with error 'ResponseCode'. + | ParseError (ResponseMessage GValue) String + -- ^ The client fails to parse the \"data\" field of the + -- 'ResponseMessage'. The error message is kept in the 'String' + -- field. + deriving (Show,Typeable) + +instance Exception SubmitException + + +-- | Get the next value from the 'ResultHandle'. If you have got all +-- values, it returns 'Nothing'. This function may block for a new +-- response to come. +-- +-- On error, it may throw all sorts of exceptions including +-- 'SubmitException' and 'Conn.RequestException'. For example, if the +-- submitted Gremlin script throws an exception, 'nextResult' throws +-- 'ResponseError' with 'ResponseCode' of 'Res.ScriptEvaluationError'. +nextResult :: ResultHandle v -> IO (Maybe v) +nextResult = atomically . nextResultSTM + +-- | 'STM' version of 'nextResult'. +nextResultSTM :: ResultHandle v -> STM (Maybe v) +nextResultSTM rh = do + cur_state <- readTVar $ rhState rh + case cur_state of + HandleError err -> throw err + HandleClose -> return Nothing + HandleOpen -> doNext `withExceptionSTM` gotoError + where + doNext = do + mret <- getNext + case mret of + Nothing -> writeTVar (rhState rh) HandleClose + _ -> return () + return mret + getNext = do + mnext <- getNextCachedResult rh + case mnext of + Just v -> return $ Just v + Nothing -> loadResponse rh + -- 'withException' function is for MonadMask and STM is not a + -- MonadMask. So we use catch-and-rethrow by hand. + withExceptionSTM main finish = + main `catch` (\ex -> finish ex >> throw ex) + gotoError ex = writeTVar (rhState rh) $ HandleError ex + + +getNextCachedResult :: ResultHandle v -> STM (Maybe v) +getNextCachedResult rh = do + (cache, index) <- (,) <$> (readTVar $ rhResultCache rh) <*> (readTVar $ rhNextResultIndex rh) + if index < length cache + then fromCache cache index + else return Nothing + where + fromCache cache index = do + writeTVar (rhNextResultIndex rh) $ index + 1 + return $ Just (cache ! index) + +loadResponse :: ResultHandle v -> STM (Maybe v) +loadResponse rh = parseResponse =<< (Conn.nextResponseSTM $ rhResHandle rh) + where + parseResponse Nothing = return Nothing + parseResponse (Just res) = + case Res.code $ Res.status res of + Res.Success -> parseData res + Res.NoContent -> return Nothing + Res.PartialContent -> parseData res + _ -> throw $ ResponseError res -- TODO: handle Authenticate code + parseData res = + case rhParseGValue rh $ Res.resultData $ Res.result res of + Left err -> throw $ ParseError res err + Right parsed -> do + writeTVar (rhResultCache rh) parsed + if length parsed == 0 + then do + writeTVar (rhNextResultIndex rh) 0 + return Nothing + else do + writeTVar (rhNextResultIndex rh) 1 + return $ Just (parsed ! 0) + +-- | Get all remaining results from 'ResultHandle'. +slurpResults :: ResultHandle v -> IO (Vector v) +slurpResults h = slurp $ nextResult h diff --git a/src/Network/Greskell/WebSocket/Client/Options.hs b/src/Network/Greskell/WebSocket/Client/Options.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Client/Options.hs @@ -0,0 +1,58 @@ +-- | +-- Module: Network.Greskell.WebSocket.Client.Options +-- Description: Options to create a Client +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Client.Options + ( -- * Options + Options, + defOptions, + -- ** accessor functions + connectionSettings, + batchSize, + language, + aliases, + scriptEvaluationTimeout, + -- * Settings + module Network.Greskell.WebSocket.Connection.Settings + ) where + +import Data.Greskell.GraphSON (GValue) +import Data.HashMap.Strict (HashMap) +import Data.Text (Text) + +import Network.Greskell.WebSocket.Connection (Connection) +import Network.Greskell.WebSocket.Connection.Settings + +-- | Configuration options to create a client for Gremlin Server. +-- +-- You can get the default 'Options' by 'defOptions' function, and +-- customize its fields by accessor functions. +data Options = + Options + { connectionSettings :: !(Settings GValue), + -- ^ Settings for the underlying 'Connection'. Default: + -- 'defJSONSettings'. + batchSize :: !(Maybe Int), + -- ^ \"batchSize\" field for \"eval\" operation. Default: + -- 'Nothing'. + language :: !(Maybe Text), + -- ^ \"language\" field for \"eval\" operation. Default: + -- 'Nothing'. + aliases :: !(Maybe (HashMap Text Text)), + -- ^ \"aliases\" field for \"eval\" operation. Default: 'Nothing'. + scriptEvaluationTimeout :: !(Maybe Int) + -- ^ \"scriptEvaluationTimeout\" field for \"eval\" + -- operation. Default: 'Nothing'. + } + +defOptions :: Options +defOptions = + Options + { connectionSettings = defJSONSettings, + batchSize = Nothing, + language = Nothing, + aliases = Nothing, + scriptEvaluationTimeout = Nothing + } diff --git a/src/Network/Greskell/WebSocket/Codec.hs b/src/Network/Greskell/WebSocket/Codec.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Codec.hs @@ -0,0 +1,69 @@ +-- | +-- Module: Network.Greskell.WebSocket.Codec +-- Description: Encoder\/decoder of Request\/Response +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Codec + ( -- * Codec + Codec(..), + -- * Request encoder + encodeBinaryWith, + messageHeader, + -- * Request decoder + decodeBinary + ) where + +import Control.Monad (when) +import qualified Data.ByteString.Lazy as BSL +import Data.Monoid ((<>)) +import Data.Text (Text) +import Data.Text.Encoding (encodeUtf8, decodeUtf8') + +import Network.Greskell.WebSocket.Request (RequestMessage) +import Network.Greskell.WebSocket.Response (ResponseMessage) + +-- | Encoder of 'RequestMessage' and decoder of 'ResponseMessage', +-- associated with a MIME type. +-- +-- Type @s@ is the body of Response. +data Codec s = + Codec + { mimeType :: Text, -- ^ MIME type sent to the server + encodeWith :: RequestMessage -> BSL.ByteString, -- ^ Request encoder + decodeWith :: BSL.ByteString -> Either String (ResponseMessage s) -- ^ Response decoder + } + +instance Functor Codec where + fmap f c = c { decodeWith = (fmap . fmap . fmap) f $ decodeWith c } + +-- | Make a request message header. +messageHeader :: Text -- ^ MIME type + -> BSL.ByteString +messageHeader mime = BSL.singleton size <> mime_bin + where + size = fromIntegral $ BSL.length mime_bin -- what if 'mime' is too long?? + mime_bin = BSL.fromStrict $ encodeUtf8 mime + +-- | Encode a 'RequestMessage' into a \"binary\" format of Gremlin +-- Server. The result includes the message \"header\" and the +-- \"payload\". +encodeBinaryWith :: Codec s -> RequestMessage -> BSL.ByteString +encodeBinaryWith c req = messageHeader (mimeType c) <> encodeWith c req + +-- | Decode a message in the \"binary\" format. This is mainly for +-- testing purposes. +decodeBinary :: BSL.ByteString + -> Either String (Text, BSL.ByteString) -- ^ (mimeType, payload) +decodeBinary raw_msg = do + case BSL.uncons raw_msg of + Nothing -> Left "Length of MIME type is missing in the header." + Just (mime_len, rest) -> decodeMimeAndPayload mime_len rest + where + decodeMimeAndPayload mime_lenw msg = do + when (BSL.length mime_field /= mime_len) $ Left ("Too short MIME field: " <> show mime_field) + mime_text <- either (Left . show) Right $ decodeUtf8' $ BSL.toStrict $ mime_field + return (mime_text, payload) + where + (mime_field, payload) = BSL.splitAt mime_len msg + mime_len = fromIntegral mime_lenw diff --git a/src/Network/Greskell/WebSocket/Codec/JSON.hs b/src/Network/Greskell/WebSocket/Codec/JSON.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Codec/JSON.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | +-- Module: Network.Greskell.WebSocket.Codec.JSON +-- Description: application\/json codec +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Codec.JSON + ( jsonCodec + ) where + +import Data.Aeson (ToJSON, FromJSON) +import qualified Data.Aeson as A +import Data.Aeson.Types (parseEither) + +import Data.Greskell.GraphSON (FromGraphSON(..)) + +import Network.Greskell.WebSocket.Codec (Codec(..)) + +-- | Simple \"application\/json\" codec. +-- +-- The encoder uses GraphSON v1 format. The decoder supports all +-- GraphSON v1, v2 and v3. +jsonCodec :: (FromGraphSON s) => Codec s +jsonCodec = Codec { mimeType = "application/json", + encodeWith = encode, + decodeWith = decode + } + where + encode = A.encode + decode bs = parseEither parseGraphSON =<< A.eitherDecode' bs diff --git a/src/Network/Greskell/WebSocket/Connection.hs b/src/Network/Greskell/WebSocket/Connection.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Connection.hs @@ -0,0 +1,30 @@ +-- | +-- Module: Network.Greskell.WebSocket.Connection +-- Description: WebSocket Connection to Gremlin Server +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Connection + ( -- * Make a Connection + connect, + close, + Connection, + Host, + Port, + -- ** Settings for Connection + module Network.Greskell.WebSocket.Connection.Settings, + -- * Make a request + sendRequest, + sendRequest', + ResponseHandle, + nextResponse, + nextResponseSTM, + slurpResponses, + -- * Exceptions + GeneralException(..), + RequestException(..) + ) where + +import Network.Greskell.WebSocket.Connection.Impl +import Network.Greskell.WebSocket.Connection.Settings +import Network.Greskell.WebSocket.Connection.Type diff --git a/src/Network/Greskell/WebSocket/Connection/Impl.hs b/src/Network/Greskell/WebSocket/Connection/Impl.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Connection/Impl.hs @@ -0,0 +1,441 @@ +{-# LANGUAGE DuplicateRecordFields, CPP #-} +-- | +-- Module: Network.Greskell.WebSocket.Connection.Impl +-- Description: internal implementation of Connection +-- Maintainer: Toshio Ito +-- +-- This is an internal module. It deliberately exports everything. The +-- upper module is responsible to make a proper export list. +module Network.Greskell.WebSocket.Connection.Impl where + +import Control.Applicative ((<$>), (<|>), empty) +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (withAsync, Async, async, waitCatchSTM, waitAnySTM) +import qualified Control.Concurrent.Async as Async +import Control.Concurrent.STM + ( TBQueue, readTBQueue, newTBQueueIO, writeTBQueue, + TQueue, writeTQueue, newTQueueIO, readTQueue, + TVar, newTVarIO, readTVar, writeTVar, + TMVar, tryPutTMVar, tryReadTMVar, putTMVar, newEmptyTMVarIO, readTMVar, + STM, atomically, retry + ) +import qualified Control.Concurrent.STM as STM +import Control.Exception.Safe + ( Exception(toException), SomeException, withException, throw, try, finally + ) +import Control.Monad (when, void, forM_) +import Data.Aeson (Value) +import qualified Data.ByteString.Lazy as BSL +import Data.Foldable (toList) +import qualified Data.HashTable.IO as HT +import Data.Monoid (mempty) +import Data.Typeable (Typeable) +import Data.UUID (UUID) +import Data.Vector (Vector) +import qualified Network.WebSockets as WS + +import Network.Greskell.WebSocket.Codec (Codec(decodeWith, encodeWith), encodeBinaryWith) +import Network.Greskell.WebSocket.Connection.Settings (Settings) +import qualified Network.Greskell.WebSocket.Connection.Settings as Settings +import Network.Greskell.WebSocket.Connection.Type + ( Connection(..), ConnectionState(..), + ResPack, ReqID, ReqPack(..), RawRes, + GeneralException(..) + ) +import Network.Greskell.WebSocket.Request + ( RequestMessage(RequestMessage, requestId), + Operation, makeRequestMessage + ) +import Network.Greskell.WebSocket.Response + ( ResponseMessage(ResponseMessage, requestId, status), + ResponseStatus(ResponseStatus, code), + isTerminating + ) +import Network.Greskell.WebSocket.Util (slurp) + + +flushTBQueue :: TBQueue a -> STM [a] +#if MIN_VERSION_stm(2,4,5) +flushTBQueue = STM.flushTBQueue +#else +flushTBQueue q = fmap toList $ slurp $ STM.tryReadTBQueue q +#endif + + +-- | Host name or an IP address. +type Host = String + +-- | TCP port number. +type Port = Int + +-- | Make a 'Connection' to a Gremlin Server. +-- +-- If it fails to connect to the specified server, it throws an +-- exception. +connect :: Settings s -> Host -> Port -> IO (Connection s) +connect settings host port = do + req_pool <- HT.new -- Do not manipulate req_pool in this thread. It belongs to runWSConn thread. + qreq <- newTBQueueIO qreq_size + var_connect_result <- newEmptyTMVarIO + var_conn_state <- newTVarIO ConnOpen + ws_thread <- async $ runWSConn settings host port ws_path req_pool qreq var_connect_result var_conn_state + eret <- atomically $ readTMVar var_connect_result + case eret of + Left e -> throw e + Right () -> return $ Connection { connQReq = qreq, + connState = var_conn_state, + connWSThread = ws_thread, + connCodec = codec + } + where + codec = Settings.codec settings + qreq_size = Settings.requestQueueSize settings + ws_path = Settings.endpointPath settings + +-- | Close the 'Connection'. +-- +-- If there are pending requests in the 'Connection', 'close' function +-- blocks for them to complete or time out. +-- +-- Calling 'close' on a 'Connection' already closed (or waiting to +-- close) does nothing. +close :: Connection s -> IO () +close conn = do + need_wait <- atomically $ do + cur_state <- readTVar $ connState conn + case cur_state of + ConnClosed -> return False + ConnClosing -> return True + ConnOpen -> do + writeTVar (connState conn) ConnClosing + return True + if need_wait then waitForClose else return () + where + waitForClose = atomically $ do + cur_state <- readTVar $ connState conn + if cur_state == ConnClosed + then return () + else retry + +type Path = String + +-- | A thread taking care of a WS connection. +runWSConn :: Settings s + -> Host -> Port -> Path + -> ReqPool s -> TBQueue (ReqPack s) + -> TMVar (Either SomeException ()) -> TVar ConnectionState + -> IO () +runWSConn settings host port path req_pool qreq var_connect_result var_conn_state = + (doConnect `withException` reportFatalEx) `finally` finalize + where + doConnect = WS.runClient host port path $ \wsconn -> do + is_success <- checkAndReportConnectSuccess + if not is_success + then return () -- result is already reported at var_connect_result + else setupMux wsconn + setupMux wsconn = do + qres <- newTQueueIO + withAsync (runRxLoop wsconn qres) $ \rx_thread -> + runMuxLoop wsconn req_pool settings qreq qres (readTVar var_conn_state) rx_thread + checkAndReportConnectSuccess = atomically $ do + mret <- tryReadTMVar var_connect_result + case mret of + -- usually, mret should be Nothing. + Nothing -> do + putTMVar var_connect_result $ Right () + return True + Just (Right _) -> return True + Just (Left _) -> return False + reportFatalEx :: SomeException -> IO () + reportFatalEx cause = do + reportToConnectCaller cause + reportToReqPool req_pool cause + reportToQReq qreq cause + reportToConnectCaller cause = void $ atomically $ tryPutTMVar var_connect_result $ Left cause + finalize = do + cleanupReqPool req_pool + atomically $ writeTVar var_conn_state ConnClosed + +reportToReqPool :: ReqPool s -> SomeException -> IO () +reportToReqPool req_pool cause = HT.mapM_ forEntry req_pool + where + forEntry (_, entry) = atomically $ writeTQueue (rpeOutput entry) $ Left cause + +reportToQReq :: TBQueue (ReqPack s) -> SomeException -> IO () +reportToQReq qreq cause = atomically $ do + reqpacks <- flushTBQueue qreq + forM_ reqpacks reportToReqPack + where + reportToReqPack reqp = writeTQueue (reqOutput reqp) $ Left cause + +-- | An exception related to a specific request. +data RequestException = + AlreadyClosed + -- ^ The connection is already closed before it sends the request. + | ServerClosed + -- ^ The server closed the connection before it sends response for + -- this request. + | DuplicateRequestId UUID + -- ^ The requestId (kept in this object) is already pending in the + -- connection. + | ResponseTimeout + -- ^ The server fails to send ResponseMessages within + -- 'Settings.responseTimeout'. + deriving (Show,Eq,Typeable) + +instance Exception RequestException + +data ReqPoolEntry s = + ReqPoolEntry + { rpeReqId :: !ReqID, + rpeOutput :: !(TQueue (ResPack s)), + rpeTimer :: !(Async ReqID) + -- ^ timer thread to time out response. + } + +-- | (requestId of pending request) --> (objects related to that pending request) +type ReqPool s = HT.BasicHashTable ReqID (ReqPoolEntry s) + +-- | Multiplexed event object +data MuxEvent s = EvReq (ReqPack s) + | EvRes RawRes + | EvActiveClose + | EvRxFinish + | EvRxError SomeException + | EvResponseTimeout ReqID + +-- | HashTable's mutateIO is available since 1.2.3.0 +tryInsertToReqPool :: ReqPool s + -> ReqID + -> IO (ReqPoolEntry s) -- ^ action to create the new entry. + -> IO Bool -- ^ 'True' if insertion is successful. +tryInsertToReqPool req_pool rid makeEntry = do + mexist_entry <- HT.lookup req_pool rid + case mexist_entry of + Just _ -> return False + Nothing -> do + new_entry <- makeEntry + HT.insert req_pool rid new_entry + return True + +cleanupReqPoolEntry :: ReqPoolEntry s -> IO () +cleanupReqPoolEntry entry = Async.cancel $ rpeTimer entry + +removeReqPoolEntry :: ReqPool s -> ReqPoolEntry s -> IO () +removeReqPoolEntry req_pool entry = do + cleanupReqPoolEntry entry + HT.delete req_pool $ rpeReqId entry + +cleanupReqPool :: ReqPool s -> IO () +cleanupReqPool req_pool = HT.mapM_ forEntry req_pool + where + forEntry (_, entry) = cleanupReqPoolEntry entry + +getAllResponseTimers :: ReqPool s -> IO [Async ReqID] +getAllResponseTimers req_pool = (fmap . fmap) toTimer $ HT.toList req_pool + where + toTimer (_, entry) = rpeTimer entry + +-- | Multiplexer loop. +runMuxLoop :: WS.Connection -> ReqPool s -> Settings s + -> TBQueue (ReqPack s) -> TQueue RawRes -> STM ConnectionState + -> Async () + -> IO () +runMuxLoop wsconn req_pool settings qreq qres readConnState rx_thread = loop + where + codec = Settings.codec settings + loop = do + res_timers <- getAllResponseTimers req_pool + event <- atomically $ getEventSTM res_timers + case event of + EvReq req -> handleReq req >> loop + EvRes res -> handleRes res >> loop + EvActiveClose -> return () + EvRxFinish -> handleRxFinish + EvRxError e -> throw e + EvResponseTimeout rid -> handleResponseTimeout rid >> loop + getEventSTM res_timers = getRequest + <|> (EvRes <$> readTQueue qres) + <|> makeEvActiveClose + <|> (rxResultToEvent <$> waitCatchSTM rx_thread) + <|> (timeoutToEvent <$> waitAnySTM res_timers) + where + max_concurrency = Settings.concurrency settings + cur_concurrency = length res_timers + getRequest = if cur_concurrency < max_concurrency + then EvReq <$> readTBQueue qreq + else empty + rxResultToEvent (Right ()) = EvRxFinish + rxResultToEvent (Left e) = EvRxError e + timeoutToEvent (_, rid) = EvResponseTimeout rid + makeEvActiveClose = do + if cur_concurrency > 0 + then empty + else do + conn_state <- readConnState + if conn_state == ConnOpen then empty else return EvActiveClose + handleReq req = do + insert_ok <- tryInsertToReqPool req_pool rid makeNewEntry + if insert_ok + then WS.sendBinaryData wsconn $ reqData req + else reportError + where + rid = reqId req + qout = reqOutput req + makeNewEntry = do + timer_thread <- runTimer (Settings.responseTimeout settings) rid + return $ ReqPoolEntry { rpeReqId = rid, + rpeOutput = qout, + rpeTimer = timer_thread + } + reportError = + atomically $ writeTQueue qout $ Left $ toException $ DuplicateRequestId rid + handleRes res = case decodeWith codec res of + Left err -> Settings.onGeneralException settings $ ResponseParseFailure err + Right res_msg -> handleResMsg res_msg + handleResMsg res_msg@(ResponseMessage { requestId = rid }) = do + m_entry <- HT.lookup req_pool rid + case m_entry of + Nothing -> Settings.onGeneralException settings $ UnexpectedRequestId rid + Just entry -> do + when (isTerminatingResponse res_msg) $ do + removeReqPoolEntry req_pool entry + atomically $ writeTQueue (rpeOutput entry) $ Right res_msg + handleRxFinish = do + -- RxFinish is an error for pending requests. If there is no + -- pending requests, it's totally normal. + let ex = toException ServerClosed + reportToReqPool req_pool ex + reportToQReq qreq ex + handleResponseTimeout rid = do + mentry <- HT.lookup req_pool rid + case mentry of + Nothing -> return () -- this case may happen if the response came just before the time-out, I think. + Just entry -> do + atomically $ writeTQueue (rpeOutput entry) $ Left $ toException $ ResponseTimeout + removeReqPoolEntry req_pool entry + + +-- | Receiver thread. It keeps receiving data from WS until the +-- connection finishes cleanly. Basically every exception is raised to +-- the caller. +runRxLoop :: WS.Connection -> TQueue RawRes -> IO () +runRxLoop wsconn qres = loop + where + loop = do + mgot <- tryReceive + case mgot of + Nothing -> return () + Just got -> do + atomically $ writeTQueue qres got + loop + tryReceive = toMaybe =<< (try $ WS.receiveData wsconn) + where + toMaybe (Right d) = return $ Just d + toMaybe (Left e@(WS.CloseRequest close_status _)) = do + if close_status == 1000 -- "normal closure". See sec. 7.4, RFC 6455. + then return Nothing + else throw e + -- We allow the server to close the connection without sending Close request message. + toMaybe (Left WS.ConnectionClosed) = return Nothing + toMaybe (Left e) = throw e + +runTimer :: Int -> ReqID -> IO (Async ReqID) +runTimer wait_sec rid = async $ do + threadDelay $ wait_sec * 1000000 + return rid + + +-- | A handle associated in a 'Connection' for a pair of request and +-- response. You can retrieve 'ResponseMessage's from this object. +-- +-- Type @s@ is the body of the response. +data ResponseHandle s = + ResponseHandle + { rhGetResponse :: STM (ResPack s), + rhTerminated :: TVar Bool + } + +instance Functor ResponseHandle where + fmap f rh = rh { rhGetResponse = (fmap . fmap . fmap) f $ rhGetResponse rh } + + +-- | Make a 'RequestMessage' from an 'Operation' and send it. +-- +-- Usually this function does not throw any exception. Exceptions +-- about sending requests are reported when you operate on +-- 'ResponseHandle'. +sendRequest :: Operation o => Connection s -> o -> IO (ResponseHandle s) +sendRequest conn o = sendRequest' conn =<< makeRequestMessage o + +-- | Like 'sendRequest', but you can pass a 'RequestMessage' directly +-- to this function. +sendRequest' :: Connection s -> RequestMessage -> IO (ResponseHandle s) +sendRequest' conn req_msg = do + qout <- newTQueueIO + is_open <- getConnectionOpen + if is_open + then sendReqPack qout + else reportAlreadyClosed qout + makeResHandle qout + where + codec = connCodec conn + qreq = connQReq conn + var_conn_state = connState conn + rid = requestId (req_msg :: RequestMessage) + getConnectionOpen = fmap (== ConnOpen) $ atomically $ readTVar var_conn_state + sendReqPack qout = do + atomically $ writeTBQueue qreq reqpack + where + reqpack = ReqPack + { reqData = encodeBinaryWith codec req_msg, + reqId = rid, + reqOutput = qout + } + makeResHandle qout = do + var_term <- newTVarIO False + return $ ResponseHandle + { rhGetResponse = readTQueue qout, + rhTerminated = var_term + } + reportAlreadyClosed qout = do + atomically $ writeTQueue qout $ Left $ toException $ AlreadyClosed + + +-- | Get a 'ResponseMessage' from 'ResponseHandle'. If you have +-- already got all responses, it returns 'Nothing'. This function may +-- block for a new 'ResponseMessage' to come. +-- +-- On error, it may throw all sorts of exceptions including +-- 'RequestException'. +nextResponse :: ResponseHandle s -> IO (Maybe (ResponseMessage s)) +nextResponse = atomically . nextResponseSTM + +-- | 'STM' version of 'nextResponse'. +nextResponseSTM :: ResponseHandle s -> STM (Maybe (ResponseMessage s)) +nextResponseSTM rh = do + termed <- readTVar $ rhTerminated rh + if termed + then return Nothing + else readResponse + where + readResponse = do + eres <- rhGetResponse rh + case eres of + Left ex -> throw ex -- throw in STM. The eres is put back to the queue. + Right res -> do + updateTermed res + return $ Just res + updateTermed res = + when (isTerminatingResponse res) $ do + writeTVar (rhTerminated rh) True + +isTerminatingResponse :: ResponseMessage s -> Bool +isTerminatingResponse (ResponseMessage { status = (ResponseStatus { code = c }) }) = + isTerminating c + +-- | Get all remaining 'ResponseMessage's from 'ResponseHandle'. +slurpResponses :: ResponseHandle s -> IO (Vector (ResponseMessage s)) +slurpResponses h = slurp $ nextResponse h + + diff --git a/src/Network/Greskell/WebSocket/Connection/Settings.hs b/src/Network/Greskell/WebSocket/Connection/Settings.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Connection/Settings.hs @@ -0,0 +1,71 @@ +-- | +-- Module: Network.Greskell.WebSocket.Connection.Settings +-- Description: Settings for making Connection +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Connection.Settings + ( -- * Settings + Settings, + defSettings, + defJSONSettings, + -- ** accessor functions + codec, endpointPath, onGeneralException, responseTimeout, + concurrency, requestQueueSize + ) where + +import Data.Greskell.GraphSON (FromGraphSON) + +import Network.Greskell.WebSocket.Codec (Codec) +import Network.Greskell.WebSocket.Codec.JSON (jsonCodec) +import Network.Greskell.WebSocket.Connection.Type (GeneralException) + +import System.IO (stderr, hPutStrLn) + +-- | 'Settings' for making connection to Gremlin Server. +-- +-- You can get the default 'Settings' by 'defSettings' function, and +-- customize its fields by accessor functions. +-- +-- Type @s@ is the body of Response. +data Settings s = + Settings + { codec :: !(Codec s), + -- ^ codec for the connection. + endpointPath :: !String, + -- ^ Path of the WebSocket endpoint. Default: \"/gremlin\" + onGeneralException :: !(GeneralException -> IO ()), + -- ^ An exception handler for 'GeneralException'. This exception + -- is not fatal, so the connection survives after this handler is + -- called. You don't have to re-throw the exception. Default: + -- print the exception to stderr. + responseTimeout :: !Int, + -- ^ Time out (in seconds) for responses. It is the maximum time + -- for which the connection waits for a response to complete after + -- it sends a request. If the response consists of more than one + -- ResponseMessages, the timeout applies to the last of the + -- ResponseMessages. Default: 60 + concurrency :: !Int, + -- ^ Maximum concurrent requests the connection can make to the + -- server. If the client tries to make more concurrent requests + -- than this value, later requests are queued in the connection or + -- the client may be blocked. Default: 4 + requestQueueSize :: !Int + -- ^ Size of the internal queue of requests. Usually you don't + -- need to customize the field. See also 'concurrency'. Default: + -- 8. + } + +defSettings :: Codec s -> Settings s +defSettings c = Settings + { codec = c, + endpointPath = "/gremlin", + onGeneralException = \e -> hPutStrLn stderr $ show e, + responseTimeout = 60, + concurrency = 4, + requestQueueSize = 8 + } + +-- | 'defSettings' with 'jsonCodec'. +defJSONSettings :: FromGraphSON s => Settings s +defJSONSettings = defSettings jsonCodec diff --git a/src/Network/Greskell/WebSocket/Connection/Type.hs b/src/Network/Greskell/WebSocket/Connection/Type.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Connection/Type.hs @@ -0,0 +1,88 @@ +-- | +-- Module: Network.Greskell.WebSocket.Connection.Type +-- Description: common types for Connection +-- Maintainer: Toshio Ito +-- +-- This is an internal module. This defines and exports common types +-- used by Connection modules. The upper module is responsible to +-- limit exports from this module. +module Network.Greskell.WebSocket.Connection.Type + ( RawReq, + RawRes, + ReqID, + ResPack, + ReqPack(..), + ConnectionState(..), + Connection(..), + GeneralException(..) + ) where + +import Control.Concurrent.Async (Async) +import Control.Exception.Safe (SomeException, Typeable, Exception) +import Control.Concurrent.STM (TQueue, TBQueue, TVar) +import qualified Data.ByteString.Lazy as BSL +import Data.UUID (UUID) + +import Network.Greskell.WebSocket.Response (ResponseMessage) +import Network.Greskell.WebSocket.Codec (Codec) + +type RawReq = BSL.ByteString +type RawRes = BSL.ByteString +type ReqID = UUID + +-- | Package of Response data and related stuff. +type ResPack s = Either SomeException (ResponseMessage s) + +-- | Package of request data and related stuff. It's passed from the +-- caller thread into WS handling thread. +data ReqPack s = + ReqPack + { reqData :: !RawReq, + -- ^ Encoded request data + reqId :: !ReqID, + -- ^ request ID + reqOutput :: !(TQueue (ResPack s)) + -- ^ the output queue for incoming response for this request. + } + +-- | State of the 'Connection'. +data ConnectionState = + ConnOpen + -- ^ Connection is open and ready to use. + | ConnClosing + -- ^ Connection is closing. It rejects new requests, but keeps + -- receiving responses for pending requests. When there is no + -- pending requests, it goes to 'ConnClosed'. + | ConnClosed + -- ^ Connection is closed. It rejects requests, and it doesn't + -- expect any responses. It can close the underlying WebSocket + -- connection. + deriving (Show,Eq,Ord,Enum,Bounded) + +-- | A WebSocket connection to a Gremlin Server. +-- +-- Type @s@ is the body of Response, as in 'ResponseMessage'. +data Connection s = + Connection + { connQReq :: !(TBQueue (ReqPack s)), + -- ^ Request queue to WS (Mux) thread. + connState :: !(TVar ConnectionState), + connWSThread :: !(Async ()), + -- ^ WS (Mux) thread. It keeps the underlying WebSocket + -- connection, watches various types of events and responds to + -- those events. + connCodec :: !(Codec s) + } + +-- | Exception general to a 'Connection'. It's not related to specific +-- requests. +data GeneralException = + UnexpectedRequestId UUID + -- ^ Server sends a 'ResponseMessage' with unknown requestId, which + -- is kept in this exception. + | ResponseParseFailure String + -- ^ The 'Connection' fails to parse a data from the server. The + -- error message is kept in this exception. + deriving (Show,Eq,Typeable) + +instance Exception GeneralException diff --git a/src/Network/Greskell/WebSocket/Request.hs b/src/Network/Greskell/WebSocket/Request.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Request.hs @@ -0,0 +1,58 @@ +{-# LANGUAGE DuplicateRecordFields, DeriveGeneric #-} +-- | +-- Module: Network.Greskell.WebSocket.Request +-- Description: Request to Gremlin Server +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Request + ( -- * RequestMessage + RequestMessage(..), + Operation(..), + toRequestMessage, + makeRequestMessage + ) where + +import Control.Applicative ((<$>), (<*>)) +import Data.Aeson (Object, ToJSON(..), FromJSON(..)) +import Data.Text (Text) +import Data.UUID (UUID) +import Data.UUID.V4 (nextRandom) +import GHC.Generics (Generic) + +import qualified Network.Greskell.WebSocket.Request.Aeson as GAeson +import Network.Greskell.WebSocket.Request.Common (Operation(..)) + + +-- | RequestMessage to a Gremlin Server. See +-- . +data RequestMessage = + RequestMessage + { requestId :: !UUID, + op :: !Text, + processor :: !Text, + args :: !Object + } + deriving (Show,Eq,Generic) + +instance ToJSON RequestMessage where + toJSON = GAeson.genericToJSON GAeson.opt + toEncoding = GAeson.genericToEncoding GAeson.opt + +instance FromJSON RequestMessage where + parseJSON = GAeson.genericParseJSON GAeson.opt + +-- | Convert an 'Operation' object to 'RequestMessage'. +toRequestMessage :: Operation o => UUID -> o -> RequestMessage +toRequestMessage rid o = + RequestMessage { requestId = rid, + op = opName o, + processor = opProcessor o, + args = opArgs o + } + +-- | Create a 'RequestMessage' from an 'Operation' object. The +-- 'requestId' is generated by the random number generator of the +-- system. +makeRequestMessage :: Operation o => o -> IO RequestMessage +makeRequestMessage o = toRequestMessage <$> nextRandom <*> pure o diff --git a/src/Network/Greskell/WebSocket/Request/Aeson.hs b/src/Network/Greskell/WebSocket/Request/Aeson.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Request/Aeson.hs @@ -0,0 +1,28 @@ +-- | +-- Module: Network.Greskell.WebSocket.Request.Aeson +-- Description: parser support +-- Maintainer: Toshio Ito +-- +-- __Internal module. End-users should not use this.__ +module Network.Greskell.WebSocket.Request.Aeson + ( genericToJSON, genericToEncoding, genericParseJSON, + opt, + toObject + ) where + +import Data.Aeson + ( genericToJSON, genericToEncoding, genericParseJSON, + ToJSON(..), Object, Value(Object) + ) +import Data.Aeson.Types + ( defaultOptions, omitNothingFields, Options + ) + +opt :: Options +opt = defaultOptions { omitNothingFields = True } + +toObject :: (ToJSON a) => a -> Object +toObject = expectObject . toJSON + where + expectObject (Object o) = o + expectObject _ = error "Expect Object, but got something else" diff --git a/src/Network/Greskell/WebSocket/Request/Common.hs b/src/Network/Greskell/WebSocket/Request/Common.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Request/Common.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | +-- Module: Network.Greskell.WebSocket.Request.Common +-- Description: Common data types for Request objects +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Request.Common + ( Operation(..), + SASLMechanism(..), + Base64(..) + ) where + +import Control.Applicative (empty) +import Data.Aeson (ToJSON(..), FromJSON(..), Object, Value(String)) +import Data.ByteString (ByteString) +import qualified Data.ByteString.Base64 as B64 +import Data.Text (unpack, Text) +import Data.Text.Encoding (decodeUtf8, encodeUtf8) + + +-- | Class of operation objects. +class Operation o where + opProcessor :: o -> Text + -- ^ \"processor\" field. + opName :: o -> Text + -- ^ \"op\" field. + opArgs :: o -> Object + -- ^ \"args\" field. + +instance (Operation a, Operation b) => Operation (Either a b) where + opProcessor e = either opProcessor opProcessor e + opName e = either opName opName e + opArgs e = either opArgs opArgs e + + +-- | Possible SASL mechanisms. +data SASLMechanism = SASLPlain -- ^ \"PLAIN\" SASL + | SASLGSSAPI -- ^ \"GSSAPI\" SASL + deriving (Show,Eq,Ord,Enum,Bounded) + +instance ToJSON SASLMechanism where + toJSON = toJSON . toText + where + toText :: SASLMechanism -> Text + toText SASLPlain = "PLAIN" + toText SASLGSSAPI = "GSSAPI" + +instance FromJSON SASLMechanism where + parseJSON (String s) = case s of + "PLAIN" -> return SASLPlain + "GSSAPI" -> return SASLGSSAPI + _ -> fail ("Unknown SASLMechanism: " ++ unpack s) + parseJSON _ = empty + +-- | A raw 'ByteString' encoded to\/decoded from a base64 text. +-- +-- 'ToJSON' instance encodes the raw 'ByteString' to a base64-encoded +-- 'Text'. 'FromJSON' is its inverse. +newtype Base64 = Base64 { unByte64 :: ByteString } + deriving (Show,Eq,Ord) + +instance ToJSON Base64 where + toJSON (Base64 bs) = toJSON $ decodeUtf8 $ B64.encode bs + +instance FromJSON Base64 where + parseJSON (String t) = either fail (return . Base64) $ B64.decode $ encodeUtf8 t + parseJSON _ = empty + diff --git a/src/Network/Greskell/WebSocket/Request/Session.hs b/src/Network/Greskell/WebSocket/Request/Session.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Request/Session.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE DeriveGeneric, DuplicateRecordFields, OverloadedStrings #-} +-- | +-- Module: Network.Greskell.WebSocket.Request.Session +-- Description: Operation objects for session OpProcessor +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Request.Session + ( -- * OpAuthentication + OpAuthentication(..), + -- * OpEval + SessionID, + OpEval(..), + -- * OpClose + OpClose(..) + ) where + +import Data.Aeson (ToJSON(..), FromJSON(..), Object) +import Data.UUID (UUID) +import Data.Text (Text) +import Data.HashMap.Strict (HashMap) +import GHC.Generics (Generic) + +import qualified Network.Greskell.WebSocket.Request.Aeson as GAeson +import Network.Greskell.WebSocket.Request.Common + (Base64, SASLMechanism, Operation(..)) + +data OpAuthentication = + OpAuthentication + { batchSize :: !(Maybe Int), + sasl :: !Base64, + saslMechanism :: !SASLMechanism + } + deriving (Show,Eq,Ord,Generic) + +instance ToJSON OpAuthentication where + toJSON = GAeson.genericToJSON GAeson.opt + toEncoding = GAeson.genericToEncoding GAeson.opt + +instance FromJSON OpAuthentication where + parseJSON = GAeson.genericParseJSON GAeson.opt + +instance Operation OpAuthentication where + opProcessor _ = "session" + opName _ = "authentication" + opArgs = GAeson.toObject + + +type SessionID = UUID + +data OpEval = + OpEval + { batchSize :: !(Maybe Int), + gremlin :: !Text, + bindings :: !(Maybe Object), + language :: !(Maybe Text), + aliases :: !(Maybe (HashMap Text Text)), + scriptEvaluationTimeout :: !(Maybe Int), + session :: !SessionID, + manageTransaction :: !(Maybe Bool) + } + deriving (Show,Eq,Generic) + +instance ToJSON OpEval where + toJSON = GAeson.genericToJSON GAeson.opt + toEncoding = GAeson.genericToEncoding GAeson.opt + +instance FromJSON OpEval where + parseJSON = GAeson.genericParseJSON GAeson.opt + +instance Operation OpEval where + opProcessor _ = "session" + opName _ = "eval" + opArgs = GAeson.toObject + + +data OpClose = + OpClose + { batchSize :: !(Maybe Int), + session :: !SessionID, + force :: !(Maybe Bool) + } + deriving (Show,Eq,Ord,Generic) + +instance ToJSON OpClose where + toJSON = GAeson.genericToJSON GAeson.opt + toEncoding = GAeson.genericToEncoding GAeson.opt + +instance FromJSON OpClose where + parseJSON = GAeson.genericParseJSON GAeson.opt + +instance Operation OpClose where + opProcessor _ = "session" + opName _ = "close" + opArgs = GAeson.toObject + diff --git a/src/Network/Greskell/WebSocket/Request/Standard.hs b/src/Network/Greskell/WebSocket/Request/Standard.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Request/Standard.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE DeriveGeneric, DuplicateRecordFields, OverloadedStrings #-} +-- | +-- Module: Network.Greskell.WebSocket.Request.Standard +-- Description: Operation objects for standard OpProcessor +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Request.Standard + ( -- * OpAuthentication + OpAuthentication(..), + -- * OpEval + OpEval(..) + ) where + +import Data.Aeson (ToJSON(..), FromJSON(..), Object) +import Data.Text (Text) +import Data.HashMap.Strict (HashMap) +import GHC.Generics (Generic) + +import qualified Network.Greskell.WebSocket.Request.Aeson as GAeson +import Network.Greskell.WebSocket.Request.Common + (Base64, SASLMechanism, Operation(..)) + +data OpAuthentication = + OpAuthentication + { batchSize :: !(Maybe Int), + sasl :: !Base64, + saslMechanism :: !SASLMechanism + } + deriving (Show,Eq,Ord,Generic) + +instance ToJSON OpAuthentication where + toJSON = GAeson.genericToJSON GAeson.opt + toEncoding = GAeson.genericToEncoding GAeson.opt + +instance FromJSON OpAuthentication where + parseJSON = GAeson.genericParseJSON GAeson.opt + +instance Operation OpAuthentication where + opProcessor _ = "" + opName _ = "authentication" + opArgs = GAeson.toObject + +data OpEval = + OpEval + { batchSize :: !(Maybe Int), + gremlin :: !Text, + bindings :: !(Maybe Object), + language :: !(Maybe Text), + aliases :: !(Maybe (HashMap Text Text)), + scriptEvaluationTimeout :: !(Maybe Int) + } + deriving (Show,Eq,Generic) + +instance ToJSON OpEval where + toJSON = GAeson.genericToJSON GAeson.opt + toEncoding = GAeson.genericToEncoding GAeson.opt + +instance FromJSON OpEval where + parseJSON = GAeson.genericParseJSON GAeson.opt + +instance Operation OpEval where + opProcessor _ = "" + opName _ = "eval" + opArgs = GAeson.toObject + diff --git a/src/Network/Greskell/WebSocket/Response.hs b/src/Network/Greskell/WebSocket/Response.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Response.hs @@ -0,0 +1,208 @@ +{-# LANGUAGE DeriveGeneric, OverloadedStrings #-} +-- | +-- Module: Network.Greskell.WebSocket.Response +-- Description: Response from Gremlin Server +-- Maintainer: Toshio Ito +-- +-- +module Network.Greskell.WebSocket.Response + ( -- * ResponseMessage + ResponseMessage(..), + -- * ResponseStatus + ResponseStatus(..), + -- * ResponseResult + ResponseResult(..), + -- * ResponseCode + ResponseCode(..), + codeToInt, + codeFromInt, + isTerminating, + isSuccess, + isClientSideError, + isServerSideError + ) where + +import Control.Applicative ((<$>), (<*>)) +import Data.Aeson + ( Object, ToJSON(..), FromJSON(..), Value(Number, Object), + defaultOptions, genericParseJSON + ) +import Data.Greskell.GraphSON + ( gsonValue, FromGraphSON(..), parseUnwrapAll, (.:), + GValueBody(..) + ) +import Data.Greskell.GraphSON.GValue (gValueBody) +import Data.Text (Text) +import Data.UUID (UUID) +import GHC.Generics (Generic) + + + +-- | Response status code +data ResponseCode = + Success + | NoContent + | PartialContent + | Unauthorized + | Authenticate + | MalformedRequest + | InvalidRequestArguments + | ServerError + | ScriptEvaluationError + | ServerTimeout + | ServerSerializationError + deriving (Show,Eq,Ord,Enum,Bounded) + +codeToInt :: ResponseCode -> Int +codeToInt c = case c of + Success -> 200 + NoContent -> 204 + PartialContent -> 206 + Unauthorized -> 401 + Authenticate -> 407 + MalformedRequest -> 498 + InvalidRequestArguments -> 499 + ServerError -> 500 + ScriptEvaluationError -> 597 + ServerTimeout -> 598 + ServerSerializationError -> 599 + +codeFromInt :: Int -> Maybe ResponseCode +codeFromInt i = case i of + 200 -> Just Success + 204 -> Just NoContent + 206 -> Just PartialContent + 401 -> Just Unauthorized + 407 -> Just Authenticate + 498 -> Just MalformedRequest + 499 -> Just InvalidRequestArguments + 500 -> Just ServerError + 597 -> Just ScriptEvaluationError + 598 -> Just ServerTimeout + 599 -> Just ServerSerializationError + _ -> Nothing + +-- | Returns 'True' if the 'ResponseCode' is a terminating code. +isTerminating :: ResponseCode -> Bool +isTerminating PartialContent = False +isTerminating _ = True + +isCodeClass :: Int -> ResponseCode -> Bool +isCodeClass n c = (codeToInt c `div` 100) == n + +-- | Returns 'True' if the 'ResponseCode' is a success. +-- +-- >>> isSuccess Success +-- True +-- >>> isSuccess Unauthorized +-- False +-- >>> isSuccess ServerError +-- False +isSuccess :: ResponseCode -> Bool +isSuccess = isCodeClass 2 + +-- | Returns 'True' if the 'ResponseCode' is a client-side failure. +-- +-- >>> isClientSideError Success +-- False +-- >>> isClientSideError Unauthorized +-- True +-- >>> isClientSideError ServerError +-- False +isClientSideError :: ResponseCode -> Bool +isClientSideError = isCodeClass 4 + +-- | Returns 'True' if the 'ResponseCode' is a server-side failure. +-- +-- >>> isServerSideError Success +-- False +-- >>> isServerSideError Unauthorized +-- False +-- >>> isServerSideError ServerError +-- True +isServerSideError :: ResponseCode -> Bool +isServerSideError = isCodeClass 5 + +instance FromJSON ResponseCode where + parseJSON (Number n) = maybe err return $ codeFromInt $ floor n + where + err = fail ("Unknown response code: " ++ show n) + parseJSON v = fail ("Expected Number, but got " ++ show v) + +instance FromGraphSON ResponseCode where + parseGraphSON = parseUnwrapAll + +instance ToJSON ResponseCode where + toJSON = toJSON . codeToInt + +-- | \"status\" field. +data ResponseStatus = + ResponseStatus + { code :: !ResponseCode, + message :: !Text, + attributes :: !Object + } + deriving (Show,Eq,Generic) + +instance FromJSON ResponseStatus where + parseJSON v = parseGraphSON =<< parseJSON v + +instance FromGraphSON ResponseStatus where + parseGraphSON gv = case gValueBody gv of + GObject o -> + ResponseStatus + <$> o .: "code" + <*> o .: "message" + <*> o .: "attributes" + gb -> fail ("Expected GObject, but got " ++ show gb) + + +-- | \"result\" field. +data ResponseResult s = + ResponseResult + { resultData :: !s, + -- ^ \"data\" field. + meta :: !Object + } + deriving (Show,Eq,Generic) + +instance FromGraphSON s => FromJSON (ResponseResult s) where + parseJSON v = parseGraphSON =<< parseJSON v + +instance FromGraphSON s => FromGraphSON (ResponseResult s) where + parseGraphSON gv = case gValueBody gv of + GObject o -> + ResponseResult + <$> o .: "data" + <*> o .: "meta" + gb -> fail ("Expected GObject, but got " ++ show gb) + +instance Functor ResponseResult where + fmap f rr = rr { resultData = f $ resultData rr } + +-- | ResponseMessage object from Gremlin Server. See +-- . +-- +-- Type @s@ is the type of the response data. +data ResponseMessage s = + ResponseMessage + { requestId :: !UUID, + status :: !ResponseStatus, + result :: !(ResponseResult s) + } + deriving (Show,Eq,Generic) + +instance FromGraphSON s => FromJSON (ResponseMessage s) where + parseJSON v = parseGraphSON =<< parseJSON v + +instance FromGraphSON s => FromGraphSON (ResponseMessage s) where + parseGraphSON gv = case gValueBody gv of + GObject o -> + ResponseMessage + <$> (o .: "requestId") + <*> (o .: "status") + <*> (o .: "result") + gb -> fail ("Expected GObject, but got " ++ show gb) + +instance Functor ResponseMessage where + fmap f rm = rm { result = fmap f $ result rm } diff --git a/src/Network/Greskell/WebSocket/Util.hs b/src/Network/Greskell/WebSocket/Util.hs new file mode 100644 --- /dev/null +++ b/src/Network/Greskell/WebSocket/Util.hs @@ -0,0 +1,21 @@ +-- | +-- Module: Network.Greskell.WebSocket.Util +-- Description: Common utility +-- Maintainer: Toshio Ito +-- +-- __Internal module__. +module Network.Greskell.WebSocket.Util + ( slurp + ) where + +import Data.Monoid ((<>)) +import qualified Data.Vector as V + +slurp :: Monad m => m (Maybe a) -> m (V.Vector a) +slurp act = go mempty + where + go got = do + mres <- act + case mres of + Nothing -> return got + Just res -> go $! (V.snoc got res) diff --git a/test/Network/Greskell/WebSocket/Codec/JSONSpec.hs b/test/Network/Greskell/WebSocket/Codec/JSONSpec.hs new file mode 100644 --- /dev/null +++ b/test/Network/Greskell/WebSocket/Codec/JSONSpec.hs @@ -0,0 +1,180 @@ +{-# LANGUAGE OverloadedStrings, DuplicateRecordFields, NoMonomorphismRestriction #-} +module Network.Greskell.WebSocket.Codec.JSONSpec (main,spec) where + +import Control.Applicative ((<$>)) +import Control.Monad (forM_) +import Data.Aeson (Value(Null, Number), (.=)) +import qualified Data.Aeson as A +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as BSL +import Data.Greskell.GraphSON + ( GValue, GValueBody(..), + nonTypedGValue, typedGValue' + ) +import Data.Greskell.Greskell (unsafeGreskell, Greskell) +import qualified Data.HashMap.Strict as HM +import Data.Maybe (fromJust) +import Data.Monoid (mempty, (<>)) +import qualified Data.UUID as UUID +import qualified Data.Vector as V +import Test.Hspec + +import Network.Greskell.WebSocket.Request + ( RequestMessage(..), toRequestMessage + ) +import Network.Greskell.WebSocket.Request.Common (SASLMechanism(..), Operation, Base64(..)) +import Network.Greskell.WebSocket.Request.Standard + ( OpAuthentication(..), OpEval(..) + ) +import qualified Network.Greskell.WebSocket.Request.Session as S +import Network.Greskell.WebSocket.Response + (ResponseMessage(..), ResponseStatus(..), ResponseResult(..), ResponseCode(..)) +import Network.Greskell.WebSocket.Codec (Codec(..)) +import Network.Greskell.WebSocket.Codec.JSON (jsonCodec) + +main :: IO () +main = hspec spec + +spec :: Spec +spec = do + decode_spec + encode_spec + +loadSample :: FilePath -> IO BSL.ByteString +loadSample filename = BSL.readFile ("test/samples/" ++ filename) + +loadSampleValue :: FilePath -> IO Value +loadSampleValue filename = do + json_text <- loadSample filename + case A.eitherDecode' json_text of + Left e -> error e + Right v -> return v + +uuidFromString :: String -> UUID.UUID +uuidFromString = fromJust . UUID.fromString + +decode_spec :: Spec +decode_spec = describe "decodeWith" $ do + describe "authentication challenge" $ do + let codec :: Codec Value + codec = jsonCodec + exp_msg = ResponseMessage { requestId = uuidFromString "41d2e28a-20a4-4ab0-b379-d810dede3786", + status = exp_status, + result = exp_result + } + exp_status = ResponseStatus { code = Authenticate, + message = "", + attributes = mempty + } + exp_result = ResponseResult { resultData = Null, + meta = mempty + } + forM_ ["v1", "v2", "v3"] $ \graphson_ver -> specify graphson_ver $ do + got <- decodeWith codec <$> loadSample ("response_auth_" ++ graphson_ver ++ ".json") + got `shouldBe` Right exp_msg + + describe "standard response" $ do + let codec :: Codec GValue + codec = jsonCodec + expMsg d = ResponseMessage { requestId = uuidFromString "41d2e28a-20a4-4ab0-b379-d810dede3786", + status = exp_status, + result = expResult d + } + exp_status = ResponseStatus { code = Success, + message = "", + attributes = mempty + } + expResult d = ResponseResult { resultData = d, + meta = mempty + } + exp_v1 = nonTypedGValue $ GArray $ V.fromList + [ nonTypedGValue $ GObject $ HM.fromList + [ ("id", nonTypedGValue $ GNumber 1), + ("label", nonTypedGValue $ GString "person"), + ("type", nonTypedGValue $ GString "vertex") + ] + ] + exp_typed_vertex = typedGValue' "g:Vertex" $ GObject $ HM.fromList + [ ("id", typedGValue' "g:Int32" $ GNumber 1), + ("label", nonTypedGValue $ GString "person") + ] + exp_v2 = nonTypedGValue $ GArray $ V.fromList [exp_typed_vertex] + exp_v3 = typedGValue' "g:List" $ GArray $ V.fromList [exp_typed_vertex] + sampleFile v = "response_standard_" ++ v ++ ".json" + specify "v1" $ do + got <- decodeWith codec <$> loadSample (sampleFile "v1") + got `shouldBe` Right (expMsg exp_v1) + specify "v2" $ do + got <- decodeWith codec <$> loadSample (sampleFile "v2") + got `shouldBe` Right (expMsg exp_v2) + specify "v3" $ do + got <- decodeWith codec <$> loadSample (sampleFile "v3") + got `shouldBe` Right (expMsg exp_v3) + +encodedValue :: Codec s -> RequestMessage -> Value +encodedValue c req = case A.eitherDecode $ encodeWith c req of + Left e -> error e + Right v -> v + +encodeCase :: String -> RequestMessage -> Spec +encodeCase filename input = specify filename $ do + expected <- loadSampleValue filename + encodedValue codec input `shouldBe` expected + where + codec :: Codec Value + codec = jsonCodec + +encode_spec :: Spec +encode_spec = describe "encodeWith" $ do + encodeCase "request_auth_v1.json" $ toRequestMessage (uuidFromString "cb682578-9d92-4499-9ebc-5c6aa73c5397") + $ OpAuthentication + { batchSize = Nothing, + sasl = Base64 (BS.singleton 0 <> "stephphen" <> BS.singleton 0 <> "password"), + saslMechanism = SASLPlain + } + encodeCase "request_sessionless_eval_v1.json" $ toRequestMessage (uuidFromString "cb682578-9d92-4499-9ebc-5c6aa73c5397") + $ OpEval + { batchSize = Nothing, + gremlin = "g.V(x)", + bindings = Just $ HM.fromList [("x", Number 1)], + language = Just "gremlin-groovy", + aliases = Nothing, + scriptEvaluationTimeout = Nothing + } + encodeCase "request_sessionless_eval_aliased_v1.json" $ toRequestMessage (uuidFromString "cb682578-9d92-4499-9ebc-5c6aa73c5397") + $ OpEval + { batchSize = Nothing, + gremlin = "social.V(x)", + bindings = Just $ HM.fromList [("x", Number 1)], + language = Just "gremlin-groovy", + aliases = Just $ HM.fromList [("g", "social")], + scriptEvaluationTimeout = Nothing + } + encodeCase "request_session_eval_v1.json" $ toRequestMessage (uuidFromString "cb682578-9d92-4499-9ebc-5c6aa73c5397") + $ S.OpEval + { S.batchSize = Nothing, + S.gremlin = "g.V(x)", + S.bindings = Just $ HM.fromList [("x", Number 1)], + S.language = Just "gremlin-groovy", + S.aliases = Nothing, + S.scriptEvaluationTimeout = Nothing, + S.session = uuidFromString "41d2e28a-20a4-4ab0-b379-d810dede3786", + S.manageTransaction = Nothing + } + encodeCase "request_session_eval_aliased_v1.json" $ toRequestMessage (uuidFromString "cb682578-9d92-4499-9ebc-5c6aa73c5397") + $ S.OpEval + { S.batchSize = Nothing, + S.gremlin = "social.V(x)", + S.bindings = Just $ HM.fromList [("x", Number 1)], + S.language = Just "gremlin-groovy", + S.aliases = Just $ HM.fromList [("g", "social")], + S.scriptEvaluationTimeout = Nothing, + S.session = uuidFromString "41d2e28a-20a4-4ab0-b379-d810dede3786", + S.manageTransaction = Nothing + } + encodeCase "request_session_close_v1.json" $ toRequestMessage (uuidFromString "cb682578-9d92-4499-9ebc-5c6aa73c5397") + $ S.OpClose + { S.batchSize = Nothing, + S.session = uuidFromString "41d2e28a-20a4-4ab0-b379-d810dede3786", + S.force = Nothing + } diff --git a/test/ServerTest.hs b/test/ServerTest.hs new file mode 100644 --- /dev/null +++ b/test/ServerTest.hs @@ -0,0 +1,14 @@ +{-# LANGUAGE OverloadedStrings, DuplicateRecordFields #-} +module Main (main,spec) where +import Test.Hspec + +import qualified ServerTest.Connection as Conn +import qualified ServerTest.Client as Client + +main :: IO () +main = hspec spec + +spec :: Spec +spec = do + Conn.spec + Client.spec diff --git a/test/ServerTest/Client.hs b/test/ServerTest/Client.hs new file mode 100644 --- /dev/null +++ b/test/ServerTest/Client.hs @@ -0,0 +1,167 @@ +{-# LANGUAGE OverloadedStrings #-} +module ServerTest.Client (main,spec) where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (withAsync, mapConcurrently) +import Control.Exception.Safe (bracket) +import Control.Monad (forM_) +import Data.Aeson (Value(Number)) +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HM +import Data.Maybe (isNothing, isJust, fromJust) +import Data.Text (Text) +import qualified Data.Vector as V +import qualified Network.WebSockets as WS +import Test.Hspec + +import Data.Greskell.Greskell (Greskell) +import qualified Data.Greskell.Greskell as G +import Data.Greskell.GMap (GMap, GMapEntry, unGMapEntry) + +import Network.Greskell.WebSocket.Client + ( Host, Port, Client, Options, + connectWith, close, submit, + defOptions, batchSize, connectionSettings, responseTimeout, + nextResult, slurpResults, + SubmitException(ResponseError, ParseError) + ) +import Network.Greskell.WebSocket.Connection + (RequestException(ResponseTimeout)) +import Network.Greskell.WebSocket.Request (RequestMessage(requestId)) +import qualified Network.Greskell.WebSocket.Response as Res + +import TestUtil.Env (withEnvForExtServer, withEnvForIntServer) +import TestUtil.MockServer + ( wsServer, receiveRequest, simpleRawResponse, waitForServer + ) + +main :: IO () +main = hspec spec + +spec :: Spec +spec = describe "Client" $ do + withEnvForExtServer $ do + client_basic_spec + withEnvForIntServer $ do + bad_server_spec + +withClient :: (Client -> IO a) -> (Host, Port) -> IO a +withClient act (host, port) = forClient' defOptions host port act + +forClient :: Host -> Port -> (Client -> IO a) -> IO a +forClient = forClient' defOptions + +forClient' :: Options -> Host -> Port -> (Client -> IO a) -> IO a +forClient' opt host port act = bracket (connectWith opt host port) close act + +client_basic_spec :: SpecWith (Host, Port) +client_basic_spec = do + specify "eval Int" $ withClient $ \client -> do + let g = 108 :: Greskell Int + rh <- submit client g Nothing + nextResult rh `shouldReturn` Just 108 + nextResult rh `shouldReturn` Nothing + nextResult rh `shouldReturn` Nothing + specify "eval Text" $ withClient $ \client -> do + let g = "hoge" :: Greskell Text + rh <- submit client g Nothing + nextResult rh `shouldReturn` Just "hoge" + nextResult rh `shouldReturn` Nothing + nextResult rh `shouldReturn` Nothing + specify "eval [Int]" $ withClient $ \client -> do + let g = G.list $ map fromInteger [1..20] :: Greskell [Int] + rh <- submit client g Nothing + forM_ [1..20] $ \n -> + nextResult rh `shouldReturn` Just n + nextResult rh `shouldReturn` Nothing + nextResult rh `shouldReturn` Nothing + specify "eval (GMap HashMap Int Text)" $ withClient $ \client -> do + let g :: Greskell (GMap HashMap Int Text) + g = G.unsafeGreskell "[100: 'hoge', 200: 'foo', 300: 'bar']" + rh <- submit client g Nothing + got <- fmap (fmap unGMapEntry) $ slurpResults rh + (V.toList got) `shouldMatchList` [(100, "hoge"), (200, "foo"), (300, "bar")] + specify "eval (Maybe Int)" $ withClient $ \client -> do + let g :: Greskell (Maybe Int) + g = G.unsafeGreskell "100" + rh <- submit client g Nothing + slurpResults rh `shouldReturn` V.fromList [Just 100] + specify "eval (bound Double)" $ withClient $ \client -> do + let gx = G.unsafeGreskell "x" + g = 92.125 + gx :: Greskell Double + b = HM.fromList [("x", Number 22.25)] + rh <- submit client g (Just b) + slurpResults rh `shouldReturn` V.fromList [114.375] + specify "multiple response messages" $ \(host, port) -> do + let opt = defOptions { batchSize = Just 3 + } + g = G.list $ map fromInteger [1..31] :: Greskell [Int] + forClient' opt host port $ \client -> do + rh <- submit client g Nothing + slurpResults rh `shouldReturn` V.fromList [1..31] + specify "ParseError exception" $ withClient $ \client -> do + let g :: Greskell Int + g = G.unsafeGreskell "\"some string\"" + expEx (ParseError _ _) = True + expEx _ = False + rh <- submit client g Nothing + nextResult rh `shouldThrow` expEx + nextResult rh `shouldThrow` expEx + specify "multiple threads reading on a single ResultHandle" $ \(host, port) -> do + let opt = defOptions { batchSize = Just 3 + } + g = G.list $ map fromInteger [1..50] :: Greskell [Int] + forClient' opt host port $ \client -> do + rh <- submit client g Nothing + mgot <- mapConcurrently (const $ nextResult rh) ([1..52] :: [Int]) + nextResult rh `shouldReturn` Nothing + (length $ filter isNothing mgot) `shouldBe` 2 + let got = map fromJust $ filter isJust mgot + got `shouldMatchList` [1..50] + specify "Gremlin script throwing exception" $ withClient $ \client -> do + let g :: Greskell Int + g = G.unsafeGreskell "throw new Exception(\"BOOM\")" + expEx (ResponseError res) = (Res.code $ Res.status res) == Res.ScriptEvaluationError + expEx _ = False + rh <- submit client g Nothing + nextResult rh `shouldThrow` expEx + nextResult rh `shouldThrow` expEx + + +bad_server_spec :: SpecWith Port +bad_server_spec = do + specify "error ResponseMessage" $ \port -> do + let server = wsServer port $ \wsconn -> do + req <- receiveRequest wsconn + WS.sendBinaryData wsconn $ simpleRawResponse (requestId req) 500 "null" + expEx (ResponseError res) = (Res.code $ Res.status res) == Res.ServerError + expEx _ = False + withAsync server $ \_ -> do + waitForServer + forClient "localhost" port $ \client -> do + let g = 100 :: Greskell Int + rh <- submit client g Nothing + nextResult rh `shouldThrow` expEx + nextResult rh `shouldThrow` expEx + nextResult rh `shouldThrow` expEx + specify "response timeout" $ \port -> do + let server = wsServer port $ \wsconn -> do + req <- receiveRequest wsconn + WS.sendBinaryData wsconn $ simpleRawResponse (requestId req) 206 "[99]" + threadDelay 10000000 + opt = defOptions { connectionSettings = sett + } + sett = (connectionSettings defOptions) + { responseTimeout = 1 + } + expEx ResponseTimeout = True + expEx _ = False + withAsync server $ \_ -> do + waitForServer + forClient' opt "localhost" port $ \client -> do + let g = G.list [99, 100, 101, 102] :: Greskell [Int] + rh <- submit client g Nothing + nextResult rh `shouldReturn` Just 99 + nextResult rh `shouldThrow` expEx + nextResult rh `shouldThrow` expEx + diff --git a/test/ServerTest/Connection.hs b/test/ServerTest/Connection.hs new file mode 100644 --- /dev/null +++ b/test/ServerTest/Connection.hs @@ -0,0 +1,373 @@ +{-# LANGUAGE OverloadedStrings, DuplicateRecordFields #-} +module ServerTest.Connection (main,spec) where + +import Control.Exception.Safe (bracket, Exception, withException, SomeException, throwString) +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (mapConcurrently, Async, withAsync, async, wait) +import Control.Concurrent.STM + ( newEmptyTMVarIO, putTMVar, takeTMVar, atomically, + TVar, newTVarIO, modifyTVar, readTVar, + TQueue, newTQueueIO, writeTQueue, flushTQueue + ) +import Control.Monad (when, forever, forM_, mapM) +import Data.Aeson (Value(Number), FromJSON(..), ToJSON(toJSON), Object) +import qualified Data.Aeson as Aeson +import qualified Data.Aeson.Types as Aeson (parseEither) +import qualified Data.ByteString.Lazy as BSL +import Data.Maybe (isNothing, isJust, fromJust) +import Data.Monoid ((<>)) +import qualified Data.HashMap.Strict as HM +import Data.Greskell.GraphSON (GraphSON, gsonValue) +import Data.Text (Text, pack) +import qualified Data.Vector as V +import Data.UUID (UUID) +import qualified Data.UUID as UUID +import Data.UUID.V4 (nextRandom) +import qualified Network.WebSockets as WS +import System.IO (stderr, hPutStrLn) +import Test.Hspec + +import Network.Greskell.WebSocket.Connection + ( Host, Port, Connection, ResponseHandle, + close, connect, sendRequest', sendRequest, slurpResponses, + nextResponse, + RequestException(..), GeneralException(..), + Settings(onGeneralException, responseTimeout, concurrency, requestQueueSize), + defJSONSettings + ) +import Network.Greskell.WebSocket.Request + ( RequestMessage(requestId), toRequestMessage, makeRequestMessage + ) +import Network.Greskell.WebSocket.Request.Standard (OpEval(..)) +import Network.Greskell.WebSocket.Response + ( ResponseMessage(requestId, status, result), ResponseStatus(code), ResponseCode(..), + ResponseResult(resultData) + ) +import qualified Network.Greskell.WebSocket.Response as Response + +import qualified TestUtil.TCounter as TCounter +import TestUtil.Env (withEnvForExtServer, withEnvForIntServer) +import TestUtil.MockServer (wsServer, receiveRequest, simpleRawResponse, waitForServer) + + +main :: IO () +main = hspec spec + +spec :: Spec +spec = describe "Connection" $ do + no_external_server_spec + withEnvForExtServer $ do + conn_basic_spec + conn_error_spec + conn_close_spec + withEnvForIntServer $ do + conn_bad_server_spec + +ourSettings :: Settings Value +ourSettings = defJSONSettings + { onGeneralException = \e -> error (show e) + -- basically we don't expect any GeneralException. + } + +withConn :: (Connection Value -> IO a) -> (Host, Port) -> IO a +withConn act (host, port) = forConn host port act + +forConn :: Host -> Port -> (Connection Value -> IO a) -> IO a +forConn = forConn' defJSONSettings + +forConn' :: Settings Value -> Host -> Port -> (Connection Value -> IO a) -> IO a +forConn' settings host port act = bracket makeConn close act + where + makeConn = connect settings host port + +parseValue :: FromJSON a => Value -> Either String a +parseValue v = Aeson.parseEither parseJSON v + +opEval :: Text -> OpEval +opEval g = OpEval { batchSize = Nothing, + gremlin = g, + bindings = Nothing, + language = Nothing, + aliases = Nothing, + scriptEvaluationTimeout = Nothing + } + +---- 'resultData' should be parsed with GraphSON wrappers to support v1, v2 and v3. +slurpParseEval :: FromJSON a => ResponseHandle Value -> IO [ResponseMessage (Either String (GraphSON [GraphSON a]))] +slurpParseEval rh = (fmap . fmap . fmap) parseValue $ fmap V.toList $ slurpResponses rh + +responseValues :: ResponseMessage (Either String (GraphSON [GraphSON a])) -> Either String [a] +responseValues = fmap (map gsonValue . gsonValue) . resultData . result + +slurpEvalValues :: FromJSON a => ResponseHandle Value -> IO [Either String [a]] +slurpEvalValues rh = (fmap . map) responseValues $ slurpParseEval rh + + +opSleep :: Int -> OpEval +opSleep time_ms = opSleep' time_ms time_ms + +opSleep' :: Int -> Int -> OpEval +opSleep' time_ms val = opEval ("sleep " <> tshow time_ms <> "; " <> tshow val) + where + tshow = pack . show + +inspectException :: IO a -> IO a +inspectException act = act `withException` printE + where + printE :: SomeException -> IO () + printE e = hPutStrLn stderr ("[DEBUG] exception thrown: " ++ show e) + +no_external_server_spec :: Spec +no_external_server_spec = describe "connect" $ do + it "should throw exception on failure" $ do + let act = connect ourSettings "this.should.not.be.real.server.com" 8000 + expectSomeEx :: SomeException -> Bool + expectSomeEx _ = True + inspectException act `shouldThrow` expectSomeEx + + + +conn_basic_spec :: SpecWith (Host, Port) +conn_basic_spec = do + specify "basic transaction" $ withConn $ \conn -> do + rid <- nextRandom + let op = opEval "123" + exp_val :: [Int] + exp_val = [123] + res <- sendRequest' conn $ toRequestMessage rid op + got <- slurpParseEval res + map (Response.requestId) got `shouldBe` [rid] + map (code . status) got `shouldBe` [Success] + map responseValues got `shouldBe` [Right exp_val] + specify "continuous response with bindings" $ withConn $ \conn -> do + rid <- nextRandom + let op = (opEval "x") { batchSize = Just 2, + bindings = Just $ HM.fromList [("x", toJSON ([1 .. 10] :: [Int]))] + } + exp_vals :: [Either String [Int]] + exp_vals = map Right [[1,2], [3,4], [5,6], [7,8], [9,10]] + got <- slurpParseEval =<< (sendRequest' conn $ toRequestMessage rid op) + map (Response.requestId) got `shouldBe` replicate 5 rid + map (code . status) got `shouldBe` ((replicate 4 PartialContent) ++ [Success]) + map responseValues got `shouldBe` exp_vals + specify "concurrent requests" $ withConn $ \conn -> do + handles <- mapM (sendRequest conn) $ map opSleep $ map (* 100) $ reverse [1..5] + got <- mapM slurpEvalValues handles :: IO [[Either String [Int]]] + got `shouldBe` [ [Right [500]], + [Right [400]], + [Right [300]], + [Right [200]], + [Right [100]] + ] + specify "make requests from multiple threads" $ withConn $ \conn -> do + let reqAndRes val = slurpEvalValues =<< (sendRequest conn $ opSleep' 200 val) + got <- mapConcurrently reqAndRes [1 .. 5] + :: IO [[Either String [Int]]] + got `shouldBe` [ [Right [1]], + [Right [2]], + [Right [3]], + [Right [4]], + [Right [5]] + ] + specify "requestId should be cleared from the request pool once completed" $ withConn $ \conn -> do + req <- makeRequestMessage $ opSleep 30 + let evalReq = slurpEvalValues =<< sendRequest' conn req :: IO [Either String [Int]] + evalReq `shouldReturn` [Right [30]] + evalReq `shouldReturn` [Right [30]] + specify "client concurrency should be bounded by 'concurrency' + 'requestQueueSize'" $ \(host, port) -> do + let input_concurrency = 2 + input_qreq_size = 1 + settings = defJSONSettings { concurrency = input_concurrency, + requestQueueSize = input_qreq_size + } + exp_concurrency_max = input_concurrency + input_qreq_size + forConn' settings host port $ \conn -> do + tcounter <- TCounter.new + let makeReq v = TCounter.count tcounter + (sendRequest conn $ opSleep' 500 v) + (\rh -> slurpEvalValues rh :: IO [Either String [Int]]) + got <- mapConcurrently makeReq [1..10] + got `shouldBe` map (\v -> [Right [v]]) [1..10] + got_hist <- TCounter.history tcounter + length got_hist `shouldBe` (2 * 10) + forM_ got_hist $ \conc -> conc `shouldSatisfy` (<= exp_concurrency_max) + specify "multiple threads reading a single ResponseHandle" $ withConn $ \conn -> do + let op = (opEval "[1,2,3,4,5,6,7,8,9,10]") + { batchSize = Just 1 + } + rh <- sendRequest conn op + mgot <- mapConcurrently (const $ nextResponse rh) ([1..15] :: [Int]) + (length $ filter isNothing mgot) `shouldBe` 5 + let got = map fromJust $ filter isJust mgot + got_data = map (responseValues . fmap parseValue) got :: [Either String [Int]] + map (code . status) got `shouldMatchList` (Success : replicate 9 PartialContent) + got_data `shouldMatchList` map (\v -> Right [v]) [1..10] + +conn_error_spec :: SpecWith (Host, Port) +conn_error_spec = do + specify "duplicate requests" $ withConn $ \conn -> do + req <- makeRequestMessage $ opSleep 300 + let expEx (DuplicateRequestId got_rid) = got_rid == requestId (req :: RequestMessage) + expEx _ = False + ok_rh <- sendRequest' conn req + ng_rh <- sendRequest' conn req + ok_res <- slurpEvalValues ok_rh :: IO [Either String [Int]] + ok_res `shouldBe` [Right [300]] + nextResponse ng_rh `shouldThrow` expEx + specify "request timeout" $ \(host, port) -> do + let settings = ourSettings { responseTimeout = 1 } + expEx ResponseTimeout = True + expEx _ = False + forConn' settings host port $ \conn -> do + rh <- sendRequest conn $ opSleep 2000 + (nextResponse rh) `shouldThrow` expEx + +conn_close_spec :: SpecWith (Host, Port) +conn_close_spec = describe "close" $ do + it "should wait for a pending request to finish" $ withConn $ \conn -> do + rh <- sendRequest conn $ opSleep 100 + close conn + got <- slurpEvalValues rh :: IO [Either String [Int]] + got `shouldBe` [Right [100]] + it "should block for all pending requests" $ withConn $ \conn -> do + tc <- TCounter.new + let makeReq :: Int -> IO [Either String [Int]] + makeReq t = do + ret <- TCounter.count tc (sendRequest conn $ opSleep t) slurpEvalValues + close conn -- more than one close don't hurt anybody (after all requests are sent.) + return ret + req_threads <- mapM (async . makeReq) [200, 400, 600] + TCounter.waitFor tc (== 3) + close conn + TCounter.now tc `shouldReturn` 0 + got_hist <- TCounter.history tc + length got_hist `shouldBe` 6 + got <- mapM wait req_threads + got `shouldBe` map (\v -> [Right [v]]) [200, 400, 600] + it "should make nextResponse throw AlreadyClosed exception. nextResponse should throw it every time it's called" $ withConn $ \conn -> do + let expEx AlreadyClosed = True + expEx _ = False + close conn + rh <- sendRequest conn $ opEval "999" + nextResponse rh `shouldThrow` expEx + nextResponse rh `shouldThrow` expEx + nextResponse rh `shouldThrow` expEx + + +succUUID :: UUID -> UUID +succUUID orig = UUID.fromWords a b c d' + where + (a, b, c ,d) = UUID.toWords orig + d' = succ d + +conn_bad_server_spec :: SpecWith Port +conn_bad_server_spec = do + describe "ResponseHandle" $ describe "nextResponse" $ do + it "should throw exception when the server closes the connection while there is a pending request" $ \port -> do + let server = wsServer port $ \wsconn -> do + _ <- WS.receiveDataMessage wsconn + throwString ( "Server connection abort. Seeing this message in a test console is OK." + ++ " It's because of WS.runServer internals." + ) + exp_ex ServerClosed = True + exp_ex _ = False + withAsync server $ \_ -> do + waitForServer + forConn "localhost" port $ \conn -> do + rh <- sendRequest conn $ opEval "100" + (inspectException $ nextResponse rh) `shouldThrow` exp_ex + it "should throw exception when the server sends Close request while there is a pending request" $ \port -> do + let server = wsServer port $ \wsconn -> do + _ <- WS.receiveDataMessage wsconn + WS.sendClose wsconn ("" :: Text) + (WS.ControlMessage (WS.Close wscode _)) <- WS.receive wsconn + when (wscode /= 1000) $ do + throwString ( "Fatal: expects WebSocket status code 1000, but got " + ++ show wscode + ) + exp_ex ServerClosed = True + exp_ex _ = False + withAsync server $ \_ -> do + waitForServer + forConn "localhost" port $ \conn -> do + rh <- sendRequest conn $ opEval "100" + (inspectException $ nextResponse rh) `shouldThrow` exp_ex + it "should be ok if the server actively closes the connection" $ \port -> do + let server = wsServer port $ \wsconn -> do + req <- receiveRequest wsconn + WS.sendBinaryData wsconn $ simpleRawResponse (requestId (req :: RequestMessage)) 200 "[99]" + withAsync server $ \_ -> do + waitForServer + forConn "localhost" port $ \conn -> do + rh <- sendRequest conn $ opEval "99" + got <- slurpEvalValues rh :: IO [Either String [Int]] + got `shouldBe` [Right [99]] + it "should throw AlreadyClosed exception after the server actively closes the connection" $ \port -> do + let server = wsServer port $ \wsconn -> do + threadDelay 10000 + WS.sendClose wsconn ("" :: Text) + expEx AlreadyClosed = True + expEx _ = False + withAsync server $ \_ -> do + waitForServer + forConn "localhost" port $ \conn -> do + threadDelay 40000 + rh <- sendRequest conn $ opEval "256" + nextResponse rh `shouldThrow` expEx + describe "Settings" $ describe "onGeneralException" $ do + it "should be called on unexpected requestId" $ \port -> do + report_gex <- newEmptyTMVarIO + let server = wsServer port $ \wsconn -> do + req <- receiveRequest wsconn + let res_id = succUUID $ requestId (req :: RequestMessage) -- deliberately send wrong requestId. + res = simpleRawResponse res_id 200 "[333]" + WS.sendBinaryData wsconn res + reportEx ex = atomically $ putTMVar report_gex ex + settings = defJSONSettings { onGeneralException = reportEx } + withAsync server $ \_ -> do + waitForServer + forConn' settings "localhost" port $ \conn -> do + req <- makeRequestMessage $ opEval "333" + let exp_res_id = succUUID $ requestId (req :: RequestMessage) + _ <- sendRequest' conn req + got <- atomically $ takeTMVar report_gex + got `shouldBe` UnexpectedRequestId exp_res_id + it "should be called on failure to parse response" $ \port -> do + report_gex <- newEmptyTMVarIO + let server = wsServer port $ \wsconn -> do + _ <- receiveRequest wsconn + WS.sendBinaryData wsconn ("hoge hoge hoge" :: Text) + settings = defJSONSettings { onGeneralException = \e -> atomically $ putTMVar report_gex e } + withAsync server $ \_ -> do + waitForServer + forConn' settings "localhost" port $ \conn -> do + let expEx (ResponseParseFailure _) = True + expEx _ = False + _ <- sendRequest conn $ opEval "100" + got <- atomically $ takeTMVar report_gex + got `shouldSatisfy` expEx + it "should throw ResponseTimeout exception when the server endlessly sends responses" $ \port -> do + let server = wsServer port $ \wsconn -> do + req <- receiveRequest wsconn + let res_id = requestId (req :: RequestMessage) + sendContinuousRes wsconn res_id 20 + sendContinuousRes :: WS.Connection -> UUID -> Int -> IO () + sendContinuousRes wsconn res_id n = do + let finish = n == 0 + threadDelay 200000 + WS.sendBinaryData wsconn $ simpleRawResponse res_id (if finish then 200 else 206) "[1]" + if finish + then return () + else sendContinuousRes wsconn res_id (n - 1) + settings = ourSettings { responseTimeout = 1 } + expEx ResponseTimeout = True + expEx _ = False + withAsync server $ \_ -> do + waitForServer + forConn' settings "localhost" port $ \conn -> do + rh <- sendRequest conn $ opEval "99" + forM_ ([1..3] :: [Int]) $ \_ -> do + mgot <- (fmap . fmap) (responseValues . fmap parseValue) $ nextResponse rh :: IO (Maybe (Either String [Int])) + mgot `shouldBe` (Just $ Right [1]) + slurpResponses rh `shouldThrow` expEx + diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-} diff --git a/test/TestUtil/Env.hs b/test/TestUtil/Env.hs new file mode 100644 --- /dev/null +++ b/test/TestUtil/Env.hs @@ -0,0 +1,27 @@ +module TestUtil.Env + ( requireEnv, + withEnvForExtServer, + withEnvForIntServer + ) where + +import System.Environment (lookupEnv) +import Test.Hspec + +import Network.Greskell.WebSocket.Connection (Host, Port) + +requireEnv :: String -> IO String +requireEnv env_key = maybe bail return =<< lookupEnv env_key + where + bail = expectationFailure msg >> return "" + where + msg = "Set environment variable "++ env_key ++ " for Server test. " + +withEnvForExtServer :: SpecWith (Host, Port) -> Spec +withEnvForExtServer = before $ do + hostname <- requireEnv "GRESKELL_TEST_HOST" + port <- fmap read $ requireEnv "GRESKELL_TEST_PORT" + return (hostname, port) + +withEnvForIntServer :: SpecWith Port -> Spec +withEnvForIntServer = before $ fmap read $ requireEnv "GRESKELL_TEST_INTERNAL_PORT" + diff --git a/test/TestUtil/MockServer.hs b/test/TestUtil/MockServer.hs new file mode 100644 --- /dev/null +++ b/test/TestUtil/MockServer.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE OverloadedStrings #-} +module TestUtil.MockServer + ( wsServer, + parseRequest, + receiveRequest, + simpleRawResponse, + waitForServer + ) where + +import Control.Concurrent (threadDelay) +import Control.Exception.Safe (throwString) +import qualified Data.Aeson as Aeson +import qualified Data.ByteString.Lazy as BSL +import Data.Monoid ((<>)) +import Data.Text (Text, pack) +import Data.UUID (UUID) +import qualified Data.UUID as UUID +import qualified Network.WebSockets as WS + +import Network.Greskell.WebSocket.Codec (decodeBinary) +import Network.Greskell.WebSocket.Request (RequestMessage) + + +wsServer :: Int -- ^ port number + -> (WS.Connection -> IO ()) + -> IO () +wsServer p act = WS.runServer "localhost" p $ \pending_conn -> + act =<< WS.acceptRequest pending_conn + +parseRequest :: BSL.ByteString -> Either String RequestMessage +parseRequest raw_msg = do + (_, payload) <- decodeBinary raw_msg + Aeson.eitherDecode payload + +receiveRequest :: WS.Connection -> IO RequestMessage +receiveRequest wsconn = do + raw_msg <- WS.receiveData wsconn + case parseRequest raw_msg of + Left e -> throwString e + Right r -> return r + +simpleRawResponse :: UUID -> Int -> Text -> Text +simpleRawResponse request_id status_code data_content = + "{\"requestId\":\"" <> UUID.toText request_id <> "\"," + <> "\"status\":{\"code\":" <> (pack $ show status_code) <> ",\"message\":\"\",\"attributes\":{}}," + <> "\"result\":{\"data\":" <> data_content <> ",\"meta\":{}}}" + +waitForServer :: IO () +waitForServer = threadDelay 100000 diff --git a/test/TestUtil/TCounter.hs b/test/TestUtil/TCounter.hs new file mode 100644 --- /dev/null +++ b/test/TestUtil/TCounter.hs @@ -0,0 +1,59 @@ +-- | +-- Module: TestUtil.TCounter +-- Description: +-- Maintainer: Toshio Ito +-- +-- +module TestUtil.TCounter + ( TCounter, + new, + modify, + now, + waitFor, + history, + count + ) where + +import Control.Applicative ((<$>), (<*>)) +import Control.Concurrent.STM + ( TVar, newTVarIO, modifyTVar, readTVar, + STM, atomically, retry + ) + +-- | Transaction counter. +data TCounter = TCounter { tcCurrent :: TVar Int, + tcHistory :: TVar [Int] + } + +new :: IO TCounter +new = TCounter <$> newTVarIO 0 <*> newTVarIO [] + +modify :: TCounter -> (Int -> Int) -> IO () +modify tc f = atomically $ do + modifyTVar (tcCurrent tc) f + conc <- readTVar (tcCurrent tc) + modifyTVar (tcHistory tc) (conc :) + +now :: TCounter -> IO Int +now tc = atomically $ nowSTM tc + +nowSTM :: TCounter -> STM Int +nowSTM tc = readTVar $ tcCurrent tc + +waitFor :: TCounter -> (Int -> Bool) -> IO () +waitFor tc p = atomically $ do + cur <- nowSTM tc + if p cur + then return () + else retry + +history :: TCounter -> IO [Int] +history tc = reverse <$> (atomically $ readTVar $ tcHistory tc) + +count :: TCounter -> IO a -> (a -> IO b) -> IO b +count tc start_act finish_act = do + tx <- start_act + modify tc (+ 1) + ret <- finish_act tx + modify tc (subtract 1) + return ret diff --git a/test/samples/request_auth_v1.json b/test/samples/request_auth_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/request_auth_v1.json @@ -0,0 +1,9 @@ +{ + "requestId" : "cb682578-9d92-4499-9ebc-5c6aa73c5397", + "op" : "authentication", + "processor" : "", + "args" : { + "saslMechanism" : "PLAIN", + "sasl" : "AHN0ZXBocGhlbgBwYXNzd29yZA==" + } +} diff --git a/test/samples/request_session_close_v1.json b/test/samples/request_session_close_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/request_session_close_v1.json @@ -0,0 +1,8 @@ +{ + "requestId" : "cb682578-9d92-4499-9ebc-5c6aa73c5397", + "op" : "close", + "processor" : "session", + "args" : { + "session" : "41d2e28a-20a4-4ab0-b379-d810dede3786" + } +} diff --git a/test/samples/request_session_eval_aliased_v1.json b/test/samples/request_session_eval_aliased_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/request_session_eval_aliased_v1.json @@ -0,0 +1,16 @@ +{ + "requestId" : "cb682578-9d92-4499-9ebc-5c6aa73c5397", + "op" : "eval", + "processor" : "session", + "args" : { + "gremlin" : "social.V(x)", + "language" : "gremlin-groovy", + "aliases" : { + "g" : "social" + }, + "session" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "bindings" : { + "x" : 1 + } + } +} diff --git a/test/samples/request_session_eval_v1.json b/test/samples/request_session_eval_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/request_session_eval_v1.json @@ -0,0 +1,13 @@ +{ + "requestId" : "cb682578-9d92-4499-9ebc-5c6aa73c5397", + "op" : "eval", + "processor" : "session", + "args" : { + "gremlin" : "g.V(x)", + "language" : "gremlin-groovy", + "session" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "bindings" : { + "x" : 1 + } + } +} diff --git a/test/samples/request_sessionless_eval_aliased_v1.json b/test/samples/request_sessionless_eval_aliased_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/request_sessionless_eval_aliased_v1.json @@ -0,0 +1,15 @@ +{ + "requestId" : "cb682578-9d92-4499-9ebc-5c6aa73c5397", + "op" : "eval", + "processor" : "", + "args" : { + "gremlin" : "social.V(x)", + "language" : "gremlin-groovy", + "aliases" : { + "g" : "social" + }, + "bindings" : { + "x" : 1 + } + } +} diff --git a/test/samples/request_sessionless_eval_v1.json b/test/samples/request_sessionless_eval_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/request_sessionless_eval_v1.json @@ -0,0 +1,12 @@ +{ + "requestId" : "cb682578-9d92-4499-9ebc-5c6aa73c5397", + "op" : "eval", + "processor" : "", + "args" : { + "gremlin" : "g.V(x)", + "language" : "gremlin-groovy", + "bindings" : { + "x" : 1 + } + } +} diff --git a/test/samples/response_auth_v1.json b/test/samples/response_auth_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/response_auth_v1.json @@ -0,0 +1,12 @@ +{ + "requestId" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "status" : { + "message" : "", + "code" : 407, + "attributes" : { } + }, + "result" : { + "data" : null, + "meta" : { } + } +} diff --git a/test/samples/response_auth_v2.json b/test/samples/response_auth_v2.json new file mode 100644 --- /dev/null +++ b/test/samples/response_auth_v2.json @@ -0,0 +1,12 @@ +{ + "requestId" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "status" : { + "message" : "", + "code" : 407, + "attributes" : { } + }, + "result" : { + "data" : null, + "meta" : { } + } +} diff --git a/test/samples/response_auth_v3.json b/test/samples/response_auth_v3.json new file mode 100644 --- /dev/null +++ b/test/samples/response_auth_v3.json @@ -0,0 +1,18 @@ +{ + "requestId" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "status" : { + "message" : "", + "code" : 407, + "attributes" : { + "@type" : "g:Map", + "@value" : [ ] + } + }, + "result" : { + "data" : null, + "meta" : { + "@type" : "g:Map", + "@value" : [ ] + } + } +} diff --git a/test/samples/response_standard_v1.json b/test/samples/response_standard_v1.json new file mode 100644 --- /dev/null +++ b/test/samples/response_standard_v1.json @@ -0,0 +1,16 @@ +{ + "requestId" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "status" : { + "message" : "", + "code" : 200, + "attributes" : { } + }, + "result" : { + "data" : [ { + "id" : 1, + "label" : "person", + "type" : "vertex" + } ], + "meta" : { } + } +} diff --git a/test/samples/response_standard_v2.json b/test/samples/response_standard_v2.json new file mode 100644 --- /dev/null +++ b/test/samples/response_standard_v2.json @@ -0,0 +1,21 @@ +{ + "requestId" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "status" : { + "message" : "", + "code" : 200, + "attributes" : { } + }, + "result" : { + "data" : [ { + "@type" : "g:Vertex", + "@value" : { + "id" : { + "@type" : "g:Int32", + "@value" : 1 + }, + "label" : "person" + } + } ], + "meta" : { } + } +} diff --git a/test/samples/response_standard_v3.json b/test/samples/response_standard_v3.json new file mode 100644 --- /dev/null +++ b/test/samples/response_standard_v3.json @@ -0,0 +1,30 @@ +{ + "requestId" : "41d2e28a-20a4-4ab0-b379-d810dede3786", + "status" : { + "message" : "", + "code" : 200, + "attributes" : { + "@type" : "g:Map", + "@value" : [ ] + } + }, + "result" : { + "data" : { + "@type" : "g:List", + "@value" : [ { + "@type" : "g:Vertex", + "@value" : { + "id" : { + "@type" : "g:Int32", + "@value" : 1 + }, + "label" : "person" + } + } ] + }, + "meta" : { + "@type" : "g:Map", + "@value" : [ ] + } + } +}