remote: implement Logger_Read

Adds `setDataSource` which can be used to set
a function to be polled when daemon asks
for data using `Logger_Read`.

Function should return `Nothing` when all
data was read.

`clearDataSource` should be used after
the operation using the data source is finished.

Related to #265

Author: sorki

hnix-store-remote/src/System/Nix/Store/Remote/Logger.hs

@@ -3,18 +3,19 @@ module System.Nix.Store.Remote.Logger
   ) where
 
 import Control.Monad.Except (throwError)
+import Control.Monad.IO.Class (liftIO)
 import Data.ByteString (ByteString)
 import Data.Serialize (Result(..))
-import System.Nix.Store.Remote.Serialize.Prim (putByteString)
 import System.Nix.Store.Remote.Serializer (LoggerSError, logger, runSerialT)
-import System.Nix.Store.Remote.Socket (sockGet8, sockPut)
-import System.Nix.Store.Remote.MonadStore (MonadRemoteStore, RemoteStoreError(..), appendLog, clearData, getData, getProtoVersion, setError)
+import System.Nix.Store.Remote.Socket (sockGet8)
+import System.Nix.Store.Remote.MonadStore (MonadRemoteStore, RemoteStoreError(..), appendLog, getDataSource, getStoreSocket, getProtoVersion, setError)
 import System.Nix.Store.Remote.Types.Logger (Logger(..))
 import System.Nix.Store.Remote.Types.ProtoVersion (ProtoVersion)
 
 import qualified Control.Monad
 import qualified Data.Serialize.Get
 import qualified Data.Serializer
+import qualified Network.Socket.ByteString
 
 processOutput
   :: MonadRemoteStore m
@@ -55,16 +56,18 @@ processOutput = do
           Logger_Last -> appendLog Logger_Last
 
           -- Read data from source
-          Logger_Read _n -> do
-            mdata <- getData
-            case mdata of
-              Nothing   -> throwError RemoteStoreError_NoDataProvided
-              Just part -> do
-                -- XXX: we should check/assert part size against n of (Read n)
-                -- ^ not really, this is just an indicator how big of a chunk
-                -- to read from the source
-                sockPut $ putByteString part
-                clearData
+          Logger_Read size -> do
+            mSource <- getDataSource
+            case mSource of
+              Nothing   ->
+                throwError RemoteStoreError_NoDataSourceProvided
+              Just source -> do
+                mChunk <- liftIO $ source size
+                case mChunk of
+                  Nothing -> throwError RemoteStoreError_DataSourceExhausted
+                  Just chunk -> do
+                    sock <- getStoreSocket
+                    liftIO $ Network.Socket.ByteString.sendAll sock chunk
 
             loop
 

hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs

