Merge pull request #277 from squalus/addtostorenar2

remote: add AddToStoreNar operation

Author: sorki

hnix-store-remote/hnix-store-remote.cabal

@@ -85,6 +85,7 @@ library
     , System.Nix.Store.Remote.Types.GC
     , System.Nix.Store.Remote.Types.Handshake
     , System.Nix.Store.Remote.Types.Logger
+    , System.Nix.Store.Remote.Types.NoReply
     , System.Nix.Store.Remote.Types.ProtoVersion
     , System.Nix.Store.Remote.Types.Query
     , System.Nix.Store.Remote.Types.Query.Missing
@@ -93,6 +94,7 @@ library
     , System.Nix.Store.Remote.Types.StoreReply
     , System.Nix.Store.Remote.Types.StoreText
     , System.Nix.Store.Remote.Types.SubstituteMode
+    , System.Nix.Store.Remote.Types.SuccessCodeReply
     , System.Nix.Store.Remote.Types.TrustedFlag
     , System.Nix.Store.Remote.Types.Verbosity
     , System.Nix.Store.Remote.Types.WorkerMagic
@@ -195,6 +197,7 @@ test-suite remote-io
   ghc-options:       -rtsopts -fprof-auto "-with-rtsopts -V0"
   other-modules:
       NixDaemonSpec
+    , SampleNar
   build-depends:
       base >=4.12 && <5
     , hnix-store-core
@@ -205,6 +208,8 @@ test-suite remote-io
     , concurrency
     , containers
     , crypton
+    , data-default-class
+    , dependent-sum
     , directory
     , exceptions
     , filepath
@@ -215,5 +220,6 @@ test-suite remote-io
     , some
     , temporary
     , text
+    , time
     , unix
     , unordered-containers

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

@@ -1,5 +1,6 @@
 module System.Nix.Store.Remote.Client
   ( addToStore
+  , addToStoreNar
   , addTextToStore
   , addSignatures
   , addTempRoot
@@ -26,12 +27,14 @@ module System.Nix.Store.Remote.Client
   , module System.Nix.Store.Remote.Client.Core
   ) where
 
-import Control.Monad (when)
+import Control.Monad (void, when)
 import Control.Monad.Except (throwError)
+import Data.ByteString (ByteString)
 import Data.HashSet (HashSet)
 import Data.Map (Map)
 import Data.Set (Set)
 import Data.Some (Some)
+import Data.Word (Word64)
 
 import System.Nix.Build (BuildMode, BuildResult)
 import System.Nix.DerivedPath (DerivedPath)
@@ -73,6 +76,19 @@ addToStore name source method hashAlgo repair = do
   setNarSource source
   doReq (AddToStore name method hashAlgo repair)
 
+addToStoreNar
+  :: MonadRemoteStore m
+  => StorePath
+  -> Metadata StorePath
+  -> RepairMode
+  -> CheckMode
+  -> (Word64 -> IO(Maybe ByteString))
+  -> m ()
+addToStoreNar path metadata repair checkSigs source = do
+  setDataSource source
+  void $ doReq (AddToStoreNar path metadata repair checkSigs)
+  pure ()
+
 -- | Add @StoreText@ to the store
 -- Reference accepts repair but only uses it
 -- to throw error in case of remote talking to nix-daemon.
@@ -96,7 +112,7 @@ addSignatures
   => StorePath
   -> Set Signature
   -> m ()
-addSignatures p signatures = doReq (AddSignatures p signatures)
+addSignatures p signatures = (void . doReq) (AddSignatures p signatures)
 
 -- | Add temporary garbage collector root.
 --
@@ -105,14 +121,14 @@ addTempRoot
   :: MonadRemoteStore m
   => StorePath
   -> m ()
-addTempRoot = doReq . AddTempRoot
+addTempRoot = void . doReq . AddTempRoot
 
 -- | Add indirect garbage collector root.
 addIndirectRoot
   :: MonadRemoteStore m
   => StorePath
   -> m ()
-addIndirectRoot = doReq . AddIndirectRoot
+addIndirectRoot = void . doReq . AddIndirectRoot
 
 -- | Build a derivation available at @StorePath@
 buildDerivation
