@@ -12,7 +12,7 @@ import Codec.Winery
1212import Codec.Winery.Internal (varInt, decodeVarInt, getBytes)
1313import Codec.Winery.Class (mkExtractor)
1414import GHC.Generics
15−import Control.Concurrent.MVar (MVar)
15+import Control.Concurrent.MVar (MVar, newMVar, withMVar)
1616import Control.Exception
1717import Data.Function ((&))
1818import Data.Word
@@ -50,6 +50,20 @@ traceBytes msg bs = traceShowM (msg, BS.length bs, bs)
5050traceBytes _ _ = pure ()
5151#endif
5252
53+data Locking a = Locking (MVar ()) a
54+
55+newLock :: a -> IO (Locking a)
56+newLock x = do
57+ lock <- newMVar ()
58+ pure (Locking lock x)
59+
60+withLock :: Locking a -> (a -> IO b) -> IO b
61+withLock (Locking mvar v) m =
62+ withMVar mvar $ \_ -> m v
63+
64+lockless :: Locking a -> a
65+lockless (Locking _ a) = a
66+
5367type Timeout = Int
5468
5569data Message a = Response UUID a
@@ -131,23 +145,24 @@ serve connhandler userMsgHandler hostaddr port mSockLock = do
131145 Right dec -> dec
132146 handleSock sock = do
133147 mResp <- connhandler
148+ lockingSocket <- newLock sock
134149 case mResp of
135150 Nothing -> pure ()
136151 Just resp -> do
137152 requestID <- UUID <$> UUIDBase.nextRandom
138− sendMessage (AsyncRequest requestID resp) sock
139− drainSocketMessages sock (decoder userMsgHandler) (serverMessageHandler sock userMsgHandler)
153+ sendMessage (AsyncRequest requestID resp) lockingSocket
154+ drainSocketMessages sock (decoder userMsgHandler) (serverMessageHandler lockingSocket userMsgHandler)
140155 serially (S.unfold (SA.acceptOnAddrWith [(ReuseAddr,1)] mSockLock) (hostaddr, port)) & parallely . S.mapM (handleWithM handleSock) & S.drain
141156 pure True
142157
143158serverMessageHandler :: forall req resp. (Show req,
144159 Serialise req,
145160 Serialise resp)
146− => Socket
161+ => Locking Socket
147162 -> NewMessageHandler req resp
148163 -> Message req
149164 -> IO ()
150−serverMessageHandler sock requestMessageHandler msg = do
165+serverMessageHandler sockLock requestMessageHandler msg = do
151166 --putStrLn $ "GOT MSG " ++ show msg
152167 let --runTimeout :: IO (HandlerResponse resp) -> IO (Maybe (HandlerResponse resp))
153168 runTimeout mTimeout m = case mTimeout of
@@ -159,18 +174,18 @@ serverMessageHandler sock requestMessageHandler msg = do
159174 let normalResponder = do
160175 resp <- runTimeout mTimeout (requestMessageHandler val)
161176 case resp of
162− Just (HandlerResponse responseVal) -> sendMessage (Response requestID responseVal) sock
177+ Just (HandlerResponse responseVal) -> sendMessage (Response requestID responseVal) sockLock
163178 Just NoResponse -> error "attempt to return non-response to expected response message"
164− Just (HandlerException exc) -> sendMessage (ExceptionResponse @(Message resp) requestID exc) sock
179+ Just (HandlerException exc) -> sendMessage (ExceptionResponse @(Message resp) requestID exc) sockLock
165180 Nothing ->
166− sendMessage (TimedOutResponse @(Message resp) requestID) sock
181+ sendMessage (TimedOutResponse @(Message resp) requestID) sockLock
167182 excHandler :: SomeException -> IO ()
168183 excHandler e = do
169184 --send exception to client
170− sendMessage (ExceptionResponse @(Message resp) requestID (show e)) sock
185+ sendMessage (ExceptionResponse @(Message resp) requestID (show e)) sockLock
171186 throwIO e
172187 catch normalResponder excHandler
173− AsyncRequest _ val -> do
188+ AsyncRequest _ val ->
174189 void $ requestMessageHandler val
175190 --no response necessary
176191 ExceptionResponse{} -> putStrLn "client sent exception response"
@@ -192,36 +207,15 @@ drainSocketMessages sock decoder msgHandler = do
192207 S.drain $ serially $ S.parseMany messageBoundaryP sockStream & S.mapM handler
193208
194209--send length-tagged bytestring, perhaps should be in network byte order?
195−sendMessage :: Serialise a => Message a -> Socket -> IO ()
196−sendMessage msg socket' = do
210+sendMessage :: Serialise a => Message a -> Locking Socket -> IO ()
211+sendMessage msg sockLock = do
197212 let
198213 msgbytes = serialiseOnly msg
199214 fullbytes = lenbytes <> msgbytes
200215 len = BS.length msgbytes
201216 lenbytes = BO.bytestring32 (fromIntegral len)
202−
203− byteCount <- Socket.send socket' fullbytes
204− when (byteCount /= BS.length fullbytes) (error "bytes sent mismatch")
205− traceBytes "sent bytes" fullbytes
206−
207−
208−
209−--from streamly-bytestring
210−-- | Convert a 'ByteString' to an array of 'Word8'. This function unwraps the
211−-- 'ByteString' and wraps it with 'Array' constructors and hence the operation
212−-- is performed in constant time.
213−{-
214−{-# INLINE toArrayS #-}
215−toArrayS :: BS.ByteString -> SA.Array Word8
216−--slow path
217−--toArrayS = SA.fromList . BS.unpack
218−toArrayS (BSI.PS fp off len) = SA.Array nfp endPtr
219− where
220− nfp = fp `plusForeignPtr` off
221− endPtr = unsafeForeignPtrToPtr nfp `plusPtr` len `plusPtr` 1 -- ? why +1?
222−
223−toArraySlow :: BS.ByteString -> SA.Array Word8
224−toArraySlow bs = arr
225− where
226− arr = SA.fromList (BS.unpack bs)
227−-}
217+
218+ --Socket.sendAll syscalls send() on a loop until all the bytes are sent, so we need socket locking here to account for serialized messages of size > PIPE_BUF
219+ withLock sockLock $ \socket' ->
220+ Socket.sendAll socket' fullbytes
221+ traceBytes "sendMessage" fullbytes