Skip to content

Commit

Permalink
remote: implement Logger_Read
Browse files Browse the repository at this point in the history
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
  • Loading branch information
sorki committed Dec 7, 2023
1 parent 82262a1 commit 57cc9e3
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 48 deletions.
29 changes: 16 additions & 13 deletions hnix-store-remote/src/System/Nix/Store/Remote/Logger.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
75 changes: 41 additions & 34 deletions hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand All @@ -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
Expand Down Expand Up @@ -122,7 +127,7 @@ runRemoteStoreT r =
emptyState = RemoteStoreState
{ remoteStoreState_logs = mempty
, remoteStoreState_gotError = False
, remoteStoreState_mData = Nothing
, remoteStoreState_mDataSource = Nothing
, remoteStoreState_mNarSource = Nothing
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 57cc9e3

Please sign in to comment.