Amazon S3 backup support

Uses amazonka-s3, which auto discovers credentials.

Did some manual testing:

```
glean --tier test --db-root ~/glean backup myrepo/0 s3:mybucket/jade-test
glean --tier test --db-root ~/glean2 restore s3:mybucket/jade-test/myrepo.0
```

Configuration ~/.config/glean/server/test, though I think this doesn't
actually matter?

```
{
    "backup": {
        "allowed": ["myrepo"],
        "location": "s3:mybucket/jade-test"
    },
    "restore": {
        "enabled": true
    }
}
```

This feature really needs some documentation (especially given I don't
know if fb expected anyone external to immediately go and use it so it's
probably in a wiki somewhere?); I would appreciate the docs that exist
getting released, or I can maybe write some.

Author: lf-

glean.cabal.in

@@ -84,6 +84,10 @@ flag opt
 flag benchmarks
      default: False
 
+-- Support for backups to Amazon S3
+flag s3-support
+     default: True
+
 -- run tests that require clang
 flag clang-tests
      default: True
@@ -623,6 +627,17 @@ library db
 
         Glean.Logger
 
+    if flag(s3-support)
+        cpp-options: -DENABLE_S3=1
+        exposed-modules: Glean.Database.Backup.S3
+        build-depends:
+            amazonka,
+            amazonka-s3,
+            conduit,
+            unliftio,
+            microlens,
+            http-client,
+
     other-modules:
         Paths_glean
     autogen-modules:
@@ -691,6 +706,7 @@ library backend-local
         Glean.Backend.Local
         Glean.Backend.Logging
         Glean.Dump
+
     build-depends:
         glean:angle,
         glean:backend-api,
@@ -1100,6 +1116,8 @@ executable glean-server
         Glean.Handler
         Glean.Server.Config
         Glean.Server.Sharding
+    if flag(s3-support)
+        cpp-options: -DENABLE_S3=1
     extra-libraries: stdc++
     build-depends:
         glean:backend-api,
@@ -1225,6 +1243,8 @@ executable glean
         GleanCLI.Restore
         GleanCLI.Write
     ghc-options: -main-is GleanCLI -with-rtsopts=-I0
+    if flag(s3-support)
+        cpp-options: -DENABLE_S3=1
     extra-libraries: stdc++
     build-depends:
         glean:cli-types,
@@ -2130,6 +2150,26 @@ test-suite backup
         glean:stubs,
         glean:util
 
+test-suite backup-s3
+    import: test
+    type: exitcode-stdio-1.0
+    main-is: BackupTestS3.hs
+    ghc-options: -main-is BackupTestS3
+    build-depends:
+        glean:config,
+        glean:core,
+        glean:db,
+        glean:if-glean-hs,
+        glean:if-internal-hs,
+        glean:schema,
+        glean:stubs,
+        glean:util
+
+    if flag(s3-support)
+        buildable: True
+    else
+        buildable: False
+
 test-suite catalog
     import: test
     type: exitcode-stdio-1.0

glean/db/Glean/Database/Backup/S3.hs

