module TahoeLAFS.Storage.Backend.Memory (
    MemoryBackend (MemoryBackend),
    memoryBackend,
) where

import Prelude hiding (
    lookup,
    map,
 )

import Network.HTTP.Types (
    ByteRanges,
 )

import Control.Exception (
    throwIO,
 )
import Data.Maybe (fromMaybe)

import Data.IORef (
    IORef,
    atomicModifyIORef',
    modifyIORef,
    newIORef,
    readIORef,
 )
import Data.Map.Strict (
    Map,
    adjust,
    filterWithKey,
    fromList,
    insert,
    keys,
    lookup,
    map,
    toList,
 )
import qualified Data.Set as Set

import TahoeLAFS.Storage.API (
    AllocateBuckets,
    AllocationResult (..),
    CBORSet (..),
    CorruptionDetails,
    Offset,
    QueryRange,
    ReadResult,
    ReadTestWriteResult (..),
    ReadTestWriteVectors (..),
    ShareData,
    ShareNumber,
    Size,
    StorageIndex,
    TestWriteVectors (..),
    Version (..),
    Version1Parameters (..),
    WriteVector (..),
    shareNumbers,
 )

import TahoeLAFS.Storage.Backend (
    Backend (..),
    ImmutableShareAlreadyWritten (ImmutableShareAlreadyWritten),
 )

type ShareStorage = Map StorageIndex (Map ShareNumber ShareData)
type BucketStorage = Map StorageIndex (Map ShareNumber (Size, ShareData))

data MemoryBackend = MemoryBackend
    { MemoryBackend -> IORef ShareStorage
immutableShares :: IORef ShareStorage -- Completely written immutable shares
    , MemoryBackend -> IORef ShareStorage
mutableShares :: IORef ShareStorage -- Completely written mutable shares
    , MemoryBackend -> IORef BucketStorage
buckets :: IORef BucketStorage -- In-progress immutable share uploads
    }

instance Show MemoryBackend where
    show :: MemoryBackend -> String
show MemoryBackend
_ = String
"<MemoryBackend>"

instance Backend MemoryBackend where
    version :: MemoryBackend -> IO Version
version MemoryBackend
backend = do
        Size
totalSize <- MemoryBackend -> IO Size
totalShareSize MemoryBackend
backend
        Version -> IO Version
forall (m :: * -> *) a. Monad m => a -> m a
return
            Version :: Version1Parameters -> ShareData -> Version
Version
                { applicationVersion :: ShareData
applicationVersion = ShareData
"(memory)"
                , parameters :: Version1Parameters
parameters =
                    Version1Parameters :: Size -> Size -> Size -> Version1Parameters
Version1Parameters
                        { maximumImmutableShareSize :: Size
maximumImmutableShareSize = Size
1024 Size -> Size -> Size
forall a. Num a => a -> a -> a
* Size
1024 Size -> Size -> Size
forall a. Num a => a -> a -> a
* Size
64
                        , maximumMutableShareSize :: Size
maximumMutableShareSize = Size
1024 Size -> Size -> Size
forall a. Num a => a -> a -> a
* Size
1024 Size -> Size -> Size
forall a. Num a => a -> a -> a
* Size
64
                        , availableSpace :: Size
availableSpace = (Size
1024 Size -> Size -> Size
forall a. Num a => a -> a -> a
* Size
1024 Size -> Size -> Size
forall a. Num a => a -> a -> a
* Size
1024) Size -> Size -> Size
forall a. Num a => a -> a -> a
- Size
totalSize
                        }
                }

    createMutableStorageIndex :: MemoryBackend -> StorageIndex -> AllocateBuckets -> IO AllocationResult
    createMutableStorageIndex :: MemoryBackend -> String -> AllocateBuckets -> IO AllocationResult
createMutableStorageIndex MemoryBackend
_backend String
_storageIndex AllocateBuckets
params =
        AllocationResult -> IO AllocationResult
forall (m :: * -> *) a. Monad m => a -> m a
return
            AllocationResult :: [ShareNumber] -> [ShareNumber] -> AllocationResult
AllocationResult
                { alreadyHave :: [ShareNumber]
alreadyHave = [ShareNumber]
forall a. Monoid a => a
mempty
                , allocated :: [ShareNumber]
allocated = AllocateBuckets -> [ShareNumber]
shareNumbers AllocateBuckets
params
                }

    getMutableShareNumbers :: MemoryBackend -> StorageIndex -> IO (CBORSet ShareNumber)
    getMutableShareNumbers :: MemoryBackend -> String -> IO (CBORSet ShareNumber)
