diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fbff6f1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +/dist/ +/dist-newstyle/ +/cabal.project.local +/cabal.project.local~* diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 index 0000000..702a706 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE ApplicativeDo #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE QuasiQuotes #-} + +module Main (main) where + +import Data.List.NonEmpty (NonEmpty) +import Numeric.Natural (Natural) + +import qualified Algebra.Graph.Export.Dot +import qualified Data.List.NonEmpty as NonEmpty +import qualified Nix.Graph +import qualified Options.Applicative as Options + +data Options = Options + { paths :: NonEmpty FilePath + , excludeBuilt :: Bool + , excludeCached :: Bool + , maxFiles :: Natural + } + +parseOptions :: Options.Parser Options +parseOptions = do + paths <- + fmap NonEmpty.fromList $ + Options.some $ Options.strArgument (Options.metavar "PATH") + + excludeBuilt <- + Options.switch . mconcat $ + [ Options.long "exclude-built" + , Options.help . mconcat $ + [ "Reduce size of graph by excluding derivations which have already" + , "been built, and have their output path in the local Nix store." + ] + , Options.hidden + ] + + excludeCached <- + Options.switch . mconcat $ + [ Options.long "exclude-cached" + , Options.help . mconcat $ + [ "Reduce size of graph by excluding derivations which do not need to" + , "be built (because they have already been built locally) or can be" + , "substituted from another cache. Implies `--exclude-built`." + ] + , Options.hidden + ] + + maxFiles <- + Options.option Options.auto . mconcat $ + [ Options.long "max-files" + , Options.metavar "COUNT" + , Options.help "Limit number of open files" + , Options.value 100 + , Options.showDefault + , Options.hidden + ] + + pure Options{paths, excludeBuilt, excludeCached, maxFiles} + +getOptions :: IO Options +getOptions = do + let parserPrefs = + Options.prefs . mconcat $ + [ Options.showHelpOnError + , Options.multiSuffix ".." + ] + let parserInfo = Options.info (Options.helper <*> parseOptions) mempty + Options.customExecParser parserPrefs parserInfo + +main :: IO () +main = do + Options{paths, excludeBuilt, excludeCached, maxFiles} <- getOptions + + let exclude + | excludeCached = Nix.Graph.ExcludeCached + | excludeBuilt = Nix.Graph.ExcludeBuilt + | otherwise = Nix.Graph.ExcludeNothing + + let config = + Nix.Graph.Config + { Nix.Graph.exclude + , Nix.Graph.maxFiles + } + + adjacencyMap <- Nix.Graph.build config (NonEmpty.toList paths) + + putStrLn (Algebra.Graph.Export.Dot.exportAsIs adjacencyMap) diff --git a/nix-graph.cabal b/nix-graph.cabal new file mode 100644 index 0000000..d313895 --- /dev/null +++ b/nix-graph.cabal @@ -0,0 +1,66 @@ +cabal-version: 2.2 + +name: nix-graph +version: 1.0.0.0 +synopsis: Reify the Nix build graph into a Haskell graph data structure +description: Reify the Nix build graph into a Haskell graph data structure +category: Nix, Graphs +author: Arista Networks +maintainer: Arista Networks +homepage: https://github.com/awakesecurity/nix-graph +license: BSD-3-Clause +copyright: 2021 Arista Networks +tested-with: GHC == 8.8.4, GHC == 8.10.4 + +license-file: LICENSE +extra-source-files: + CHANGELOG.md + LICENSE + README.md + +common common + build-depends: base >= 4.13 && < 5.0 + default-language: Haskell2010 + ghc-options: + -Wall + -Wcompat + -Werror=incomplete-record-updates + -Werror=incomplete-uni-patterns + -Werror=missing-fields + -Werror=partial-fields + -Widentities + -Wmissing-home-modules + -Wredundant-constraints + -fshow-warning-groups + +common executable + ghc-options: + -threaded + -rtsopts + -with-rtsopts=-N + +library + import: common + hs-source-dirs: src + build-depends: + , algebraic-graphs >= 0.5 && < 0.6 + , attoparsec >= 0.13.1 && < 0.15 + , containers + , hashable + , nix-derivation >= 1.1.0 && < 2.0 + , stm >= 2.4.2 + , text + , ttrie >= 0.1.2 && < 0.2 + , unliftio >= 0.2.6 && < 0.3 + exposed-modules: + Nix.Graph + Nix.Graph.Internal + +executable nix-graph + import: common, executable + hs-source-dirs: app + main-is: Main.hs + build-depends: + , algebraic-graphs + , nix-graph + , optparse-applicative >= 0.15.0.0 diff --git a/src/Nix/Graph.hs b/src/Nix/Graph.hs new file mode 100644 index 0000000..70bc9b8 --- /dev/null +++ b/src/Nix/Graph.hs @@ -0,0 +1,7 @@ +module Nix.Graph ( + Config (..), + Exclude (..), + build, +) where + +import Nix.Graph.Internal diff --git a/src/Nix/Graph/Internal.hs b/src/Nix/Graph/Internal.hs new file mode 100644 index 0000000..0a79e82 --- /dev/null +++ b/src/Nix/Graph/Internal.hs @@ -0,0 +1,198 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module Nix.Graph.Internal where + +import Algebra.Graph.AdjacencyMap (AdjacencyMap) +import Control.Applicative ((<|>)) +import Control.Concurrent.STM.TSem (TSem) +import Control.Monad (unless, when) +import Control.Monad.IO.Class (MonadIO (..)) +import Data.Attoparsec.Text (()) +import Data.Hashable (Hashable) +import Data.Set (Set) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) +import System.Exit (ExitCode (..)) + +import qualified Algebra.Graph.AdjacencyMap as AdjacencyMap +import qualified Control.Concurrent.STM.Map as STM.Map +import qualified Control.Concurrent.STM.TSem as TSem +import qualified Data.Attoparsec.Text as Attoparsec +import qualified Data.Map as Map +import qualified Data.Set as Set +import qualified Data.Text as Text +import qualified Data.Text.IO as Text.IO +import qualified Nix.Derivation +import qualified UnliftIO.Async as Async +import qualified UnliftIO.Directory as Directory +import qualified UnliftIO.Exception as Exception +import qualified UnliftIO.IO as IO +import qualified UnliftIO.Process as Process +import qualified UnliftIO.STM as STM + +data Derivation = Derivation + { derivationPath :: FilePath + , derivationInputDrvs :: [FilePath] + , derivationBuilt :: Bool + } + deriving stock (Eq, Ord, Generic) + deriving anyclass (Hashable) + +readDerivation :: (MonadIO m, MonadFail m) => TSem -> FilePath -> m Derivation +readDerivation tSem derivationPath = do + let acquire = STM.atomically $ TSem.waitTSem tSem + let release = STM.atomically $ TSem.signalTSem tSem + + drv <- liftIO . Exception.bracket_ acquire release $ + IO.withFile derivationPath IO.ReadMode $ \handle -> do + fileContents <- Text.IO.hGetContents handle + case Attoparsec.parseOnly Nix.Derivation.parseDerivation fileContents of + Left err -> fail ("Failed to parse derivation: " <> err) + Right drv -> pure drv + + outputPath <- + case Map.lookup "out" (Nix.Derivation.outputs drv) of + Nothing -> fail "Failed to lookup output path" + Just output -> pure (Nix.Derivation.path output) + + derivationBuilt <- Directory.doesPathExist outputPath + + let derivationInputDrvs = Map.keys (Nix.Derivation.inputDrvs drv) + + pure Derivation{derivationPath, derivationBuilt, derivationInputDrvs} + +buildAdjacencyMap :: + MonadIO m => + Eq k => + Hashable k => + (k -> IO (Set k)) -> + [k] -> + m [(k, Set k)] +buildAdjacencyMap getNeighbors roots = liftIO $ do + stmMap <- STM.atomically STM.Map.empty + + let go key = do + isMember <- STM.atomically $ do + isMember <- STM.Map.member key stmMap + unless isMember $ STM.Map.insert key Set.empty stmMap + pure isMember + unless isMember $ do + neighbors <- getNeighbors key + STM.atomically $ STM.Map.insert key neighbors stmMap + Async.mapConcurrently_ go neighbors + + Async.mapConcurrently_ go roots + + STM.Map.unsafeToList stmMap + +filterUnbuilt :: (MonadIO m, MonadFail m) => [FilePath] -> m (Set FilePath) +filterUnbuilt derivationPaths = do + (exitCode, _nixStdOut, nixStdErr) <- + Process.readProcessWithExitCode + "nix-store" + ( [ "--realize" + , "--dry-run" + ] + <> derivationPaths + ) + "" + + when (exitCode /= ExitSuccess) $ do + fail ("Failed to run 'nix-store --realize --dry-run " <> unwords derivationPaths <> "':\n" <> nixStdErr) + + case Attoparsec.parseOnly willBeBuilt (Text.pack nixStdErr) of + Left parseError -> + fail ("Failed to parse output from 'nix-store --realize --dry-run ...':\n" <> parseError) + Right derivationPathsToBuild -> do + pure (Set.fromList derivationPathsToBuild) + where + willBeBuilt :: Attoparsec.Parser [FilePath] + willBeBuilt = Attoparsec.option [] $ do + willBeBuiltHeading + Attoparsec.many' nixStorePath + + -- Slightly different headings depending on Nix version + -- + -- Nix 2: https://github.com/NixOS/nix/blob/2.3.10/src/libmain/shared.cc#L45-L71 + -- (uses the same strings from Nix 2.0 to 2.3.10) + + -- + -- Nix 3: https://github.com/NixOS/nix/blob/8e758d4/src/libmain/shared.cc#L48-L86 + -- (latest as of 2020-02-25) + + willBeBuiltHeading :: Attoparsec.Parser () + willBeBuiltHeading = do + let nix2 = "these derivations will be built:" + let nix3 = + "this derivation will be built:" + <|> ("these " *> Attoparsec.decimal @Int *> " derivations will be built:") + _ <- nix2 <|> nix3 + Attoparsec.endOfLine + + nixStorePath :: Attoparsec.Parser FilePath + nixStorePath = ( "/nix/store path") $ do + _ <- " " + nixStore <- "/nix/store" + rest <- Attoparsec.takeTill Attoparsec.isEndOfLine + Attoparsec.endOfLine + pure (Text.unpack (nixStore <> rest)) + +data Config = Config + { exclude :: Exclude + , maxFiles :: Natural + } + +data Exclude + = ExcludeNothing + | ExcludeBuilt + | ExcludeCached + deriving stock (Eq) + +-- | Build graph of dependencies +build :: + MonadIO m => + -- | Configure how the graph is built + Config -> + -- | Derivations to build graph from + [FilePath] -> + m (AdjacencyMap FilePath) +build _ [] = pure (AdjacencyMap.empty) +build Config{exclude, maxFiles} roots = liftIO $ do + tSem <- STM.atomically $ TSem.newTSem (toInteger maxFiles) + + process :: [FilePath] -> IO [Derivation] <- do + case exclude of + ExcludeCached -> do + unbuiltSet <- filterUnbuilt roots + + pure \paths -> do + let unbuiltPaths = filter (`Set.member` unbuiltSet) paths + + Async.mapConcurrently (readDerivation tSem) unbuiltPaths + ExcludeBuilt -> do + pure \paths -> do + derivations <- Async.mapConcurrently (readDerivation tSem) paths + + pure (filter (not . derivationBuilt) derivations) + ExcludeNothing -> do + pure \paths -> do + Async.mapConcurrently (readDerivation tSem) paths + + let getInputDrvs :: Derivation -> IO (Set Derivation) + getInputDrvs derivation = do + fmap Set.fromList (process (derivationInputDrvs derivation)) + + rootDrvs <- process roots + + adjacencySets <- buildAdjacencyMap getInputDrvs rootDrvs + + let adjacencyMap = AdjacencyMap.fromAdjacencySets adjacencySets + + pure (AdjacencyMap.gmap derivationPath adjacencyMap)