@@ -0,0 +1,312 @@
+{-
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+  All rights reserved.
+
+  This source code is licensed under the BSD-style license found in the
+  LICENSE file in the root directory of this source tree.
+-}
+{-# LANGUAGE TypeApplications #-}
+
+-- | S3 backup backend for Glean.
+module Glean.Database.Backup.S3 (withS3Backups, fakeS3Backend) where
+
+import Control.Exception.Safe (throwIO)
+import Control.Monad ((<=<), (>=>))
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Builder as BB
+import qualified Data.ByteString.Lazy as LBS
+import Data.Default (Default (..))
+import Data.Functor (void)
+import qualified Data.HashMap.Strict as HashMap
+import qualified Data.Map as Map
+import Data.Maybe (catMaybes)
+import qualified Data.Set as Set
+import Data.Text (Text)
+import qualified Data.Text as Text
+import qualified Data.Text.Encoding as Text
+import qualified Network.HTTP.Client as HTTP
+import System.IO.Unsafe (unsafePerformIO)
+import UnliftIO (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)
+import UnliftIO.Async (forConcurrently)
+import UnliftIO.Exception.Lens
+
+import qualified Amazonka as AWS
+import qualified Amazonka.S3 as S3
+import Amazonka.S3.DeleteObject
+import Amazonka.S3.GetObject
+import Amazonka.S3.HeadObject
+import Amazonka.S3.ListObjectsV2
+import Amazonka.S3.PutObject
+import Amazonka.S3.Types.Object
+import Conduit
+import Lens.Micro
+
+import qualified Glean.Database.Config as DBConfig
+import Glean.Database.Backup.Backend
+import Glean.Database.Exception
+import Glean.Internal.Types (Meta (..))
+import Glean.Types (Repo (..))
+import qualified Glean.Types as Thrift hiding (Exception)
+import Glean.Util.Some
+import qualified Thrift.Protocol.JSON as Thrift
+
+withS3Backups :: DBConfig.Config -> DBConfig.Config
+withS3Backups cfg@DBConfig.Config{..} =
+  cfg{DBConfig.cfgBackupBackends = cfgBackupBackends <> HashMap.fromList [("s3", Some (def @S3Backend))]}
+
+-- | S3 backup backend, which auto discovers its environment configuration.
+data S3Backend = S3Backend {s3AwsEnvLazy :: LazyInit AWS.Env, s3BucketFactory :: (LazyInit AWS.Env -> Text -> Some S3BucketApi)}
+
+instance Default S3Backend where
+  def = newS3Backend (\env name -> GenuineS3Bucket env (S3.BucketName name))
+
+newS3Backend :: (S3BucketApi bucket) => (LazyInit AWS.Env -> Text -> bucket) -> S3Backend
+newS3Backend factory =
+  S3Backend
+    { s3AwsEnvLazy = LazyInit{initializer = AWS.newEnv AWS.discover, value = unsafePerformIO . newIORef $ Nothing}
+    , s3BucketFactory = \env name -> Some (factory env name)
+    }
+
+-- | Creates a new fake in-memory S3 backend.
+fakeS3Backend :: (MonadIO m) => m S3Backend
+fakeS3Backend = do
+  fakeFiles <- newIORef Map.empty
+  pure $ newS3Backend (\_env name -> FakeS3Bucket{fakeBucketName = name, fakeFiles})
+
+instance Backend S3Backend where
+  fromPath S3Backend{s3AwsEnvLazy, s3BucketFactory} path = do
+    let (bucketName, bucketBasePath') = Text.breakOn "/" path
+    (_, bucketBasePath) <- Text.uncons bucketBasePath'
+    Just . Some $
+      S3Site
+        { s3Client = s3BucketFactory s3AwsEnvLazy bucketName
+        , bucketBasePath
+        }
+
+-- | File in an S3 bucket.
+data Item = Item
+  { itemPath :: Text
+  }
+
+type Metadata = HashMap.HashMap Text Text
+
+-- | A client to an S3 bucket, as a seam for unit testing purposes.
+class S3BucketApi a where
+  -- | Gets the name of the bucket this is a client for.
+  bucketName :: a -> Text
+
+  {- | Lists items with a prefix.
+
+  We aren't using the delimiter functionality (which does dedupe by common
+  prefixes) as we need both the meta and the file itself to exist for
+  consistency, so we need to catch both of those.
+
+  This also doesn't concern itself with pagination because our consuming API
+  doesn't either.
+  -}
+  listItemsWithPrefix :: (MonadResource m, MonadUnliftIO m) => a -> Text -> m [Item]
+
+  -- | Uploads a file on disk to the given path in the bucket.
+  uploadFile :: (MonadResource m, MonadUnliftIO m) => a -> Text -> Metadata -> AWS.RequestBody -> m ()
+
+  {- | Downloads a file from S3.
+
+  Throws for any error but absent key.
+  -}
+  downloadFile :: (MonadResource m, MonadUnliftIO m) => a -> Text -> m (Maybe (Metadata, AWS.ResponseBody))
+
+  {- | Checks a file exists on S3.
+
+  Throws for any error but absent key.
+  -}
+  headFile :: (MonadResource m, MonadUnliftIO m) => a -> Text -> m (Maybe Metadata)
+
+  -- | Deletes a file in the bucket.
+  deleteFile :: (MonadResource m, MonadUnliftIO m) => a -> Text -> m ()
+
+instance S3BucketApi (Some S3BucketApi) where
+  bucketName (Some bucket) = bucketName bucket
+  listItemsWithPrefix (Some bucket) = listItemsWithPrefix bucket
+  uploadFile (Some bucket) = uploadFile bucket
+  downloadFile (Some bucket) = downloadFile bucket
+  headFile (Some bucket) = headFile bucket
+  deleteFile (Some bucket) = deleteFile bucket
+
+-- TODO(jadel): surely there's a better way
+data LazyInit a = LazyInit {initializer :: IO a, value :: IORef (Maybe a)}
+
+getOrInit :: LazyInit b -> IO b
+getOrInit LazyInit{initializer, value} = do
+  actual <- readIORef value
+  case actual of
+    Just v -> pure v
+    Nothing -> do
+      v <- initializer
+      writeIORef value (Just v)
+      pure v
+
+data GenuineS3Bucket = GenuineS3Bucket
+  { awsEnvLazy :: LazyInit AWS.Env
+  , s3BucketName :: S3.BucketName
+  }
+
+instance S3BucketApi GenuineS3Bucket where
+  bucketName = (^. S3._BucketName) . s3BucketName
+
+  listItemsWithPrefix GenuineS3Bucket{awsEnvLazy, s3BucketName} prefix = do
+    awsEnv <- liftIO $ getOrInit awsEnvLazy
+    let request =
+          newListObjectsV2 s3BucketName
+            & listObjectsV2_prefix
+            ?~ prefix
+    runConduit $
+      AWS.paginate awsEnv request
+        .| concatMapC (^. listObjectsV2Response_contents)
+        .| concatC
+        .| mapC (\obj -> Item{itemPath = obj ^. object_key . S3._ObjectKey})
+        .| sinkList
+
+  uploadFile GenuineS3Bucket{awsEnvLazy, s3BucketName} name metadata body = do
+    awsEnv <- liftIO $ getOrInit awsEnvLazy
+    let req =
+          newPutObject s3BucketName (S3.ObjectKey name) body
+            & (putObject_metadata .~ metadata)
+    void . AWS.send awsEnv $ req
+
+  headFile GenuineS3Bucket{awsEnvLazy, s3BucketName} name = do
+    awsEnv <- liftIO $ getOrInit awsEnvLazy
+    let req = newHeadObject s3BucketName (S3.ObjectKey name)
+    handling_ S3._NoSuchKey (pure Nothing) $ Just . (^. headObjectResponse_metadata) <$> AWS.send awsEnv req
+
+  downloadFile GenuineS3Bucket{awsEnvLazy, s3BucketName} name = do
+    awsEnv <- liftIO $ getOrInit awsEnvLazy
+    let req = newGetObject s3BucketName (S3.ObjectKey name)
+    handling_ S3._NoSuchKey (pure Nothing) $ Just . (\r -> (r ^. getObjectResponse_metadata, r ^. getObjectResponse_body)) <$> AWS.send awsEnv req
+
+  deleteFile GenuineS3Bucket{awsEnvLazy, s3BucketName} name = do
+    awsEnv <- liftIO $ getOrInit awsEnvLazy
+    let req = newDeleteObject s3BucketName (S3.ObjectKey name)
+    void $ AWS.send awsEnv req
+
+data FakeS3Bucket = FakeS3Bucket
+  { fakeFiles :: IORef (Map.Map Text (Metadata, ByteString))
+  , fakeBucketName :: Text
+  }
+
+-- FIXME(jadel): upstream this to http-client
+materializeRequestBody :: HTTP.RequestBody -> IO LBS.ByteString
+materializeRequestBody = \case
+  HTTP.RequestBodyLBS lbs -> pure lbs
+  HTTP.RequestBodyBS bs -> pure (BS.fromStrict bs)
+  HTTP.RequestBodyBuilder _ b -> pure (BB.toLazyByteString b)
+  HTTP.RequestBodyStream _ g -> withPopper g
+  HTTP.RequestBodyStreamChunked g -> withPopper g
+  HTTP.RequestBodyIO io -> materializeRequestBody =<< io
+ where
+  withPopper :: HTTP.GivesPopper () -> IO LBS.ByteString
+  withPopper giver = do
+    r <- newIORef []
+    giver $ do writeIORef r <=< getChunks
+    LBS.fromChunks <$> readIORef r
+
+  getChunks :: IO ByteString -> IO [ByteString]
+  getChunks io =
+    io >>= \case
+      "" -> pure []
+      chunk -> (chunk :) <$> getChunks io
+
+instance S3BucketApi FakeS3Bucket where
+  bucketName = fakeBucketName
+
+  listItemsWithPrefix FakeS3Bucket{fakeFiles} prefix = do
+    fmap Item . filter (prefix `Text.isPrefixOf`) . Map.keys <$> readIORef fakeFiles
+
+  uploadFile FakeS3Bucket{fakeFiles} name metadata body = do
+    body <- liftIO . materializeRequestBody $ AWS.toRequestBody body
+    atomicModifyIORef' fakeFiles (\files -> (files <> Map.singleton name (metadata, BS.toStrict body), ()))
+
+  headFile FakeS3Bucket{fakeFiles} name = do
+    fmap fst . Map.lookup name <$> readIORef fakeFiles
+
+  downloadFile FakeS3Bucket{fakeFiles} name = do
+    fmap (\(meta, body) -> (meta, AWS.ResponseBody $ yield body)) . Map.lookup name <$> readIORef fakeFiles
+
+  deleteFile FakeS3Bucket{fakeFiles} name = do
+    atomicModifyIORef' fakeFiles ((,()) . (`Map.withoutKeys` (Set.singleton name)))
+
+data S3Site = S3Site
+  { s3Client :: Some S3BucketApi
+  , bucketBasePath :: Text
+  }
+
+makeRepoPath :: Text -> Repo -> Text
+makeRepoPath bucketBasePath Repo{repo_name, repo_hash} = Text.intercalate "/" [bucketBasePath, repo_name, repo_hash]
+
+dbPath :: Text -> Text
+dbPath = (<> ".tar.gz")
+
+metadataKey :: Text
+metadataKey = "glean-metadata"
+
+parseMeta :: (MonadIO m, MonadThrow m) => Repo -> Text -> m Meta
+parseMeta repo header = do
+  parsed <- either (dbError repo) pure . Thrift.deserializeJSON . Text.encodeUtf8 $ header
+  pure parsed
+
+instance Site S3Site where
+  backup S3Site{s3Client, bucketBasePath} repo meta _ttl fromPath = runResourceT $ do
+    let repoPath = makeRepoPath bucketBasePath repo
+    body <- AWS.chunkedFile AWS.defaultChunkSize fromPath
+
+    -- https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingMetadata.html#UserMetadata
+    -- Arbitrary printable Unicode characters only means that any binary
+    -- encoding would force us to base64 it, which seems like a hassle and
+    -- maybe not even more compact.
+    let meta' = HashMap.singleton metadataKey (Text.decodeUtf8 $ Thrift.serializeJSON meta)
+    _ <- uploadFile s3Client (dbPath repoPath) meta' body
+    pure $ Data (fromIntegral $ AWS.contentLength body)
+
+  delete S3Site{s3Client, bucketBasePath} repo = runResourceT $ do
+    let repoPath = makeRepoPath bucketBasePath repo
+    deleteFile s3Client (dbPath repoPath)
+
+  restore S3Site{s3Client, bucketBasePath} repo intoPath = runResourceT $ do
+    let repoPath = makeRepoPath bucketBasePath repo
+    res <- downloadFile s3Client (dbPath repoPath)
+    case res of
+      Just (meta, repoStream)
+        | Just metaJson <- HashMap.lookup metadataKey meta -> do
+            meta <- parseMeta repo metaJson
+            runConduit $ AWS.sinkBody repoStream (sinkFile intoPath)
+            pure meta
+      _ -> throwIO . Thrift.InvalidLocator $ "locator is missing either metadata or db.tar.gz" <> (Text.pack . show) repo
+
+  inspect S3Site{s3Client, bucketBasePath} repo = runResourceT $ do
+    let repoPath = makeRepoPath bucketBasePath repo
+    res <- headFile s3Client (dbPath repoPath)
+
+    case res of
+      Just meta
+        | Just metaJson <- HashMap.lookup metadataKey meta ->
+            parseMeta repo metaJson
+      _ -> throwIO . Thrift.InvalidLocator $ "locator is missing either metadata or db.tar.gz: " <> (Text.pack . show) repo
+
+  enumerate site@S3Site{s3Client, bucketBasePath} = runResourceT $ do
+    items <- listItemsWithPrefix s3Client bucketBasePath
+    let parsed = catMaybes $ map parseItemFilename items
+    forConcurrently parsed $ \repo -> (repo,) <$> liftIO (inspect site repo)
+   where
+    parseItemFilename =
+      (Text.stripPrefix (bucketBasePath <> "/") . itemPath)
+        >=> Text.stripSuffix ".tar.gz"
+        >=> splitFilename
+    -- >>> splitFilename "myrepo/123"
+    -- Just (Repo "myrepo" "123")
+    splitFilename name
+      | let (withTrailingSlash, repo_hash) = Text.breakOnEnd "/" name
+      , Just (repo_name, _slash) <- Text.unsnoc withTrailingSlash =
+          Just Repo{repo_name, repo_hash}
+    splitFilename _name = Nothing
+
+  toPath S3Site{s3Client, bucketBasePath} = bucketName s3Client <> "/" <> bucketBasePath