@@ -36,7 +36,11 @@ import System.Nix.Store.Remote.Types.StoreConfig (HasStoreSocket(..), StoreConfi
 data RemoteStoreState = RemoteStoreState {
     remoteStoreState_logs :: [Logger]
   , remoteStoreState_gotError :: Bool
-  , remoteStoreState_mData :: Maybe ByteString
+  , remoteStoreState_mDataSource :: Maybe (Word64 -> IO (Maybe ByteString))
+  -- ^ Source for @Logger_Read@, this will be called repeatedly
+  -- as the daemon requests chunks of size @Word64@.
+  -- If the function returns Nothing and daemon tries to read more
+  -- data an error is thrown.
   , remoteStoreState_mNarSource :: Maybe (NarSource IO)
   }
 
@@ -55,7 +59,8 @@ data RemoteStoreError
   | RemoteStoreError_IOException SomeException
   | RemoteStoreError_LoggerLeftovers String ByteString -- when there are bytes left over after incremental logger parser is done, (Done x leftover), first param is show x
   | RemoteStoreError_LoggerParserFail String ByteString -- when incremental parser returns ((Fail msg leftover) :: Result)
-  | RemoteStoreError_NoDataProvided
+  | RemoteStoreError_NoDataSourceProvided -- remoteStoreState_mDataSource is required but it is Nothing
+  | RemoteStoreError_DataSourceExhausted -- remoteStoreState_mDataSource returned Nothing but more data was requested
   | RemoteStoreError_NoNarSourceProvided
   | RemoteStoreError_OperationFailed
   | RemoteStoreError_ProtocolMismatch
@@ -122,7 +127,7 @@ runRemoteStoreT r =
     emptyState = RemoteStoreState
       { remoteStoreState_logs = mempty
       , remoteStoreState_gotError = False
-      , remoteStoreState_mData = Nothing
+      , remoteStoreState_mDataSource = Nothing
       , remoteStoreState_mNarSource = Nothing
       }
 
@@ -182,34 +187,6 @@ class ( MonadIO m
     => m Bool
   gotError = lift gotError
 
-  setData :: ByteString -> m ()
-  default setData
-   :: ( MonadTrans t
-      , MonadRemoteStoreR r m'
-      , m ~ t m'
-      )
-   => ByteString
-   -> m ()
-  setData = lift . setData
-
-  getData :: m (Maybe ByteString)
-  default getData
-    :: ( MonadTrans t
-       , MonadRemoteStoreR r m'
-       , m ~ t m'
-       )
-    => m (Maybe ByteString)
-  getData = lift getData
-
-  clearData :: m ()
-  default clearData
-    :: ( MonadTrans t
-       , MonadRemoteStoreR r m'
-       , m ~ t m'
-       )
-    => m ()
-  clearData = lift clearData
-
   getStoreDir :: m StoreDir
   default getStoreDir
     :: ( MonadTrans t
@@ -247,6 +224,36 @@ class ( MonadIO m
    => m (Maybe (NarSource IO))
   takeNarSource = lift takeNarSource
 
+  setDataSource :: (Word64 -> IO (Maybe ByteString)) -> m ()
+  default setDataSource
+   :: ( MonadTrans t
+      , MonadRemoteStoreR r m'
+      , m ~ t m'
+      )
+   => (Word64 -> IO (Maybe ByteString))
+   -> m ()
+  setDataSource x = lift (setDataSource x)
+
+  getDataSource :: m (Maybe (Word64 -> IO (Maybe ByteString)))
+  default getDataSource
+   :: ( MonadTrans t
+      , MonadRemoteStoreR r m'
+      , m ~ t m'
+      )
+   => m (Maybe (Word64 -> IO (Maybe ByteString)))
+  getDataSource = lift getDataSource
+
+  clearDataSource :: m ()
+  default clearDataSource
+    :: ( MonadTrans t
+       , MonadRemoteStoreR r m'
+       , m ~ t m'
+       )
+    => m ()
+  clearDataSource = lift clearDataSource
+
+
+
 instance MonadRemoteStoreR r m => MonadRemoteStoreR r (StateT s m)
 instance MonadRemoteStoreR r m => MonadRemoteStoreR r (ReaderT r m)
 instance MonadRemoteStoreR r m => MonadRemoteStoreR r (ExceptT RemoteStoreError m)
@@ -271,9 +278,9 @@ instance ( MonadIO m
   clearError = RemoteStoreT $ modify $ \s -> s { remoteStoreState_gotError = False }
   gotError = remoteStoreState_gotError <$> RemoteStoreT get
 
-  getData = remoteStoreState_mData <$> RemoteStoreT get
-  setData x = RemoteStoreT $ modify $ \s -> s { remoteStoreState_mData = pure x }
-  clearData = RemoteStoreT $ modify $ \s -> s { remoteStoreState_mData = Nothing }
+  setDataSource x = RemoteStoreT $ modify $ \s -> s { remoteStoreState_mDataSource = pure x }
+  getDataSource = remoteStoreState_mDataSource <$> RemoteStoreT get
+  clearDataSource = RemoteStoreT $ modify $ \s -> s { remoteStoreState_mDataSource = Nothing }
 
   setNarSource x = RemoteStoreT $ modify $ \s -> s { remoteStoreState_mNarSource = pure x }
   takeNarSource = RemoteStoreT $ do

hnix-store-remote/src/System/Nix/Store/Remote/Types/Logger.hs

@@ -81,7 +81,7 @@ word64ToLoggerOpCode = \case
 
 data Logger
   = Logger_Next Text
-  | Logger_Read Int         -- data needed from source
+  | Logger_Read Word64      -- data needed from source
   | Logger_Write ByteString -- data for sink
   | Logger_Last
   | Logger_Error (Either BasicError ErrorInfo)