path: root/src/Lib.hs

{-# LANGUAGE NamedFieldPuns, RecordWildCards #-}
{-
 - Copyright 2018 Raymond Hogenson

 - This file is part of BlockCopy

 - BlockCopy is free software: you can redistribute it and/or modify
 - it under the terms of the GNU General Public License as published by
 - the Free Software Foundation, either version 3 of the License, or
 - (at your option) any later version.

 - BlockCopy is distributed in the hope that it will be useful,
 - but WITHOUT ANY WARRANTY; without even the implied warranty of
 - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 - GNU General Public License for more details.

 - You should have received a copy of the GNU General Public License
 - along with BlockCopy.  If not, see <https://www.gnu.org/licenses/>.
 -}

module Lib (blockCopyMain) where

import qualified Foreign.C.Error as E
import qualified System.Fuse as F
import qualified System.Posix.Types as T
import qualified Data.ByteString as B
import qualified System.FilePath as FP
import qualified Data.IORef as IOR
import qualified Data.Map as M
import qualified Control.Concurrent.ReadWriteLock as RWL
import qualified System.Posix.User as U
import qualified Control.Exception as Ex
import qualified System.Posix.Time as Ti
import qualified Data.ByteString.Char8 as C8
import qualified System.Environment as En
import qualified System.Directory as D
import qualified Control.Concurrent.ParallelIO as P
import qualified Data.List as L
import qualified Data.Maybe as Ma
import qualified System.Timeout as Tim
import qualified System.IO.Error as Er

newtype Handle = Handle ()

data BlockCopy = BlockCopy
  { bLocks :: IOR.IORef (M.Map Int RWL.RWLock)
  , bStore :: FilePath
  , bSize :: T.FileOffset
  , bBlockSize :: T.ByteCount }

statsMessage :: BlockCopy -> IO B.ByteString
statsMessage BlockCopy { bStore, bSize, bBlockSize, ..} = do
  let stats = "store: " ++ bStore ++ "\nsize: " ++ show bSize
        ++ "\nblock-size: " ++ show bBlockSize ++ "\n"
  return $ C8.pack stats

mkFuseGetFileStat :: BlockCopy -> FilePath -> IO (Either E.Errno F.FileStat)
mkFuseGetFileStat b@BlockCopy { bSize, bBlockSize, .. } s = do
  ln ("fileStat on " ++ s)
  userID <- U.getEffectiveUserID
  groupID <- U.getEffectiveGroupID
  statsSize <- fmap (fromIntegral . B.length) $ statsMessage b
  case s of
    "/" ->
      return $ Right $ F.FileStat F.Directory 0o777 1 userID groupID 0 0 0 0 0 0
    "/disk" ->
      return $ Right
        $ F.FileStat F.RegularFile 0o666 1 userID groupID 0 bSize
          (fromIntegral bSize `div` fromIntegral bBlockSize) 0 0 0
    "/stats" ->
      return $ Right
        $ F.FileStat F.RegularFile 0o666 1 userID groupID 0 statsSize 0 0 0 0
    _ -> return $ Left $ E.eNOENT

fuseReadSymbolicLink :: FilePath -> IO (Either E.Errno FilePath)
fuseReadSymbolicLink _ = return $ Left $ E.eNOENT

fuseCreateDevice :: FilePath -> F.EntryType -> T.FileMode -> T.DeviceID -> IO E.Errno
fuseCreateDevice _ _ _ _ = do
  ln "createDevice"
  return E.eNOTSUP

fuseCreateDirectory :: FilePath -> T.FileMode -> IO E.Errno
fuseCreateDirectory _ _ = do
  ln "createDirectory"
  return E.eNOTSUP

fuseRemoveLink :: FilePath -> IO E.Errno
fuseRemoveLink _ = do
  ln "removeLink"
  return E.eNOTSUP

fuseRemoveDirectory :: FilePath -> IO E.Errno
fuseRemoveDirectory _ = return E.eNOTSUP

fuseCreateSymbolicLink :: FilePath -> FilePath -> IO E.Errno
fuseCreateSymbolicLink _ _ = return E.eNOTSUP

fuseRename :: FilePath -> FilePath -> IO E.Errno
fuseRename _ _ = do
  ln "rename"
  return E.eNOTSUP

fuseCreateLink :: FilePath -> FilePath -> IO E.Errno
fuseCreateLink _ _ = return E.eNOTSUP

fuseSetFileMode :: FilePath -> T.FileMode -> IO E.Errno
fuseSetFileMode _ _ = return E.eNOTSUP

fuseSetOwnerAndGroup :: FilePath -> T.UserID -> T.GroupID -> IO E.Errno
fuseSetOwnerAndGroup _ _ _ = return E.eNOTSUP

mkFuseSetFileSize :: BlockCopy -> FilePath -> T.FileOffset -> IO E.Errno
mkFuseSetFileSize _ _ _ = return E.eOK

fuseSetFileTimes :: FilePath -> T.EpochTime -> T.EpochTime -> IO E.Errno
fuseSetFileTimes _ _ _ = return E.eNOTSUP

fuseOpen :: FilePath -> F.OpenMode -> F.OpenFileFlags
  -> IO (Either E.Errno Handle)
fuseOpen "/disk" _ _ = do
  ln "open /disk"
  return $ Right $ Handle ()
fuseOpen "/stats" _ _ = return $ Right $ Handle ()
fuseOpen _ _ _ = return $ Left $ E.eNOENT

blockPath :: BlockCopy -> Int -> FilePath
blockPath BlockCopy { bStore, .. } block = bStore FP.</> show block

getLock :: BlockCopy -> Int -> IO RWL.RWLock
getLock BlockCopy { bLocks, .. } block = do
  newLock <- RWL.new
  IOR.atomicModifyIORef'
    bLocks (\m ->
      case M.lookup block m of
        Nothing -> (M.insert block newLock m, newLock)
        Just a -> (m, a))

withReadLocks :: BlockCopy -> [Int] -> IO a -> IO a
withReadLocks b blocks op = do
  ls <- mapM (getLock b) blocks
  mapM_ RWL.acquireRead ls
  x <- op
  mapM_ RWL.releaseRead ls
  return x

withWriteLocks :: BlockCopy -> [Int] -> IO a -> IO a
withWriteLocks b blocks op = do
  ls <- mapM (getLock b) blocks
  mapM_ RWL.acquireWrite ls
  x <- op
  mapM_ RWL.releaseWrite ls
  return x

ioTimeout :: Int
ioTimeout = 30000000

timeoutRead :: BlockCopy -> Int -> IO (Maybe B.ByteString)
timeoutRead b i = Tim.timeout ioTimeout (B.readFile (blockPath b i)) 

maxRetries :: Int
maxRetries = 5

readBlock :: BlockCopy -> Int -> IO B.ByteString
readBlock b@BlockCopy { bBlockSize, .. } block = do
  let go r
        | r >= maxRetries = Ex.throwIO (userError "Timeout")
        | otherwise = do
          x <- Ex.try (timeoutRead b block)
          case x of
            Left e
              | Er.isDoesNotExistError e -> do
                ln $ show e ++ ": read of non-existent block " ++ show block
                return $ B.replicate (fromIntegral bBlockSize) 0
              | otherwise -> Ex.throwIO e
            Right Nothing -> do
              ln $ "Timeout reading block " ++ show block ++ ". Retrying "
                ++ show r
              go (r + 1)
            Right (Just s) -> return $ B.take (fromIntegral bBlockSize) s
  go 0

timeoutWrite :: BlockCopy -> Int -> B.ByteString -> IO Bool
timeoutWrite b@BlockCopy { bBlockSize, ..} block s
  | B.length s == fromIntegral bBlockSize = do
    x <- Tim.timeout ioTimeout (B.writeFile (blockPath b block) s)
    case x of
      Nothing -> return False
      Just _ -> return True
  | otherwise = do
    ln $ "Write of incorrectly sized block! length s = "
      ++ show (B.length s) ++ " but blockSize = " ++ show bBlockSize
    timeoutWrite b block (blit (B.replicate (fromIntegral bBlockSize) 0) 0 s)

writeBlock :: BlockCopy -> Int -> B.ByteString -> IO ()
writeBlock b block s = do
  let go r
        | r >= maxRetries = Ex.throw (userError "Timeout")
        | otherwise = do
          x <- timeoutWrite b block s
          case x of
            True -> return ()
            False -> do
              ln $ "Timeout writing block " ++ show block ++ ". Retrying "
                ++ show r
              go (r + 1)
  go 0

mkFuseRead :: BlockCopy -> FilePath -> Handle -> T.ByteCount -> T.FileOffset
  -> IO (Either E.Errno B.ByteString)
mkFuseRead store@BlockCopy { bBlockSize, .. } "/disk" _ count offset = do
  ln $ "Read /disk count=" ++ show count ++ " offset=" ++ show offset
  let block = fromIntegral offset `div` fromIntegral bBlockSize
  let lastBlock = fromIntegral (offset + fromIntegral count - 1)
        `div` fromIntegral bBlockSize
  let blockOffset = offset `mod` fromIntegral bBlockSize
  result <- fmap
    (Right . B.take (fromIntegral count) . B.drop (fromIntegral blockOffset))
      $ withReadLocks store [block..lastBlock] $ fmap (B.intercalate B.empty)
      $ P.parallel $ map (readBlock store) [block..lastBlock]
  ln $ "Completed read count=" ++ show count ++ " offset=" ++ show offset
  return result
mkFuseRead b "/stats" _ count offset = do
  bs <- statsMessage b
  return $ Right $ B.take (fromIntegral count) $ B.drop (fromIntegral offset) bs
mkFuseRead _ _ _ _ _ = return $ Left E.eNOENT

ln :: String -> IO ()
ln s = do
  time <- Ti.epochTime
  _ <- Ex.try
    $ appendFile "/home/ray/log"
    $ "[" ++ show time ++ "] " ++ s ++ "\n" :: IO (Either Ex.SomeException ())
  return ()

blit :: B.ByteString -> Int -> B.ByteString -> B.ByteString
blit a i x = B.intercalate B.empty [B.take i a, B.take (B.length a - i) x, B.drop (B.length x + i) a]

mkFuseWrite :: BlockCopy -> FilePath -> Handle -> B.ByteString -> T.FileOffset
  -> IO (Either E.Errno T.ByteCount)
mkFuseWrite store@BlockCopy { bBlockSize, .. } "/disk" _ string offset = do
  ln $ "Write offset=" ++ show offset ++ " size=" ++ show (B.length string)
  let block = fromIntegral offset `div` fromIntegral bBlockSize
  let lastBlock = fromIntegral (offset + fromIntegral (B.length string) - 1)
        `div` fromIntegral bBlockSize
  let blockOffset = offset `mod` fromIntegral bBlockSize
  withWriteLocks store [block..lastBlock] $ P.parallel_
    $ map (\bl ->
        let mySuffix =
              B.drop
                ((bl - 1) * fromIntegral bBlockSize + fromIntegral blockOffset)
                string
        in if bl == block && blockOffset /= 0
          then do
            currentContents <- readBlock store bl
            writeBlock store bl
              (blit currentContents (fromIntegral blockOffset) string)
          else if bl == lastBlock && B.length mySuffix < fromIntegral bBlockSize
            then do
              currentContents <- readBlock store bl
              writeBlock store bl (blit currentContents 0 mySuffix)
            else do
              writeBlock store bl (B.take (fromIntegral bBlockSize) mySuffix))
        [block..lastBlock]
  ln $ "Completed write offset=" ++ show offset ++ " length="
    ++ show (B.length string)
  return $ Right $ fromIntegral $ B.length string
mkFuseWrite _ _ _ _ _ = return $ Left E.eNOENT

mkFuseGetFileSystemStats :: BlockCopy -> String
  -> IO (Either E.Errno F.FileSystemStats)
mkFuseGetFileSystemStats BlockCopy { bBlockSize, bSize, .. } _ =
  return $ Right
    $ F.FileSystemStats (fromIntegral bBlockSize)
      (fromIntegral bSize `div` fromIntegral bBlockSize) 0 0 0 0 0

fuseFlush :: FilePath -> Handle -> IO E.Errno
fuseFlush _ _ = return E.eOK

fuseRelease :: FilePath -> Handle -> IO ()
fuseRelease _ _ = return ()

fuseSynchronizeFile :: FilePath -> F.SyncType -> IO E.Errno
fuseSynchronizeFile _ _ = return E.eOK

fuseOpenDirectory :: FilePath -> IO E.Errno
fuseOpenDirectory _ = return E.eOK

mkFuseReadDirectory :: BlockCopy -> FilePath
  -> IO (Either E.Errno [(FilePath, F.FileStat)])
mkFuseReadDirectory b _ = do
  s <- mkFuseGetFileStat b "/disk"
  k <- mkFuseGetFileStat b "/stats"
  case (s, k) of
    (Right a, Right g) -> return $ Right [("disk", a), ("stats", g)]
    (Left e, _) -> return $ Left e
    (_, Left e) -> return $ Left e

fuseReleaseDirectory :: FilePath -> IO E.Errno
fuseReleaseDirectory _ = return E.eOK

fuseSynchronizeDirectory :: FilePath -> F.SyncType -> IO E.Errno
fuseSynchronizeDirectory _ _ = return E.eOK

fuseAccess :: FilePath -> Int -> IO E.Errno
fuseAccess _ _ = return E.eOK

fuseInit :: IO ()
fuseInit = ln "init"

fuseDestroy :: IO ()
fuseDestroy = ln "destroy"

handler :: Ex.SomeException -> IO E.Errno
handler e = do
  ln $ "ERROR: " ++ show e
  return E.eFAULT

getArg :: Char -> [String] -> (Maybe String, [String])
getArg _ [] = (Nothing, [])
getArg x [l]
  | ['-', x] `L.isPrefixOf` l = (Just $ drop 2 l, [])
  | otherwise = (Nothing, [l])
getArg x (l : a : ls)
  | ['-', x] == l = (Just a, ls)
  | ['-', x] `L.isPrefixOf` l = (Just $ drop 2 l, a : ls)
  | otherwise = let (h, ll) = getArg x (a : ls) :: (Maybe String, [String])
                in (h, l : ll)

sizeMap :: M.Map String Integer
sizeMap =
  M.fromList
    [ ("c", 1), ("w", 2), ("b", 512), ("kB", 1000), ("K", 1024)
    , ("MB", 1000*1000), ("M", 1024*1024), ("xM", 1024*1024)
    , ("GB", 1000*1000*1000), ("G", 1024*1024*1024), ("T", 1024^(4 :: Integer))
    , ("P", 1024^(5 :: Integer)), ("E", 1024^(6 :: Integer)), ("Z", 1024^(7 :: Integer)), ("Y", 1024^(8 :: Integer)) ]

maybeRead :: Read a => String -> Maybe a
maybeRead = fmap fst . Ma.listToMaybe . reads

parseSize :: String -> Maybe Integer
parseSize s =
  case (M.lookup (drop (length s - 2) s) sizeMap,
        M.lookup (drop (length s - 1) s) sizeMap) of
    (Just e, _) -> (* e) <$> maybeRead (take (length s - 2) s)
    (_, Just e) -> (* e) <$> maybeRead (take (length s - 1) s)
    _ -> maybeRead s

mapFst :: (a -> b) -> (a, c) -> (b, c)
mapFst f (a, b) = (f a, b)

blockCopyMain :: IO ()
blockCopyMain = do
  locks <- IOR.newIORef M.empty
  prog <- En.getProgName
  args <- En.getArgs
  case getArg 'b' args of
    (Nothing, _) -> putStrLn "please store"
    (Just store, mArgs) -> do
      ln "main"
      trueStore <- D.canonicalizePath store
      let (size, nextArgs) =
            case mapFst (>>= parseSize) $ getArg 's' mArgs of
              (Nothing, _) -> (maxBound, mArgs)
              (Just ssize, aa) -> (fromIntegral ssize, aa)
      let (blockSize, fuseArgs) =
            case mapFst (>>= parseSize) $ getArg 'w' nextArgs of
              (Nothing, _) -> (4096, nextArgs)
              (Just bSize, aa) -> (fromIntegral bSize, aa)
      let bs = BlockCopy
            { bLocks=locks, bStore=trueStore, bSize=size
            , bBlockSize=blockSize }
      F.fuseRun prog fuseArgs
        (F.FuseOperations (mkFuseGetFileStat bs) fuseReadSymbolicLink
          fuseCreateDevice
          fuseCreateDirectory fuseRemoveLink
          fuseRemoveDirectory fuseCreateSymbolicLink fuseRename fuseCreateLink
          fuseSetFileMode fuseSetOwnerAndGroup
          (mkFuseSetFileSize bs) fuseSetFileTimes fuseOpen (mkFuseRead bs)
          (mkFuseWrite bs) (mkFuseGetFileSystemStats bs)
          fuseFlush fuseRelease fuseSynchronizeFile fuseOpenDirectory
          (mkFuseReadDirectory bs) fuseReleaseDirectory
          fuseSynchronizeDirectory fuseAccess fuseInit fuseDestroy)
        handler