Bryan O'Sullivan avatar Bryan O'Sullivan committed 62fcecd

Fix some nasty buffering bugs, and improve performance

* Bug 1: if a previous receive left buffered data behind, recvExactly
tried to use it, but was losing it instead. Oops.

* Bug 2: recvGetN was doing completely nonsensical buffer management.

Performance improvement: recvExactly uses a too-large buffer, just
in case there might be some, and saves any leftover data.

Comments (0)

Files changed (1)

src/Network/Riak/Connection/Internal.hs

 
 import Control.Concurrent
 import Control.Exception (Exception, IOException, throw)
-import Control.Monad (forM_, replicateM, replicateM_, unless)
+import Control.Monad (forM_, replicateM, replicateM_)
 import Data.Binary.Put (Put, putWord32be, runPut)
-import Data.IORef (modifyIORef, newIORef, readIORef, writeIORef)
+import Data.IORef (newIORef, readIORef, writeIORef)
 import Data.Int (Int64)
 import Network.Riak.Connection.NoPush (setNoPush)
 import Network.Riak.Debug as Debug
   sClose connSock
   writeIORef connBuffer L.empty
 
+recvBufferSize :: Integral a => a
+recvBufferSize = 16384
+{-# INLINE recvBufferSize #-}
+
 recvWith :: (L.ByteString -> IO L.ByteString) -> Connection -> Int64
          -> IO L.ByteString
 recvWith onError Connection{..} n0
       len = L.length h
   if len == n0
     then writeIORef connBuffer t >> return h
-    else if len == 0
-         then go [] n0
-         else go (reverse (L.toChunks t)) (n0-len)
+    else go (reverse (L.toChunks h)) (n0-len)
   where
     maxInt = fromIntegral (maxBound :: Int)
+    go (s:acc) n
+      | n < 0 = do
+        let (h,t) = B.splitAt (B.length s - fromIntegral n) s
+        writeIORef connBuffer $! L.fromChunks [t]
+        return $ L.fromChunks (reverse (h:acc))
     go acc n
-        | n <= 0 = return (L.fromChunks (reverse acc))
-        | otherwise = do
-      let n' = min n maxInt
-      bs <- B.recv connSock (fromIntegral n')
-      let len = B.length bs
-      if len == 0
-        then onError (L.fromChunks (reverse acc))
-        else go (bs:acc) (n' - fromIntegral len)
+      | n == 0 = do
+        writeIORef connBuffer L.empty
+        return $ L.fromChunks (reverse acc)
+      | otherwise = do
+        let n' = max recvBufferSize $ min n maxInt
+        bs <- B.recv connSock (fromIntegral n')
+        let len = B.length bs
+        if len == 0
+          then onError (L.fromChunks (reverse acc))
+          else go (bs:acc) (n - fromIntegral len)
 
 recvExactly :: Connection -> Int64 -> IO L.ByteString
 recvExactly = recvWith $ \_ ->
 recvGet :: Connection -> Get a -> IO a
 recvGet Connection{..} get = do
   let refill = do
-        bs <- L.recv connSock 16384
+        bs <- L.recv connSock recvBufferSize
         if L.null bs
           then shutdown connSock ShutdownReceive >> return Nothing
           else return (Just bs)
 recvGetN :: Connection -> Int64 -> Get a -> IO a
 recvGetN conn n get = do
   bs <- recvExactly conn n
-  let finish bs' r = do
-        unless (L.null bs') $ modifyIORef (connBuffer conn) (`L.append` bs')
-        return r
   case runGet get bs of
-    Finished bs' _ r -> finish bs' r
+    Finished _ _ r -> return r
     Partial k    -> case k Nothing of
-                      Finished bs' _ r -> finish bs' r
+                      Finished _ _ r -> return r
                       Failed _ err -> moduleError "recvGetN" err
                       Partial _    -> moduleError "recvGetN"
                                       "parser wants more input!?"
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.