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.IO.Error as Er
import qualified Data.Heap as H

newtype Handle = Handle ()

data BlockCopy = BlockCopy
  { bLocks :: IOR.IORef (M.Map Int RWL.RWLock)
  , bStore :: FilePath
  , bSize :: T.FileOffset
  , bBlockSize :: T.ByteCount
  , bCache
      :: IOR.IORef (H.Heap (H.Entry T.EpochTime Int), M.Map Int B.ByteString)
  }

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
    "/flush" -> return $ Right
        $ F.FileStat F.RegularFile 0o666 1 userID groupID 0 0 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 "/flush" _ _ = return $ Right $ Handle ()
fuseOpen _ _ _ = return $ Left $ E.eNOENT

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

maxCacheSize :: Int
maxCacheSize = 512

readFromDisk :: BlockCopy -> Int -> IO B.ByteString
readFromDisk BlockCopy { bStore, bBlockSize, .. } block = do
  x <- Ex.try (B.readFile (bStore FP.</> show block))
  case x of
    Left e
      | Er.isDoesNotExistError e ->
        return $ B.replicate (fromIntegral bBlockSize) 0
      | otherwise -> Ex.throw e
    Right s -> return s

writeToDisk :: BlockCopy -> Int -> B.ByteString -> IO ()
writeToDisk BlockCopy { bStore, .. } block s = do
  ln $ "Flushing " ++ show block ++ " to disk"
  B.writeFile (bStore FP.</> show block) s

cacheData :: BlockCopy -> Int -> B.ByteString -> IO (Maybe (Int, B.ByteString))
cacheData BlockCopy { bCache, .. } i s = do
  time <- Ti.epochTime
  IOR.atomicModifyIORef' bCache
    (\(lru, blockCache) ->
      let spaceCase =
            ((H.insert (H.Entry time i) lru, M.insert i s blockCache),
             Nothing)
      in case H.viewMin lru of
        Nothing -> spaceCase
        Just (H.Entry _ oldestBlock, newLru)
          | M.size blockCache >= maxCacheSize ->
            let smallerMap = M.delete oldestBlock blockCache
                evictedBlock = M.lookup oldestBlock blockCache
                newHeap = H.insert (H.Entry time i) newLru
                newMap = M.insert i s smallerMap
            in case evictedBlock of
                 Nothing -> ((newHeap, newMap), Nothing)
                 Just e -> ((newHeap, newMap), Just (oldestBlock, e))
          | otherwise -> spaceCase)

readData :: BlockCopy -> Int -> IO (Maybe B.ByteString)
readData BlockCopy { bCache, .. } block = do
  (_, cache) <- IOR.readIORef bCache
  return $ M.lookup block cache

readBlock :: BlockCopy -> Int -> IO B.ByteString
readBlock b block = do
  r <- readData b block
  case r of
    Just s -> return s
    Nothing -> do
      s <- readFromDisk b block
      e <- cacheData b block s
      case e of
        Nothing -> return s
        Just (b', s') -> do
          writeToDisk b b' s'
          return s

writeBlock :: BlockCopy -> Int -> B.ByteString -> IO ()
writeBlock b block s = do
  e <- cacheData b block s
  case e of
    Nothing -> return ()
    Just (b', s') -> writeToDisk b b' s'

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]

flushAll :: BlockCopy -> IO ()
flushAll b@BlockCopy { bCache, .. } = do
  (_, cache) <- IOR.readIORef bCache
  ln $ "flushing remaining blocks"
  withWriteLocks b (map fst (M.toList cache)) $ P.parallel_
    $ map (uncurry (writeToDisk b)) (M.toList cache)

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 b "/flush" _ s _ = do
  flushAll b
  return $ Right $ fromIntegral $ B.length s
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"
  m <- mkFuseGetFileStat b "/flush"
  case (s, k, m) of
    (Right a, Right g, Right l) -> return $ Right [("disk", a), ("stats", g), ("flush", l)]
    (Left e, _, _) -> return $ Left e
    (_, 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"

mkFuseDestroy :: BlockCopy -> IO ()
mkFuseDestroy = flushAll

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
  cache <- IOR.newIORef (H.empty, 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, bCache=cache }
      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 (mkFuseDestroy bs))
        handler