@@ -139,7 +155,7 @@ buildPaths
   => Set DerivedPath
   -> BuildMode
   -> m ()
-buildPaths a b = doReq (BuildPaths a b)
+buildPaths a b = (void . doReq) (BuildPaths a b)
 
 collectGarbage
   :: MonadRemoteStore m
@@ -151,7 +167,7 @@ ensurePath
   :: MonadRemoteStore m
   => StorePath
   -> m ()
-ensurePath = doReq . EnsurePath
+ensurePath = void . doReq . EnsurePath
 
 -- | Find garbage collector roots.
 findRoots
@@ -235,12 +251,12 @@ queryMissing = doReq . QueryMissing
 optimiseStore
   :: MonadRemoteStore m
   => m ()
-optimiseStore = doReq OptimiseStore
+optimiseStore = (void . doReq) OptimiseStore
 
 syncWithGC
   :: MonadRemoteStore m
   => m ()
-syncWithGC = doReq SyncWithGC
+syncWithGC = (void . doReq) SyncWithGC
 
 verifyStore
   :: MonadRemoteStore m

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

@@ -7,9 +7,13 @@ module System.Nix.Store.Remote.Client.Core
 import Control.Monad (unless, when)
 import Control.Monad.Except (throwError)
 import Control.Monad.IO.Class (MonadIO, liftIO)
+import Data.ByteString (ByteString)
 import Data.DList (DList)
 import Data.Some (Some(Some))
+import Data.Word (Word64)
+import Network.Socket (Socket)
 import System.Nix.Nar (NarSource)
+import System.Nix.StorePath.Metadata (Metadata(..))
 import System.Nix.Store.Remote.Logger (processOutput)
 import System.Nix.Store.Remote.MonadStore
   ( MonadRemoteStore(..)
@@ -28,11 +32,13 @@ import System.Nix.Store.Remote.Serializer
   )
 import System.Nix.Store.Remote.Types.Handshake (ClientHandshakeOutput(..))
 import System.Nix.Store.Remote.Types.Logger (Logger)
+import System.Nix.Store.Remote.Types.NoReply (NoReply(..))
 import System.Nix.Store.Remote.Types.ProtoVersion (ProtoVersion(..))
 import System.Nix.Store.Remote.Types.StoreRequest (StoreRequest(..))
 import System.Nix.Store.Remote.Types.StoreReply (StoreReply(..))
 import System.Nix.Store.Remote.Types.WorkerMagic (WorkerMagic(..))
 
+import qualified Data.ByteString
 import qualified Network.Socket.ByteString
 
 type Run m a = m (Either RemoteStoreError a, DList Logger)
@@ -69,14 +75,58 @@ doReq = \case
           Nothing ->
             throwError
               RemoteStoreError_NoNarSourceProvided
-
-      _ -> pure ()
-
-    processOutput
-    sockGetS
-      (mapErrorS RemoteStoreError_SerializerReply
-        $ getReplyS @a
-      )
+        processOutput
+        processReply
+
+      AddToStoreNar _ meta _ _ -> do
+        let narBytes = maybe 0 id $ metadataNarBytes meta
+        maybeDataSource <- takeDataSource
+        soc <- getStoreSocket
+        case maybeDataSource of
+          Nothing ->
+            if narBytes == 0 then writeFramedSource (const (pure Nothing)) soc 0
+            else throwError RemoteStoreError_NoDataSourceProvided
+          Just dataSource -> do
+            writeFramedSource dataSource soc narBytes
+        processOutput
+        pure NoReply
+
+      _ -> do
+        processOutput
+        processReply
+
+  where
+    processReply = sockGetS
+          (mapErrorS RemoteStoreError_SerializerReply
+            $ getReplyS @a
+          )
+
+writeFramedSource
+  :: forall m
+   . ( MonadIO m
+     , MonadRemoteStore m
+     )
+  => (Word64 -> IO(Maybe ByteString))
+  -> Socket
+  -> Word64
+  -> m ()
+writeFramedSource dataSource soc remainingBytes = do
+  let chunkSize = 16384
+  maybeBytes <- liftIO $ dataSource chunkSize
+  case maybeBytes of
+    Nothing -> do
+      unless (remainingBytes == 0) $ throwError RemoteStoreError_DataSourceExhausted
+      let eof :: Word64 = 0
+      sockPutS int eof
+    Just bytes -> do
+      let bytesInChunk = fromIntegral $ Data.ByteString.length bytes
+      when (bytesInChunk > chunkSize || bytesInChunk > remainingBytes) $ throwError RemoteStoreError_DataSourceReadTooLarge
+      when (bytesInChunk == 0) $ throwError RemoteStoreError_DataSourceZeroLengthRead
+      sockPutS int bytesInChunk
+      liftIO
+        $ Network.Socket.ByteString.sendAll soc bytes
+      let nextRemainingBytes = remainingBytes - bytesInChunk
+      writeFramedSource dataSource soc nextRemainingBytes
 
 greetServer
   :: MonadRemoteStore m

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

