module Cook.State.Manager
( StateManager
, createStateManager, markUsingImage
, isImageKnown
, garbageCollectImages
)
where
import Cook.Util
import Cook.State.Model
import Cook.Types
import Control.Applicative
import Control.Concurrent.STM
import Control.Monad.Logger hiding (logInfo)
import Control.Monad.State
import Control.Monad.Trans.Resource
import Data.Maybe
import Data.SafeCopy (safeGet, safePut)
import Data.Serialize.Get (runGet)
import Data.Serialize.Put (runPut)
import Data.Time.Clock
import Database.Persist.Sqlite
import System.Directory
import System.FilePath
import qualified Data.ByteString as BS
import qualified Data.HashMap.Strict as HM
import qualified Data.Graph as G
import qualified Data.Graph.NodeManager as NM
import qualified Data.Graph.Persistence as GP
import qualified Data.Text as T
import qualified Data.Traversable as T
import qualified Data.Vector.Unboxed as VU
data StateManager
= StateManager
{ sm_runSql :: forall a. SqlPersistM a -> IO a
, sm_graph :: TVar G.Graph
, sm_nodeManager :: TVar (NM.NodeManager DockerImage)
, sm_persistGraph :: IO ()
}
createStateManager :: FilePath -> IO StateManager
createStateManager stateDirectory =
do pool <- createSqlitePool (T.pack sqlLoc) 5
let runSql = runResourceT . runNoLoggingT . ((flip runSqlPool) pool)
runSql (runMigration migrateState)
exists <- doesFileExist graphFile
(g, nm) <-
case exists of
True ->
do bs <- BS.readFile graphFile
case runGet safeGet bs of
Left errMsg ->
error ("Failed to read " ++ graphFile ++ ": " ++ errMsg)
Right persistedGraph ->
let (lnm, lg) = GP.loadGraph persistedGraph
in (,) <$> newTVarIO lg <*> newTVarIO lnm
False ->
(,) <$> newTVarIO G.empty <*> newTVarIO NM.emptyNodeManager
return $ StateManager
{ sm_runSql = runSql
, sm_graph = g
, sm_nodeManager = nm
, sm_persistGraph = persistGraph g nm
}
where
sqlLoc = stateDirectory </> "database.db"
graphFile = stateDirectory </> "graph.bin"
persistGraph gVar nmVar =
do nm <- atomically $ readTVar nmVar
g <- atomically $ readTVar gVar
BS.writeFile graphFile $ runPut (safePut $ GP.persistGraph nm g)
data GCState
= GCState
{ gc_canTrash :: !(HM.HashMap NM.Node (Bool, DockerImage))
, gc_trashCount :: !Int
, gc_cache :: !(HM.HashMap DockerImage DbDockerImage)
} deriving (Show)
data SweepState
= SweepState
{ ss_graph :: !(G.Graph)
, ss_removedImages :: [DockerImage]
} deriving (Show)
garbageCollectImages :: StateManager
-> (DbDockerImage -> Bool)
-> (DockerImage -> IO Bool)
-> IO [DockerImage]
garbageCollectImages (StateManager{..}) deletePred deleteFun =
do graph <- atomically $ readTVar sm_graph
nodeManager <- atomically $ readTVar sm_nodeManager
let graphLeafs = filter (\n -> VU.null $ G.parents graph n) (G.nodes graph)
logInfo ("Found " ++ (show (length graphLeafs)) ++ " toplevel image(s). Starting mark and sweep")
gcState <- execStateT (mapM (markNode graph nodeManager) graphLeafs) (GCState HM.empty 0 HM.empty)
logInfo ("Found " ++ (show (gc_trashCount gcState)) ++ " deletable image(s).")
sweepState <- execStateT (sweepNodes gcState nodeManager graphLeafs) (SweepState graph [])
return (reverse $ ss_removedImages sweepState)
where
sweepNodes :: GCState -> NM.NodeManager DockerImage -> [NM.Node] -> StateT SweepState IO ()
sweepNodes gcState nodeManager [] =
do currentGraph <- gets ss_graph
alreadyRemoved <- gets ss_removedImages
let nextLeafs =
filter (\n ->
let isLeaf = VU.null $ G.parents currentGraph n
trashable =
case HM.lookup n (gc_canTrash gcState) of
Just (True, imageName) ->
not $ imageName `elem` alreadyRemoved
_ -> False
in isLeaf && trashable
) (G.nodes currentGraph)
if length nextLeafs == 0
then return ()
else sweepNodes gcState nodeManager nextLeafs
sweepNodes gcState nodeManager (node:rest) =
do currentGraph <- gets ss_graph
case HM.lookup node (gc_canTrash gcState) of
Just (True, imageName) ->
do let g' = G.removeNode node currentGraph
nm' = NM.removeNodeHandle node nodeManager
liftIO $
do deleteOk <- deleteFun imageName
when (not deleteOk) $
error ("Failed to delete " ++ (T.unpack $ unDockerImage imageName) ++ ". Aborting!")
atomically $
do writeTVar sm_nodeManager nm'
writeTVar sm_graph g'
sm_persistGraph
sm_runSql $ deleteBy (UniqueGraphNodeId node)
modify $ \st ->
st
{ ss_graph = g'
, ss_removedImages = (imageName : ss_removedImages st)
}
sweepNodes gcState nm' rest
_ ->
sweepNodes gcState nodeManager rest
markNode :: G.Graph -> NM.NodeManager DockerImage -> NM.Node -> StateT GCState IO Bool
markNode graph nodeManager node =
do let dockerImage =
case NM.lookupNode node nodeManager of
Nothing ->
error ("dockercook inconsistency: found node in image graph without any dockerimage!")
Just d -> d
children = G.children graph node
trashCheck (Just False) _ = return (Just False)
trashCheck x [] = return x
trashCheck _ (x:xs) =
do trashState <- gets gc_canTrash
r <-
case HM.lookup x trashState of
Nothing -> markNode graph nodeManager x
Just (v, _) -> return v
trashCheck (Just r) xs
markAs x =
do modify $ \s ->
s { gc_canTrash = HM.insert node (x, dockerImage) (gc_canTrash s)
, gc_trashCount = (if x then 1 else 0) + (gc_trashCount s)
}
return x
mCanTrash <- trashCheck Nothing $ VU.toList children
case mCanTrash of
Just False ->
markAs False
_ ->
do cacheMap <- gets gc_cache
dbInfo <-
case HM.lookup dockerImage cacheMap of
Nothing ->
do mDbInfo <- liftIO $ sm_runSql $ getBy (UniqueGraphNodeId node)
case mDbInfo of
Nothing ->
error ("dockercook inconsistency: found node in image graph without any meta data!")
Just someInfo ->
do let e = entityVal someInfo
modify $ \s ->
s { gc_cache = HM.insert dockerImage e (gc_cache s) }
return e
Just info ->
return info
markAs $ deletePred dbInfo
isImageKnown :: StateManager -> DockerImage -> IO Bool
isImageKnown (StateManager{..}) (DockerImage imageName) =
do x <- sm_runSql $ getBy (UniqueDbDockerImage imageName)
return (isJust x)
markUsingImage :: StateManager -> DockerImage -> Maybe DockerImage -> IO ()
markUsingImage (StateManager{..}) img@(DockerImage imageName) mParentImage =
do parentEntity <-
T.mapM (sm_runSql . findParentImage) mParentImage
mImageEntity <- sm_runSql $ getBy (UniqueDbDockerImage imageName)
now <- getCurrentTime
case mImageEntity of
Nothing ->
do newNodeId <-
atomically $
do nm <- readTVar sm_nodeManager
(nodeId, newNm) <-
runStateT (NM.getNodeHandle img) nm
writeTVar sm_nodeManager newNm
modifyTVar sm_graph $ \g ->
case parentEntity of
Nothing ->
G.addNode nodeId g
Just pe ->
G.addEdge nodeId (dbDockerImageNodeId $ entityVal pe) g
return nodeId
_ <- sm_runSql $ insert $ DbDockerImage imageName now now 1 newNodeId
sm_persistGraph
return ()
Just imageEntity ->
sm_runSql $ update (entityKey imageEntity) [ DbDockerImageUsageCount +=. 1
, DbDockerImageLastUsed =. now
]
where
findParentImage (DockerImage parentImageName) =
do mParentEntity <- getBy (UniqueDbDockerImage parentImageName)
case mParentEntity of
Just entity -> return entity
Nothing ->
error ("dockercook inconsistency: parent image "
++ show parentImageName ++ " referenced but unknown to database.")