added StoreContentTo

This is needed in addition to StoreContent, because retrieveKeyFile can
be used to retrieve to different destination files, not only the tmp
file for a key.

This commit was sponsored by Ole-Morten Duesund on Patreon.

diff --git a/P2P/Annex.hs b/P2P/Annex.hs
index 5e1763fc6ed8e16943cc8ba4f0c1e84e1fd292c3..d0c00def37091e8959111d5463dc8134dd2f4cbd 100644 (file)
--- a/P2P/Annex.hs
+++ b/P2P/Annex.hs
@@ -16,7 +16,6 @@ module P2P.Annex
 import Annex.Common
 import Annex.Content
 import Annex.Transfer
-import Annex.Notification
 import P2P.Protocol
 import P2P.IO
 import Logs.Location
@@ -46,6 +45,9 @@ runLocal runmode runner a = case a of
                tmp <- fromRepo $ gitAnnexTmpObjectLocation k
                size <- liftIO $ catchDefaultIO 0 $ getFileSize tmp
                runner (next (Len size))
+       FileSize f next -> do
+               size <- liftIO $ catchDefaultIO 0 $ getFileSize f
+               runner (next (Len size))
        ContentSize k next -> do
                let getsize = liftIO . catchMaybeIO . getFileSize
                size <- inAnnex' isJust Nothing getsize k
@@ -69,16 +71,15 @@ runLocal runmode runner a = case a of
                                        Left _ -> return Nothing
                                        Right b -> runner (next b)
                        _ -> return Nothing
-       WriteContent k af (Offset o) (Len l) b next -> do
+       StoreContent k af o l b next -> do
                ok <- flip catchNonAsync (const $ return False) $
                        transfer download k af $
-                               getViaTmp AlwaysVerify k $ \tmp -> liftIO $ do
-                                       withBinaryFile tmp WriteMode $ \h -> do
-                                               when (o /= 0) $
-                                                       hSeek h AbsoluteSeek o
-                                               L.hPut h b
-                                       sz <- getFileSize tmp
-                                       return (toInteger sz == l, UnVerified)
+                               getViaTmp AlwaysVerify k $ \tmp ->
+                                       unVerified $ storefile tmp o l b
+               runner (next ok)
+       StoreContentTo dest o l b next -> do
+               ok <- flip catchNonAsync (const $ return False) $
+                       storefile dest o l b
                runner (next ok)
        SetPresent k u next -> do
                v <- tryNonAsync $ logChange k u InfoPresent
@@ -111,10 +112,17 @@ runLocal runmode runner a = case a of
                                next
                        Right _ -> runner next
   where
-       transfer mk k af a = case runmode of
+       transfer mk k af ta = case runmode of
                -- Update transfer logs when serving.
                Serving theiruuid -> 
-                       mk theiruuid k af noRetry (const a) noNotification
+                       mk theiruuid k af noRetry (const ta) noNotification
                -- Transfer logs are updated higher in the stack when
                -- a client.
-               Client -> a
+               Client -> ta
+       storefile dest (Offset o) (Len l) b = liftIO $ do
+               withBinaryFile dest WriteMode $ \h -> do
+                       when (o /= 0) $
+                               hSeek h AbsoluteSeek o
+                       L.hPut h b
+               sz <- getFileSize dest
+               return (toInteger sz == l)

diff --git a/P2P/Protocol.hs b/P2P/Protocol.hs
index 53c3265efe64f93be83278a54f7ce39dd26387e8..b2b734e48ad37ae9d4114294695eaef091221e06 100644 (file)
--- a/P2P/Protocol.hs
+++ b/P2P/Protocol.hs
@@ -189,16 +189,25 @@ data LocalF c
        = TmpContentSize Key (Len -> c)
        -- ^ Gets size of the temp file where received content may have
        -- been stored. If not present, returns 0.
+       | FileSize FilePath (Len -> c)
+       -- ^ Gets size of the content of a file. If not present, returns 0.
        | ContentSize Key (Maybe Len -> c)
        -- ^ Gets size of the content of a key, when the full content is
        -- present.
        | ReadContent Key AssociatedFile Offset (L.ByteString -> c)
        -- ^ Lazily reads the content of a key. Note that the content
        -- may change while it's being sent.
-       | WriteContent Key AssociatedFile Offset Len L.ByteString (Bool -> c)
-       -- ^ Writes content to temp file starting at an offset.
+       | StoreContent Key AssociatedFile Offset Len L.ByteString (Bool -> c)
+       -- ^ Stores content to the key's 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.
+       -- temp file into place as the content of the key, and returns True.
+       --
+       -- Note: The ByteString may not contain the entire remaining content
+       -- of the key. Only once the temp file size == Len has the whole
+       -- content been transferred.
+       | StoreContentTo FilePath Offset Len L.ByteString (Bool -> c)
+       -- ^ Stores the content to a temp file starting at an offset.
+       -- Once the whole content of the key has been stored, returns True.
        --
        -- Note: The ByteString may not contain the entire remaining content
        -- of the key. Only once the temp file size == Len has the whole
@@ -246,16 +255,16 @@ checkPresent key = do
  -}
 lockContentWhile 
        :: MonadMask m 
-       => (forall r. Proto r -> m r)
+       => (forall r. r -> Proto r -> m r)
        -> Key
-       -> (Bool -> m ())
-       -> m ()
+       -> (Bool -> m a)
+       -> m a
 lockContentWhile runproto key a = bracket setup cleanup a
   where
-       setup = runproto $ do
+       setup = runproto False $ do
                net $ sendMessage (LOCKCONTENT key)
                checkSuccess
-       cleanup True = runproto $ net $ sendMessage UNLOCKCONTENT
+       cleanup True = runproto () $ net $ sendMessage UNLOCKCONTENT
        cleanup False = return ()
 
 remove :: Key -> Proto Bool
@@ -263,8 +272,11 @@ remove key = do
        net $ sendMessage (REMOVE key)
        checkSuccess
 
-get :: Key -> AssociatedFile -> Proto Bool
-get key af = receiveContent key af (\offset -> GET offset af key)
+get :: FilePath -> Key -> AssociatedFile -> Proto Bool
+get dest key af = receiveContent sizer storer (\offset -> GET offset af key)
+  where
+       sizer = fileSize dest
+       storer = storeContentTo dest
 
 put :: Key -> AssociatedFile -> Proto Bool
 put key af = do
@@ -349,7 +361,9 @@ serveAuthed myuuid = void $ serverLoop handler
                if have
                        then net $ sendMessage ALREADY_HAVE
                        else do
-                               ok <- receiveContent key af PUT_FROM
+                               let sizer = tmpContentSize key
+                               let storer = storeContent key af
+                               ok <- receiveContent sizer storer PUT_FROM
                                when ok $
                                        local $ setPresent key myuuid
                return ServerContinue
@@ -370,15 +384,15 @@ sendContent key af offset = do
        net $ sendBytes len content
        checkSuccess
 
-receiveContent :: Key -> AssociatedFile -> (Offset -> Message) -> Proto Bool
-receiveContent key af mkmsg = do
-       Len n <- local $ tmpContentSize key
+receiveContent :: Local Len -> (Offset -> Len -> L.ByteString -> Local Bool) -> (Offset -> Message) -> Proto Bool
+receiveContent sizer storer mkmsg = do
+       Len n <- local sizer
        let offset = Offset n
        net $ sendMessage (mkmsg offset)
        r <- net receiveMessage
        case r of
                DATA len -> do
-                       ok <- local . writeContent key af offset len
+                       ok <- local . storer offset len
                                =<< net (receiveBytes len)
                        sendSuccess ok
                        return ok