@@ -77,6 +77,8 @@ data RemoteStoreError
   | RemoteStoreError_LoggerParserFail String ByteString -- when incremental parser returns ((Fail msg leftover) :: Result)
   | RemoteStoreError_NoDataSourceProvided -- remoteStoreStateMDataSource is required but it is Nothing
   | RemoteStoreError_DataSourceExhausted -- remoteStoreStateMDataSource returned Nothing but more data was requested
+  | RemoteStoreError_DataSourceZeroLengthRead -- remoteStoreStateMDataSource returned a zero length ByteString
+  | RemoteStoreError_DataSourceReadTooLarge -- remoteStoreStateMDataSource returned a ByteString larger than the chunk size requested or the remaining bytes
   | RemoteStoreError_NoDataSinkProvided -- remoteStoreStateMDataSink is required but it is Nothing
   | RemoteStoreError_NoNarSourceProvided
   | RemoteStoreError_OperationFailed
@@ -250,6 +252,15 @@ class ( MonadIO m
    -> m ()
   setDataSource x = lift (setDataSource x)
 
+  takeDataSource :: m (Maybe (Word64 -> IO (Maybe ByteString)))
+  default takeDataSource
+   :: ( MonadTrans t
+      , MonadRemoteStore m'
+      , m ~ t m'
+      )
+   => m (Maybe (Word64 -> IO (Maybe ByteString)))
+  takeDataSource = lift takeDataSource
+
   getDataSource :: m (Maybe (Word64 -> IO (Maybe ByteString)))
   default getDataSource
    :: ( MonadTrans t
@@ -327,6 +338,11 @@ instance MonadIO m => MonadRemoteStore (RemoteStoreT m) where
   getDataSource = RemoteStoreT (gets remoteStoreStateMDataSource)
   clearDataSource = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSource = Nothing }
 
+  takeDataSource = RemoteStoreT $ do
+    x <- remoteStoreStateMDataSource <$> get
+    modify $ \s -> s { remoteStoreStateMDataSource = Nothing }
+    pure x
+
   setDataSink x = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSink = pure x }
   getDataSink = RemoteStoreT (gets remoteStoreStateMDataSink)
   clearDataSink = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSink = Nothing }

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

@@ -85,6 +85,7 @@ module System.Nix.Store.Remote.Serializer
   -- ** Reply
   , ReplySError(..)
   , opSuccess
+  , noop
   -- *** Realisation
   , derivationOutputTyped
   , realisation
@@ -1077,6 +1078,16 @@ storeRequest = Serializer
 
         pure $ Some (AddToStore pathName recursive hashAlgo repair)
 
