improve Local monad

This commit is contained in:
Joey Hess 2016-12-02 13:47:42 -04:00
parent 15dc63d47f
commit 7b7afbbedc
No known key found for this signature in database
GPG key ID: C910D9222512E3C7

View file

@ -127,7 +127,12 @@ data NetF c
= SendMessage Message c = SendMessage Message c
| ReceiveMessage (Message -> c) | ReceiveMessage (Message -> c)
| SendBytes Len L.ByteString c | SendBytes Len L.ByteString c
-- ^ Sends exactly Len bytes of data. (Any more or less will
-- confuse the receiver.)
| ReceiveBytes Len (L.ByteString -> c) | ReceiveBytes Len (L.ByteString -> c)
-- ^ Lazily reads bytes from peer. Stops once Len are read,
-- or if connection is lost, and in either case returns the bytes
-- that were read. This allows resuming interrupted transfers.
| CheckAuthToken UUID AuthToken (Bool -> c) | CheckAuthToken UUID AuthToken (Bool -> c)
| RelayService Service c | RelayService Service c
-- ^ Runs a service, relays its output to the peer, and data -- ^ Runs a service, relays its output to the peer, and data
@ -144,24 +149,28 @@ type Net = Free NetF
newtype RelayHandle = RelayHandle Handle newtype RelayHandle = RelayHandle Handle
data LocalF c data LocalF c
-- ^ Lazily reads bytes from peer. Stops once Len are read, = TmpContentSize Key (Len -> c)
-- or if connection is lost, and in either case returns the bytes -- ^ Gets size of the temp file where received content may have
-- that were read. This allows resuming interrupted transfers. -- been stored. If not present, returns 0.
= KeyFileSize Key (Len -> c) | ContentSize Key (Maybe Len -> c)
-- ^ Checks size of key file (dne = 0) -- ^ Gets size of the content of a key, when the full content is
| ReadKeyFile Key Offset (L.ByteString -> c) -- present.
| WriteKeyFile Key Offset Len L.ByteString (Bool -> c) | ReadContent Key Offset (L.ByteString -> c)
-- ^ Writes to key file starting at an offset. Returns True -- ^ Lazily reads the content of a key. Note that the content
-- once the whole content of the key is stored in the key file. -- may change while it's being sent.
| WriteContent Key Offset Len L.ByteString (Bool -> c)
-- ^ Writes content to temp file starting at an offset.
-- Once the whole content of the key has been stored, moves the
-- temp file into place and returns True.
-- --
-- Note: The ByteString may not contain the entire remaining content -- Note: The ByteString may not contain the entire remaining content
-- of the key. Only once the key file size == Len has the whole -- of the key. Only once the temp file size == Len has the whole
-- content been transferred. -- content been transferred.
| SetPresent Key UUID c | SetPresent Key UUID c
| CheckContentPresent Key (Bool -> c) | CheckContentPresent Key (Bool -> c)
-- ^ Checks if the whole content of the key is locally present. -- ^ Checks if the whole content of the key is locally present.
| RemoveKeyFile Key (Bool -> c) | RemoveContent Key (Bool -> c)
-- ^ If the key file is not present, still succeeds. -- ^ If the content is not present, still succeeds.
-- May fail if not enough copies to safely drop, etc. -- May fail if not enough copies to safely drop, etc.
| TryLockContent Key (Bool -> Proto ()) c | TryLockContent Key (Bool -> Proto ()) c
-- ^ Try to lock the content of a key, preventing it -- ^ Try to lock the content of a key, preventing it
@ -272,7 +281,7 @@ serve myuuid = go Nothing
UNLOCKCONTENT -> return () UNLOCKCONTENT -> return ()
_ -> net $ sendMessage (ERROR "expected UNLOCKCONTENT") _ -> net $ sendMessage (ERROR "expected UNLOCKCONTENT")
CHECKPRESENT key -> sendSuccess =<< local (checkContentPresent key) CHECKPRESENT key -> sendSuccess =<< local (checkContentPresent key)
REMOVE key -> sendSuccess =<< local (removeKeyFile key) REMOVE key -> sendSuccess =<< local (removeContent key)
PUT key -> do PUT key -> do
have <- local $ checkContentPresent key have <- local $ checkContentPresent key
if have if have
@ -289,20 +298,20 @@ serve myuuid = go Nothing
sendContent :: Key -> Offset -> Proto Bool sendContent :: Key -> Offset -> Proto Bool
sendContent key offset = do sendContent key offset = do
(len, content) <- readKeyFileLen key offset (len, content) <- readContentLen key offset
net $ sendMessage (DATA len) net $ sendMessage (DATA len)
net $ sendBytes len content net $ sendBytes len content
checkSuccess checkSuccess
receiveContent :: Key -> (Offset -> Message) -> Proto Bool receiveContent :: Key -> (Offset -> Message) -> Proto Bool
receiveContent key mkmsg = do receiveContent key mkmsg = do
Len n <- local $ keyFileSize key Len n <- local $ tmpContentSize key
let offset = Offset n let offset = Offset n
net $ sendMessage (mkmsg offset) net $ sendMessage (mkmsg offset)
r <- net receiveMessage r <- net receiveMessage
case r of case r of
DATA len -> do DATA len -> do
ok <- local . writeKeyFile key offset len ok <- local . writeContent key offset len
=<< net (receiveBytes len) =<< net (receiveBytes len)
sendSuccess ok sendSuccess ok
return ok return ok
@ -324,18 +333,20 @@ sendSuccess :: Bool -> Proto ()
sendSuccess True = net $ sendMessage SUCCESS sendSuccess True = net $ sendMessage SUCCESS
sendSuccess False = net $ sendMessage FAILURE sendSuccess False = net $ sendMessage FAILURE
-- Reads key file from an offset. The Len should correspond to -- Reads content from an offset. The Len should correspond to
-- the length of the ByteString, but to avoid buffering the content -- the length of the ByteString, but to avoid buffering the content
-- in memory, is gotten using keyFileSize. -- in memory, is gotten using contentSize.
readKeyFileLen :: Key -> Offset -> Proto (Len, L.ByteString) readContentLen :: Key -> Offset -> Proto (Len, L.ByteString)
readKeyFileLen key (Offset offset) = do readContentLen key (Offset offset) = go =<< local (contentSize key)
(Len totallen) <- local $ keyFileSize key where
let len = totallen - offset go Nothing = return (Len 0, L.empty)
if len <= 0 go (Just (Len totallen)) = do
then return (Len 0, L.empty) let len = totallen - offset
else do if len <= 0
content <- local $ readKeyFile key (Offset offset) then return (Len 0, L.empty)
return (Len len, content) else do
content <- local $ readContent key (Offset offset)
return (Len len, content)
connect :: Service -> Handle -> Handle -> Proto ExitCode connect :: Service -> Handle -> Handle -> Proto ExitCode
connect service hin hout = do connect service hin hout = do