getMutableShareNumbers MemoryBackend
backend String
storageIndex = do
        ShareStorage
shares' <- IORef ShareStorage -> IO ShareStorage
forall a. IORef a -> IO a
readIORef (IORef ShareStorage -> IO ShareStorage)
-> IORef ShareStorage -> IO ShareStorage
forall a b. (a -> b) -> a -> b
$ MemoryBackend -> IORef ShareStorage
mutableShares MemoryBackend
backend
        CBORSet ShareNumber -> IO (CBORSet ShareNumber)
forall (m :: * -> *) a. Monad m => a -> m a
return (CBORSet ShareNumber -> IO (CBORSet ShareNumber))
-> CBORSet ShareNumber -> IO (CBORSet ShareNumber)
forall a b. (a -> b) -> a -> b
$
            Set ShareNumber -> CBORSet ShareNumber
forall a. Set a -> CBORSet a
CBORSet (Set ShareNumber -> CBORSet ShareNumber)
-> ([ShareNumber] -> Set ShareNumber)
-> [ShareNumber]
-> CBORSet ShareNumber
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ShareNumber] -> Set ShareNumber
forall a. Ord a => [a] -> Set a
Set.fromList ([ShareNumber] -> CBORSet ShareNumber)
-> [ShareNumber] -> CBORSet ShareNumber
forall a b. (a -> b) -> a -> b
$
                [ShareNumber]
-> (Map ShareNumber ShareData -> [ShareNumber])
-> Maybe (Map ShareNumber ShareData)
-> [ShareNumber]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] Map ShareNumber ShareData -> [ShareNumber]
forall k a. Map k a -> [k]
keys (Maybe (Map ShareNumber ShareData) -> [ShareNumber])
-> Maybe (Map ShareNumber ShareData) -> [ShareNumber]
forall a b. (a -> b) -> a -> b
$
                    String -> ShareStorage -> Maybe (Map ShareNumber ShareData)
forall k a. Ord k => k -> Map k a -> Maybe a
lookup String
storageIndex ShareStorage
shares'

    readvAndTestvAndWritev :: MemoryBackend -> StorageIndex -> ReadTestWriteVectors -> IO ReadTestWriteResult
    readvAndTestvAndWritev :: MemoryBackend
-> String -> ReadTestWriteVectors -> IO ReadTestWriteResult
readvAndTestvAndWritev
        MemoryBackend
backend
        String
storageIndex
        (ReadTestWriteVectors SlotSecrets
_secrets Map ShareNumber TestWriteVectors
testWritev [ReadVector]
_readv) = do
            -- TODO implement readv and testv parts.  implement secrets part.
            let shares :: IORef ShareStorage
shares = MemoryBackend -> IORef ShareStorage
mutableShares MemoryBackend
backend
            IORef ShareStorage -> (ShareStorage -> ShareStorage) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef ShareStorage
shares ((ShareStorage -> ShareStorage) -> IO ())
-> (ShareStorage -> ShareStorage) -> IO ()
forall a b. (a -> b) -> a -> b
$ String
-> [(ShareNumber, ShareData)] -> ShareStorage -> ShareStorage
addShares String
storageIndex (Map ShareNumber TestWriteVectors -> [(ShareNumber, ShareData)]
shares' Map ShareNumber TestWriteVectors
testWritev)
            ReadTestWriteResult -> IO ReadTestWriteResult
forall (m :: * -> *) a. Monad m => a -> m a
return
                ReadTestWriteResult :: Bool -> ReadResult -> ReadTestWriteResult
ReadTestWriteResult
                    { success :: Bool
success = Bool
True
                    , readData :: ReadResult
readData = ReadResult
forall a. Monoid a => a
mempty
                    }
          where
            shares' ::
                Map ShareNumber TestWriteVectors ->
                [(ShareNumber, ShareData)]
            shares' :: Map ShareNumber TestWriteVectors -> [(ShareNumber, ShareData)]