+      WorkerOp_AddToStoreNar -> mapGetE $ do
+        storePath' <- getS storePath
+        metadata <- getS pathMetadata
+        repair <- getS bool
+        let repairMode = if repair then RepairMode_DoRepair else RepairMode_DontRepair
+        dontCheckSigs <- getS bool
+        let checkSigs = if dontCheckSigs then CheckMode_DontCheck else CheckMode_DoCheck
+
+        pure $ Some (AddToStoreNar storePath' metadata repairMode checkSigs)
+
       WorkerOp_AddTextToStore -> mapGetE $ do
         txt <- getS storeText
         paths <- getS (hashSet storePath)
@@ -1175,7 +1186,6 @@ storeRequest = Serializer
 
       w@WorkerOp_AddBuildLog -> notYet w
       w@WorkerOp_AddMultipleToStore -> notYet w
-      w@WorkerOp_AddToStoreNar -> notYet w
       w@WorkerOp_BuildPathsWithResults -> notYet w
       w@WorkerOp_ClearFailedPaths -> notYet w
       w@WorkerOp_ExportPath -> notYet w
@@ -1207,6 +1217,14 @@ storeRequest = Serializer
         putS bool (recursive == FileIngestionMethod_FileRecursive)
         putS someHashAlgo hashAlgo
 
+      Some (AddToStoreNar storePath' metadata repair checkSigs) -> mapPutE $ do
+        putS workerOp WorkerOp_AddToStoreNar
+
+        putS storePath storePath'
+        putS pathMetadata metadata
+        putS bool $ repair == RepairMode_DoRepair
+        putS bool $ checkSigs == CheckMode_DontCheck
+
       Some (AddTextToStore txt paths _repair) -> mapPutE $ do
         putS workerOp WorkerOp_AddTextToStore
 
@@ -1368,17 +1386,23 @@ mapPutER = mapErrorST ReplySError_PrimPut
 -- | Parse a bool returned at the end of simple operations.
 -- This is always 1 (@True@) so we assert that it really is so.
 -- Errors for these operations are indicated via @Logger_Error@.
-opSuccess :: NixSerializer r ReplySError ()
+opSuccess :: NixSerializer r ReplySError SuccessCodeReply
 opSuccess = Serializer
   { getS = do
       retCode <- mapGetER $ getS bool
       Control.Monad.unless
         (retCode == True)
         $ throwError ReplySError_UnexpectedFalseOpSuccess
-      pure ()
+      pure SuccessCodeReply
   , putS = \_ -> mapPutER $ putS bool True
   }
 
+noop :: a -> NixSerializer r ReplySError a
+noop ret = Serializer
+  { getS = pure ret
+  , putS = \_ -> pure ()
+  }
+
 -- *** Realisation
 
 derivationOutputTyped :: NixSerializer r ReplySError (System.Nix.Realisation.DerivationOutput OutputName)

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

@@ -170,6 +170,7 @@ processConnection workerHelper postGreet sock = do
           -- out of thin air
           () <- Data.Some.withSome someReq $ \case
             r@AddToStore {} -> perform r
+            r@AddToStoreNar {} -> perform r
             r@AddTextToStore {} -> perform r
             r@AddSignatures {} -> perform r
             r@AddTempRoot {} -> perform r

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

@@ -9,6 +9,7 @@ module System.Nix.Store.Remote.Types
   , module System.Nix.Store.Remote.Types.StoreRequest
   , module System.Nix.Store.Remote.Types.StoreText
   , module System.Nix.Store.Remote.Types.SubstituteMode
+  , module System.Nix.Store.Remote.Types.SuccessCodeReply
   , module System.Nix.Store.Remote.Types.TrustedFlag
   , module System.Nix.Store.Remote.Types.Verbosity
   , module System.Nix.Store.Remote.Types.WorkerMagic
@@ -25,6 +26,7 @@ import System.Nix.Store.Remote.Types.StoreConfig
 import System.Nix.Store.Remote.Types.StoreRequest
 import System.Nix.Store.Remote.Types.StoreText
 import System.Nix.Store.Remote.Types.SubstituteMode
+import System.Nix.Store.Remote.Types.SuccessCodeReply
 import System.Nix.Store.Remote.Types.TrustedFlag
 import System.Nix.Store.Remote.Types.Verbosity
 import System.Nix.Store.Remote.Types.WorkerMagic

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

@@ -0,0 +1,8 @@
+module System.Nix.Store.Remote.Types.NoReply
+  ( NoReply(..)
+  ) where
+
+-- | Reply type for the case where the server does not reply
+data NoReply = NoReply
+  deriving (Show)
+