Download remote cache with parallel ranged GETs, and report transfer speed

CLAUDE.md

@@ -84,7 +84,7 @@ stack test --test-arguments --accept
 ## Notes
 - This project uses tasty-golden for snapshot/golden file testing
 - The test suite includes integration tests that verify taskrunner behavior
-- **S3 Test Auto-Detection**: 15 tests require S3 credentials (marked with `# s3` directive in test files)
+- **S3 Test Auto-Detection**: 16 tests require S3 credentials (marked with `# s3` directive in test files)
   - `stack test` automatically skips S3 tests if credentials are missing
   - To run S3 tests, set: `TASKRUNNER_TEST_S3_ENDPOINT`, `TASKRUNNER_TEST_S3_ACCESS_KEY`, `TASKRUNNER_TEST_S3_SECRET_KEY`
   - Use `SKIP_S3_TESTS=1` to explicitly skip S3 tests even when credentials are present

README.md

@@ -67,6 +67,8 @@ To use it, first build using another system, and the run `taskrunner` with `TASK
 
 - `TASKRUNNER_DEBUG` - whether to output debug messages to toplevel output. Note that debug messages are always written to per-task logs, regardless of this setting.
 - `TASKRUNNER_LOG_INFO` - whether to output "info" messages to toplevel output. They are minimal messages, produced only when there's actually something to be done (including fetching from cache).
+- `TASKRUNNER_S3_DOWNLOAD_CONCURRENCY` (default: `1`) - how many ranged `GET` requests to use in parallel when downloading a remote cache archive. A single stream is usually limited well below the available bandwidth, so raising this (e.g. to `8`) speeds up restoring large caches. `1` means a single plain request, as before. Transfer sizes and speeds are reported as debug messages.
+- `TASKRUNNER_S3_DOWNLOAD_CHUNK_SIZE_MIB` (default: `8`) - how much a single ranged `GET` request asks for. At most `TASKRUNNER_S3_DOWNLOAD_CONCURRENCY + 1` chunks are held in memory at a time.
 - more...
 
 ## Possible features

package.yaml

@@ -1,5 +1,5 @@
 name:                taskrunner
-version:             0.18.0.10
+version:             0.18.0.11
 github:              "githubuser/taskrunner"
 license:             BSD-3-Clause
 author:              "Author name here"
@@ -31,6 +31,7 @@ dependencies:
 - unix
 - process
 - async
+- stm
 - time
 - temporary
 - optparse-applicative

src/Control/Concurrent/Prefetch.hs

@@ -0,0 +1,73 @@
+-- | Fetch a list of items concurrently, but consume the results strictly in
+-- order. Useful for turning a sequence of independent, latency-bound requests
+-- into a stream without buffering everything in memory at once.
+module Control.Concurrent.Prefetch
+  ( Prefetch
+  , startPrefetch
+  , cancelPrefetch
+  , nextPrefetch
+  ) where
+
+import Universum
+
+import Control.Concurrent.Async (Async, async, cancel, wait)
+import Control.Concurrent.STM (TBQueue, newTBQueueIO, readTBQueue, writeTBQueue)
+import Control.Exception (mask_)
+import Data.List (delete)
+
+data Prefetch a = Prefetch
+  { queue :: TBQueue (Maybe (Async a))
+    -- ^ Results in item order. 'Nothing' marks the end of the stream.
+  , producer :: Async ()
+  , inFlight :: TVar [Async a]
+    -- ^ Fetches that have been started but not yet consumed, so that
+    -- 'cancelPrefetch' can stop them. Consumed fetches are removed, otherwise
+    -- we would keep every result alive until the whole stream is done.
+  }
+
+-- | Start fetching @items@ in the background, at most @concurrency + 1@ at a
+-- time, and hand them out in order via 'nextPrefetch'.
+--
+-- Memory use is bounded by the size of @concurrency + 1@ results, since a fetch
+-- is only started once there is room for its result.
+--
+-- Must be paired with 'cancelPrefetch' (via 'bracket' or similar), which is
+-- also what makes exceptions safe: if a fetch fails, 'nextPrefetch' rethrows it
+-- and 'cancelPrefetch' stops the remaining ones.
+startPrefetch :: Int -> [i] -> (i -> IO a) -> IO (Prefetch a)
+startPrefetch concurrency items fetch = do
+  queue <- newTBQueueIO (fromIntegral (max 1 concurrency))
+  inFlight <- newTVarIO []
+  producer <- async do
+    forM_ items \item -> do
+      -- Registering the fetch must not be interruptible, or a cancellation
+      -- landing in between would leave an unreachable thread running.
+      a <- mask_ do
+        a <- async (fetch item)
+        atomically $ modifyTVar' inFlight (a:)
+        pure a
+      -- Blocks while the consumer is behind, which is what bounds concurrency.
+      atomically $ writeTBQueue queue (Just a)
+    atomically $ writeTBQueue queue Nothing
+  pure Prefetch{queue, producer, inFlight}
+
+-- | Stop the producer and any outstanding fetches. Idempotent.
+--
+-- Note that this makes any concurrent 'nextPrefetch' block forever, so only
+-- call it once the consumer is done with the stream.
+cancelPrefetch :: Prefetch a -> IO ()
+cancelPrefetch prefetch = do
+  -- Cancel the producer first, so that it cannot start anything new while we
+  -- are cancelling what is already in flight.
+  cancel prefetch.producer
+  readTVarIO prefetch.inFlight >>= mapM_ cancel
+
+-- | Next result in item order, or 'Nothing' once all items have been handed
+-- out. Rethrows whatever the corresponding fetch threw.
+nextPrefetch :: Prefetch a -> IO (Maybe a)
+nextPrefetch prefetch =
+  atomically (readTBQueue prefetch.queue) >>= \case
+    Nothing ->
+      pure Nothing
+    Just a ->
+      Just <$> wait a `finally` atomically (modifyTVar' prefetch.inFlight (delete a))

src/RemoteCache.hs

@@ -3,10 +3,10 @@ module RemoteCache where
 
 import Universum
 
-import Control.Monad.Trans.Resource (MonadResource)
+import Control.Monad.Trans.Resource (MonadResource, ResourceT, runResourceT)
 import Amazonka.Env (newEnv, Env'(..), overrideService)
 import Amazonka.S3 (BucketName(..), ObjectKey(..), newGetObject, _NoSuchKey, StorageClass (StorageClass_REDUCED_REDUNDANCY))
-import Amazonka.S3.GetObject (GetObjectResponse(..))
+import Amazonka.S3.GetObject (GetObject(..), GetObjectResponse(..))
 import qualified Data.ByteString as BS
 import Data.Conduit ((.|), ConduitT, bracketP, runConduitRes)
 import qualified Data.Conduit.Zstd as Zstd
@@ -21,20 +21,22 @@ import System.Environment (lookupEnv)
 import Amazonka.Types ( Region(..), AccessKey(..), SecretKey(..), Service, s3AddressingStyle, S3AddressingStyle(..) )
 import Types
 import System.Process (CreateProcess(..), cleanupProcess, createProcess_, StdStream (..), proc, waitForProcess)
-import Conduit (sourceHandle, sinkHandle, foldMapC)
+import Conduit (sourceHandle, sinkHandle, foldMapC, sinkList)
 import Network.URI (parseURI, URI (..), URIAuth(..))
 import System.Directory (makeAbsolute, canonicalizePath)
 import System.FilePath (makeRelative)
 import qualified System.FilePath as FP
-import Utils (bail, logDebug, logFileName, logInfo, withStderrPipe)
+import Utils (bail, bytesfmt, logDebug, logFileName, logInfo, logWarn, timed, transferSummary, withStderrPipe)
 import qualified Amazonka as AWS
 import Control.Exception.Lens (handling)
 import System.Exit (ExitCode(..))
 import qualified Data.Conduit as C
 import qualified Data.Conduit.Text as CT
+import qualified Data.Text as Text
 import qualified Data.Text.Lazy as TL
 import qualified Data.Text.Lazy.Builder as TLB
 import Amazonka.S3.PutObject (newPutObject, PutObject(..))
+import Control.Concurrent.Prefetch (cancelPrefetch, nextPrefetch, startPrefetch)
 
 
 packTar :: MonadResource m => AppState -> Handle -> FilePath -> [FilePath] -> ConduitT () BS.ByteString m ()
@@ -89,6 +91,12 @@ data RemoteCacheSettings = RemoteCacheSettings
 
   , logsPrefix :: Text
   , logsViewUrl :: Text
+
+  -- | How many ranged GET requests to use when downloading a cache archive.
+  -- 1 (the default) means a single plain request for the whole object.
+  , s3DownloadConcurrency :: Int
+  -- | How many bytes a single ranged GET request asks for.
+  , s3DownloadChunkSize :: Int
   }
 
 getRemoteCacheSettingsFromEnv :: MonadIO m => m RemoteCacheSettings
@@ -101,8 +109,23 @@ getRemoteCacheSettingsFromEnv = liftIO do
   remoteCachePrefix <- maybe "taskrunner/" toText <$> lookupEnv "TASKRUNNER_REMOTE_CACHE_PREFIX"
   logsPrefix <- maybe (error "TASKRUNNER_LOGS_PREFIX not provided") toText <$> lookupEnv "TASKRUNNER_LOGS_PREFIX"
   logsViewUrl <- maybe (error "TASKRUNNER_LOGS_VIEW_URL not provided") toText <$> lookupEnv "TASKRUNNER_LOGS_VIEW_URL"
+  s3DownloadConcurrency <- lookupPositiveIntEnv "TASKRUNNER_S3_DOWNLOAD_CONCURRENCY" 1
+  s3DownloadChunkSizeMiB <- lookupPositiveIntEnv "TASKRUNNER_S3_DOWNLOAD_CHUNK_SIZE_MIB" 8
+  let s3DownloadChunkSize = s3DownloadChunkSizeMiB * 1024 * 1024
   pure RemoteCacheSettings{..}
 
+lookupPositiveIntEnv :: String -> Int -> IO Int
+lookupPositiveIntEnv name defaultValue =
+  lookupEnv name >>= \case
+    Nothing ->
+      pure defaultValue
+    Just str ->
+      case readMaybe str of
+        Just value | value > 0 ->
+          pure value
+        _ ->
+          error $ toText name <> " must be a positive integer, got: " <> show str
+
 parseEndpoint :: Text -> Maybe (Service -> Service)
 parseEndpoint "default-aws" = Just id
 parseEndpoint s = do
@@ -114,7 +137,6 @@ parseEndpoint s = do
     . (\svc -> svc { s3AddressingStyle = S3AddressingStylePath })
 
 -- TODO:
--- - report speed, size etc.
 -- - integrate amazonka logging
 -- - handle errors
 saveCache
@@ -145,13 +167,18 @@ saveCache appState settings relativeCacheRoot files archiveName = do
 
     logDebug appState $ "Uploading to s3://" <> bucket <> "/" <> objectKey
 
-    withStderrPipe appState \stderrHandle ->
+    packedBytes <- newIORef 0
+    uploadedBytes <- newIORef 0
+
+    (_, elapsed) <- timed $ withStderrPipe appState \stderrHandle ->
       runConduitRes do
         let multipartUpload = (newCreateMultipartUpload (BucketName bucket) (ObjectKey objectKey) :: CreateMultipartUpload)
               { storageClass = Just StorageClass_REDUCED_REDUNDANCY }
         result <-
           packTar appState stderrHandle cacheRoot filesRelativeToCacheRoot
+          .| countBytes packedBytes
           .| Zstd.compress 3
+          .| countBytes uploadedBytes
           .| streamUpload env Nothing multipartUpload
         case result of
           Left (_, err) ->
@@ -160,6 +187,19 @@ saveCache appState settings relativeCacheRoot files archiveName = do
             liftIO $ logDebug appState "Upload success"
             pure ()
 
+    packed <- readIORef packedBytes
+    uploaded <- readIORef uploadedBytes
+    -- Note the rate covers the whole pipeline (tar, zstd and the upload), not
+    -- just the network part.
+    logDebug appState $ "Packed and uploaded " <> transferSummary uploaded elapsed
+      <> ", compressed from " <> toText (bytesfmt "%.2f" packed)
+
+-- | Pass data through unchanged, accumulating the total number of bytes seen.
+countBytes :: MonadIO m => IORef Int -> ConduitT BS.ByteString BS.ByteString m ()
+countBytes ref = C.awaitForever \chunk -> do
+  modifyIORef' ref (+ BS.length chunk)
+  C.yield chunk
+
 data LogMode = NoLog | Log deriving (Eq, Show)
 
 restoreCache
@@ -181,14 +221,129 @@ restoreCache appState settings cacheRoot archiveName logMode = do
       logDebug appState $ "Remote cache archive not found s3://" <> bucket <> "/" <> objectKey
       pure False
 
-  handling _NoSuchKey onNoSuchKey $ withStderrPipe appState \stderrHandle ->
-    runConduitRes do
-      response <- AWS.send env $ newGetObject (BucketName bucket) (ObjectKey objectKey)
+  handling _NoSuchKey onNoSuchKey $ withStderrPipe appState \stderrHandle -> do
+    downloadedBytes <- newIORef 0
+
+    (_, elapsed) <- timed $ runResourceT do
+      source <- startDownload appState settings env (BucketName bucket) (ObjectKey objectKey)
+
+      -- Only now that the archive is known to exist: say so, and start unpacking.
       when (logMode == Log) do
-        liftIO $ logInfo appState $ "Found remote cache " <> archiveName <> ", restoring"
-      response.body.body
-            .| unpackTar appState stderrHandle cacheRoot
-      pure True
+        logInfo appState $ "Found remote cache " <> archiveName <> ", restoring"
+
+      C.runConduit $
+        source
+          .| countBytes downloadedBytes
+          .| unpackTar appState stderrHandle cacheRoot
+
+    downloaded <- readIORef downloadedBytes
+    -- The size is that of the compressed archive, and the rate covers the whole
+    -- pipeline (the download, zstd and tar), not just the network part.
+    logDebug appState $ "Downloaded and unpacked " <> transferSummary downloaded elapsed
+
+    pure True
+
+-- | Make the initial request for an S3 object, and return a source streaming its
+-- contents. Uses several parallel ranged GET requests when
+-- @s3DownloadConcurrency@ is above 1: a single stream tends to be limited well
+-- below the available bandwidth, so fetching a few ranges at once is noticeably
+-- faster for large archives.
+--
+-- Chunks are emitted strictly in order, so downstream sees the same byte stream
+-- either way.
+--
+-- Note the first request deliberately happens before the returned source is
+-- consumed, so that a missing object is reported (as '_NoSuchKey') before the
+-- caller starts anything else. Conduit initialises sinks before pulling from the
+-- source, so folding this into the pipeline would mean 'unpackTar' had already
+-- spawned tar by the time we found out, which then complains about its empty
+-- input on every cache miss.
+startDownload
+  :: AppState
+  -> RemoteCacheSettings
+  -> AWS.Env
+  -> BucketName
+  -> ObjectKey
+  -> ResourceT IO (ConduitT () BS.ByteString (ResourceT IO) ())
+startDownload appState settings env bucket key
+  | settings.s3DownloadConcurrency <= 1 = do
+      response <- AWS.send env $ newGetObject bucket key
+      pure response.body.body
+  | otherwise = do
+      -- The first request doubles as the existence check (so that _NoSuchKey is
+      -- still thrown from here) and tells us the total size via Content-Range,
+      -- which is what lets us plan the remaining ranges without a separate
+      -- HeadObject request. Note that HeadObject would not do: S3 answers HEAD
+      -- with an empty body, so a missing object does not come back as
+      -- _NoSuchKey there.
+      firstResponse <- AWS.send env $ rangedGetObject bucket key (0, fromIntegral chunkSize - 1)
+
+      -- A 206 means the range was honoured and the body is only the first
+      -- chunk; anything else (a server ignoring Range, or an object smaller
+      -- than one chunk served whole) means we already have everything.
+      if firstResponse.httpStatus /= 206 then
+        pure firstResponse.body.body
+      else case parseContentRangeTotal =<< firstResponse.contentRange of
+        Nothing -> do
+          -- Partial response, but we cannot tell how much is left, so we cannot
+          -- safely stream this body and stop. Start over in a single request.
+          logWarn appState $ "Could not determine object size from Content-Range: "
+            <> show firstResponse.contentRange <> ", downloading in a single request"
+          response <- AWS.send env $ newGetObject bucket key
+          pure response.body.body
+        Just total -> do
+          logDebug appState $ "Object size: " <> toText (bytesfmt "%.2f" total)
+            <> ", downloading with concurrency " <> show settings.s3DownloadConcurrency
+          case chunkRanges chunkSize (fromIntegral chunkSize) total of
+            [] ->
+              -- Object fits in a single chunk, which we already have.
+              pure firstResponse.body.body
+            remainingRanges ->
+              -- Start prefetching the rest right away, so it overlaps with
+              -- streaming the first chunk downstream.
+              pure $ bracketP
+                (startPrefetch settings.s3DownloadConcurrency remainingRanges
+                  (fetchRange env bucket key))
+                cancelPrefetch
+                \prefetch -> do
+                  firstResponse.body.body
+                  let go = liftIO (nextPrefetch prefetch) >>= \case
+                        Nothing -> pure ()
+                        Just chunk -> C.yield chunk >> go
+                  go
+  where
+  chunkSize = max 1 settings.s3DownloadChunkSize
+
+-- | Download a single byte range of an object into memory.
+fetchRange :: AWS.Env -> BucketName -> ObjectKey -> (Integer, Integer) -> IO BS.ByteString
+fetchRange env bucket key range' =
+  AWS.runResourceT do
+    response <- AWS.send env $ rangedGetObject bucket key range'
+    BS.concat <$> C.runConduit (response.body.body .| sinkList)
+
+-- | A GET request for an inclusive byte range, as in the HTTP @Range@ header.
+rangedGetObject :: BucketName -> ObjectKey -> (Integer, Integer) -> GetObject
+rangedGetObject bucket key (start, end) =
+  (newGetObject bucket key)
+    { range = Just $ "bytes=" <> show start <> "-" <> show end }
+
+-- | Split @[start, total)@ into consecutive inclusive ranges of at most
+-- @chunkSize@ bytes each.
+chunkRanges :: Int -> Integer -> Integer -> [(Integer, Integer)]
+chunkRanges chunkSize start total
+  | start >= total = []
+  | otherwise =
+      (start, min (start + size) total - 1) : chunkRanges chunkSize (start + size) total
+  where
+  size = fromIntegral (max 1 chunkSize)
+
+-- | Total object size from a @Content-Range@ header value, e.g. the 52428800 in
+-- @bytes 0-8388607/52428800@. 'Nothing' if the size is unknown (@*@) or the
+-- header is malformed.
+parseContentRangeTotal :: Text -> Maybe Integer
+parseContentRangeTotal header = do
+  let total = Text.drop 1 $ Text.dropWhile (/= '/') header
+  readMaybe (toString total)
 
 getLatestBuildHash
   :: AppState