shares' Map ShareNumber TestWriteVectors
testWritevs =
                [ (ShareNumber
shareNumber, WriteVector -> ShareData
shareData WriteVector
writev)
                | (ShareNumber
shareNumber, TestWriteVectors
testWritev') <- Map ShareNumber TestWriteVectors
-> [(ShareNumber, TestWriteVectors)]
forall k a. Map k a -> [(k, a)]
toList Map ShareNumber TestWriteVectors
testWritevs
                , WriteVector
writev <- TestWriteVectors -> [WriteVector]
write TestWriteVectors
testWritev'
                ]

    createImmutableStorageIndex :: MemoryBackend -> StorageIndex -> AllocateBuckets -> IO AllocationResult
    createImmutableStorageIndex :: MemoryBackend -> String -> AllocateBuckets -> IO AllocationResult
createImmutableStorageIndex MemoryBackend
_backend String
_idx AllocateBuckets
params =
        AllocationResult -> IO AllocationResult
forall (m :: * -> *) a. Monad m => a -> m a
return
            AllocationResult :: [ShareNumber] -> [ShareNumber] -> AllocationResult
AllocationResult
                { alreadyHave :: [ShareNumber]
alreadyHave = [ShareNumber]
forall a. Monoid a => a
mempty
                , allocated :: [ShareNumber]
allocated = AllocateBuckets -> [ShareNumber]
shareNumbers AllocateBuckets
params
                }

    writeImmutableShare :: MemoryBackend -> StorageIndex -> ShareNumber -> ShareData -> Maybe ByteRanges -> IO ()
    writeImmutableShare :: MemoryBackend
-> String -> ShareNumber -> ShareData -> QueryRange -> IO ()
writeImmutableShare MemoryBackend
backend String
storageIndex ShareNumber
shareNumber ShareData
shareData QueryRange
Nothing = do
        -- shares <- readIORef (immutableShares backend) -- XXX uh, is this right?!
        Bool
changed <- IORef ShareStorage
-> (ShareStorage -> (ShareStorage, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (MemoryBackend -> IORef ShareStorage
immutableShares MemoryBackend
backend) ((ShareStorage -> (ShareStorage, Bool)) -> IO Bool)
-> (ShareStorage -> (ShareStorage, Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$
            \ShareStorage
shares ->
                case String -> ShareStorage -> Maybe (Map ShareNumber ShareData)
forall k a. Ord k => k -> Map k a -> Maybe a
lookup String
storageIndex ShareStorage
shares Maybe (Map ShareNumber ShareData)
-> (Map ShareNumber ShareData -> Maybe ShareData)
-> Maybe ShareData
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ShareNumber -> Map ShareNumber ShareData -> Maybe ShareData
forall k a. Ord k => k -> Map k a -> Maybe a
lookup ShareNumber
shareNumber of
                    Just ShareData
_ ->
                        -- It is not allowed to write new data for an immutable share that
                        -- has already been written.
                        (ShareStorage
shares, Bool
False)
                    Maybe ShareData
Nothing ->
                        (String
-> [(ShareNumber, ShareData)] -> ShareStorage -> ShareStorage
addShares String
storageIndex [(ShareNumber
shareNumber, ShareData
shareData)] ShareStorage
shares, Bool
True)
        if Bool
changed
            then () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            else ImmutableShareAlreadyWritten -> IO ()
forall e a. Exception e => e -> IO a
throwIO ImmutableShareAlreadyWritten
ImmutableShareAlreadyWritten
    writeImmutableShare MemoryBackend
_ String
_ ShareNumber
_ ShareData
_ QueryRange
_ = String -> IO ()
forall a. HasCallStack => String -> a
error String
"writeImmutableShare got bad input"

    adviseCorruptImmutableShare :: MemoryBackend -> StorageIndex -> ShareNumber -> CorruptionDetails -> IO ()
    adviseCorruptImmutableShare :: MemoryBackend
-> String -> ShareNumber -> CorruptionDetails -> IO ()
adviseCorruptImmutableShare MemoryBackend
_backend String
_ ShareNumber
_ CorruptionDetails
_ =
        () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
forall a. Monoid a => a
mempty

    getImmutableShareNumbers :: MemoryBackend -> StorageIndex -> IO (CBORSet ShareNumber)
    getImmutableShareNumbers :: MemoryBackend -> String -> IO (CBORSet ShareNumber)
getImmutableShareNumbers MemoryBackend
backend String
storageIndex = do
        ShareStorage
shares' <- IORef ShareStorage -> IO ShareStorage
forall a. IORef a -> IO a
readIORef (IORef ShareStorage -> IO ShareStorage)
-> IORef ShareStorage -> IO ShareStorage
forall a b. (a -> b) -> a -> b
$ MemoryBackend -> IORef ShareStorage
immutableShares MemoryBackend
backend
        CBORSet ShareNumber -> IO (CBORSet ShareNumber)
forall (m :: * -> *) a. Monad m => a -> m a
return (CBORSet ShareNumber -> IO (CBORSet ShareNumber))
-> CBORSet ShareNumber -> IO (CBORSet ShareNumber)
forall a b. (a -> b) -> a -> b
$ Set ShareNumber -> CBORSet ShareNumber
forall a. Set a -> CBORSet a
CBORSet (Set ShareNumber -> CBORSet ShareNumber)
-> ([ShareNumber] -> Set ShareNumber)
-> [ShareNumber]
-> CBORSet ShareNumber
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ShareNumber] -> Set ShareNumber
forall a. Ord a => [a] -> Set a
Set.fromList ([ShareNumber] -> CBORSet ShareNumber)
-> [ShareNumber] -> CBORSet ShareNumber
forall a b. (a -> b) -> a -> b
$ [ShareNumber]
-> (Map ShareNumber ShareData -> [ShareNumber])
-> Maybe (Map ShareNumber ShareData)
-> [ShareNumber]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] Map ShareNumber ShareData -> [ShareNumber]
forall k a. Map k a -> [k]
keys (Maybe (Map ShareNumber ShareData) -> [ShareNumber])
-> Maybe (Map ShareNumber ShareData) -> [ShareNumber]
forall a b. (a -> b) -> a -> b
$ String -> ShareStorage -> Maybe (Map ShareNumber ShareData)
forall k a. Ord k => k -> Map k a -> Maybe a
lookup String
storageIndex ShareStorage
shares'

    readImmutableShare :: MemoryBackend -> StorageIndex -> ShareNumber -> QueryRange -> IO ShareData
    readImmutableShare :: MemoryBackend
-> String -> ShareNumber -> QueryRange -> IO ShareData
readImmutableShare MemoryBackend
backend String
storageIndex ShareNumber
shareNum QueryRange
_qr = do
        ShareStorage
shares' <- IORef ShareStorage -> IO ShareStorage
forall a. IORef a -> IO a
readIORef (IORef ShareStorage -> IO ShareStorage)
-> IORef ShareStorage -> IO ShareStorage
forall a b. (a -> b) -> a -> b
$ MemoryBackend -> IORef ShareStorage
immutableShares MemoryBackend
backend
        let result :: Maybe ShareData
result = case String -> ShareStorage -> Maybe (Map ShareNumber ShareData)
forall k a. Ord k => k -> Map k a -> Maybe a
lookup String
storageIndex ShareStorage
shares' of
                Maybe (Map ShareNumber ShareData)
Nothing -> Maybe ShareData
forall a. Monoid a => a
mempty
                Just Map ShareNumber ShareData
shares'' -> ShareNumber -> Map ShareNumber ShareData -> Maybe ShareData
forall k a. Ord k => k -> Map k a -> Maybe a
lookup ShareNumber
shareNum Map ShareNumber ShareData
shares''
        ShareData -> IO ShareData
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ShareData -> IO ShareData) -> ShareData -> IO ShareData
forall a b. (a -> b) -> a -> b
$ ShareData -> Maybe ShareData -> ShareData
forall a. a -> Maybe a -> a
fromMaybe ShareData
forall a. Monoid a => a
mempty Maybe ShareData
result

totalShareSize :: MemoryBackend -> IO Size
totalShareSize :: MemoryBackend -> IO Size
totalShareSize MemoryBackend
backend = do
    ShareStorage
imm <- IORef ShareStorage -> IO ShareStorage
forall a. IORef a -> IO a
readIORef (IORef ShareStorage -> IO ShareStorage)
-> IORef ShareStorage -> IO ShareStorage
forall a b. (a -> b) -> a -> b
$ MemoryBackend -> IORef ShareStorage
immutableShares MemoryBackend
backend
    ShareStorage
mut <- IORef ShareStorage -> IO ShareStorage
forall a. IORef a -> IO a
readIORef (IORef ShareStorage -> IO ShareStorage)
-> IORef ShareStorage -> IO ShareStorage
forall a b. (a -> b) -> a -> b
$ MemoryBackend -> IORef ShareStorage
mutableShares MemoryBackend
backend
    let immSize :: Int
immSize = Map String Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Map String Int -> Int) -> Map String Int -> Int
forall a b. (a -> b) -> a -> b
$ (Map ShareNumber ShareData -> Int)
-> ShareStorage -> Map String Int
forall a b k. (a -> b) -> Map k a -> Map k b
map Map ShareNumber ShareData -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShareStorage
imm
    let mutSize :: Int
mutSize = Map String Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Map String Int -> Int) -> Map String Int -> Int
forall a b. (a -> b) -> a -> b
$ (Map ShareNumber ShareData -> Int)
-> ShareStorage -> Map String Int
forall a b k. (a -> b) -> Map k a -> Map k b
map Map ShareNumber ShareData -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShareStorage
mut
    Size -> IO Size
