Make some more changesHEADmain

Maybe I'll develop on my laptop for a bit

2 files changed, 120 insertions, 59 deletions

diff --git a/src/Cache.hs b/src/Cache.hs
new file mode 100644
index 0000000..d2fd49e
--- /dev/null
+++ b/src/Cache.hs
@@ -0,0 +1,34 @@
+module Cache (Cache, new, read, write) where
+
+import qualified System.Posix.Types as T
+import qualified Data.IORef as IOR
+import qualified Data.Set as S
+import qualified Data.Map as M
+
+data WetCache k v = WC
+  { cTimeStamps :: S.Set (T.EpochTime, k)
+  , cByKey :: M.Map k (T.EpochTime, v)
+  , cDirty :: S.Set k
+  }
+
+data Cache k v = Cache
+  { cWetCache :: IOR.IORef (WetCache k v)
+  , cRead :: k -> IO v
+  , cWrite :: k -> v -> IO ()
+  , cMaxSize :: Int
+  }
+
+new :: (k -> IO v) -> (k -> v -> IO ()) -> Int -> IO Cache
+new r w size = Cache <$> IOR.newIORef CI
+  { cTimeStamps=S.empty, cByKey=M.empty, cDirty=S.empty, cRead=r, cWrite=w
+  , cMaxSize=size }
+
+readMiss :: Cache k v -> k -> IO v
+readMiss r k =
+
+read :: Cache k v -> k -> IO v
+read r k = do
+  CI { cByKey, .. } <- IOR.readIORef r
+  case M.lookup k cByKey of
+    (_, v) -> return v
+    Nothing -> readMiss r k
diff --git a/src/Lib.hs b/src/Lib.hs
index 4f4202e..6cf967f 100644
--- a/src/Lib.hs
+++ b/src/Lib.hs
@@ -37,8 +37,8 @@ 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
+import qualified Data.Heap as H
 
 newtype Handle = Handle ()
 
@@ -46,7 +46,10 @@ data BlockCopy = BlockCopy
   { bLocks :: IOR.IORef (M.Map Int RWL.RWLock)
   , bStore :: FilePath
   , bSize :: T.FileOffset
-  , bBlockSize :: T.ByteCount }
+  , 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
@@ -70,6 +73,8 @@ mkFuseGetFileStat b@BlockCopy { bSize, bBlockSize, .. } s = do
     "/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)
@@ -122,11 +127,9 @@ fuseOpen "/disk" _ _ = do
   ln "open /disk"
   return $ Right $ Handle ()
 fuseOpen "/stats" _ _ = return $ Right $ Handle ()
+fuseOpen "/flush" _ _ = 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
@@ -152,59 +155,70 @@ withWriteLocks b blocks op = do
   mapM_ RWL.releaseWrite ls
   return x
 
-ioTimeout :: Int
-ioTimeout = 30000000
+maxCacheSize :: Int
+maxCacheSize = 512
 
-timeoutRead :: BlockCopy -> Int -> IO (Maybe B.ByteString)
-timeoutRead b i = Tim.timeout ioTimeout (B.readFile (blockPath b i)) 
+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
 
-maxRetries :: Int
-maxRetries = 5
+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
 
-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
+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)
 
-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)
+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
-  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
+  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)
@@ -236,6 +250,13 @@ ln s = do
 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
@@ -265,6 +286,9 @@ mkFuseWrite store@BlockCopy { bBlockSize, .. } "/disk" _ string offset = do
   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
@@ -291,10 +315,12 @@ mkFuseReadDirectory :: BlockCopy -> FilePath
 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
+  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
@@ -308,8 +334,8 @@ fuseAccess _ _ = return E.eOK
 fuseInit :: IO ()
 fuseInit = ln "init"
 
-fuseDestroy :: IO ()
-fuseDestroy = ln "destroy"
+mkFuseDestroy :: BlockCopy -> IO ()
+mkFuseDestroy = flushAll
 
 handler :: Ex.SomeException -> IO E.Errno
 handler e = do
@@ -352,6 +378,7 @@ 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
@@ -369,7 +396,7 @@ blockCopyMain = do
               (Just bSize, aa) -> (fromIntegral bSize, aa)
       let bs = BlockCopy
             { bLocks=locks, bStore=trueStore, bSize=size
-            , bBlockSize=blockSize }
+            , bBlockSize=blockSize, bCache=cache }
       F.fuseRun prog fuseArgs
         (F.FuseOperations (mkFuseGetFileStat bs) fuseReadSymbolicLink
           fuseCreateDevice
@@ -380,5 +407,5 @@ blockCopyMain = do
           (mkFuseWrite bs) (mkFuseGetFileSystemStats bs)
           fuseFlush fuseRelease fuseSynchronizeFile fuseOpenDirectory
           (mkFuseReadDirectory bs) fuseReleaseDirectory
-          fuseSynchronizeDirectory fuseAccess fuseInit fuseDestroy)
+          fuseSynchronizeDirectory fuseAccess fuseInit (mkFuseDestroy bs))
         handler