diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Logger.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Logger.hs index b92fd68a..a2d3499d 100644 --- a/hnix-store-remote/src/System/Nix/Store/Remote/Logger.hs +++ b/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 diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs b/hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs index 045edd37..75f66848 100644 --- a/hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs +++ b/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 diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Types/Logger.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Types/Logger.hs index 74dd991c..543b94c7 100644 --- a/hnix-store-remote/src/System/Nix/Store/Remote/Types/Logger.hs +++ b/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)