forall (m :: * -> *) a. Monad m => a -> m a
return (Size -> IO Size) -> Size -> IO Size
forall a b. (a -> b) -> a -> b
$ Int -> Size
forall a. Integral a => a -> Size
toInteger (Int -> Size) -> Int -> Size
forall a b. (a -> b) -> a -> b
$ Int
immSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
mutSize

addShares :: StorageIndex -> [(ShareNumber, ShareData)] -> ShareStorage -> ShareStorage
addShares :: String
-> [(ShareNumber, ShareData)] -> ShareStorage -> ShareStorage
addShares String
_storageIndex [] ShareStorage
shareStorage = ShareStorage
shareStorage
addShares String
storageIndex ((ShareNumber
shareNumber, ShareData
shareData) : [(ShareNumber, ShareData)]
rest) ShareStorage
shareStorage =
    let added :: ShareStorage
added = case String -> ShareStorage -> Maybe (Map ShareNumber ShareData)
forall k a. Ord k => k -> Map k a -> Maybe a
lookup String
storageIndex ShareStorage
shareStorage of
            Maybe (Map ShareNumber ShareData)
Nothing ->
                String -> Map ShareNumber ShareData -> ShareStorage -> ShareStorage
forall k a. Ord k => k -> a -> Map k a -> Map k a
insert String
storageIndex ([(ShareNumber, ShareData)] -> Map ShareNumber ShareData
forall k a. Ord k => [(k, a)] -> Map k a
fromList [(ShareNumber
shareNumber, ShareData
shareData)]) ShareStorage
shareStorage
            Just Map ShareNumber ShareData
_shares ->
                (Map ShareNumber ShareData -> Map ShareNumber ShareData)
-> String -> ShareStorage -> ShareStorage
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
adjust Map ShareNumber ShareData -> Map ShareNumber ShareData
addShare' String
storageIndex ShareStorage
shareStorage
              where
                addShare' :: Map ShareNumber ShareData -> Map ShareNumber ShareData
addShare' = ShareNumber
-> ShareData
-> Map ShareNumber ShareData
-> Map ShareNumber ShareData
forall k a. Ord k => k -> a -> Map k a -> Map k a
insert ShareNumber
shareNumber ShareData
shareData
     in String
-> [(ShareNumber, ShareData)] -> ShareStorage -> ShareStorage
addShares String
storageIndex [(ShareNumber, ShareData)]
rest ShareStorage
added

memoryBackend :: IO MemoryBackend
memoryBackend :: IO MemoryBackend
memoryBackend = do
    IORef ShareStorage
immutableShares <- ShareStorage -> IO (IORef ShareStorage)
forall a. a -> IO (IORef a)
newIORef ShareStorage
forall a. Monoid a => a
mempty
    IORef ShareStorage
mutableShares <- ShareStorage -> IO (IORef ShareStorage)
forall a. a -> IO (IORef a)
newIORef ShareStorage
forall a. Monoid a => a
mempty
    IORef BucketStorage
buckets <- BucketStorage -> IO (IORef BucketStorage)
forall a. a -> IO (IORef a)
newIORef BucketStorage
forall a. Monoid a => a
mempty
    MemoryBackend -> IO MemoryBackend
forall (m :: * -> *) a. Monad m => a -> m a
return (MemoryBackend -> IO MemoryBackend)
-> MemoryBackend -> IO MemoryBackend
forall a b. (a -> b) -> a -> b
$ IORef ShareStorage
-> IORef ShareStorage -> IORef BucketStorage -> MemoryBackend
MemoryBackend IORef ShareStorage
immutableShares IORef ShareStorage
mutableShares IORef BucketStorage
buckets