Allow concurrent indexing in glean index cpp-cmake

glean.cabal.in

@@ -588,6 +588,16 @@ library client-cpp
     build-depends:
         glean:rts
 
+library interprocess
+    import: fb-haskell, deps
+    visibility: public
+    hs-source-dirs: glean/interprocess/hs
+    exposed-modules:
+        Glean.Interprocess.Counters
+        Glean.Interprocess.Worklist
+    build-depends:
+        glean:client-cpp
+
 library schema
     import: fb-haskell, fb-cpp, deps
     visibility: public
@@ -993,6 +1003,7 @@ library indexers
         glean:client-hs-local,
         glean:db,
         glean:handler,
+        glean:interprocess,
         glean:lib,
         glean:lsif,
         glean:stubs,

glean/lang/clang/Glean/Indexer/Cpp.hs

@@ -11,10 +11,13 @@ module Glean.Indexer.Cpp
   ( indexerWith, indexer, indexerNoDeriv, Clang(..)
   , findExecutableRecursive ) where
 
+import Control.Concurrent (threadDelay)
 import Control.Concurrent.Async
+import Control.Exception
 import Control.Monad
 import Data.Proxy
 import Options.Applicative
+import qualified System.Console.ANSI as ANSI
 import System.Directory
 import System.Environment
 import System.Exit
@@ -23,6 +26,7 @@ import System.IO
 import System.Process
 import Thrift.Protocol (deserializeGen)
 import Thrift.Protocol.Compact (Compact)
+import Util.List (chunk)
 
 import Facebook.Fb303
 import Facebook.Service
@@ -32,6 +36,7 @@ import Glean.Indexer
 import Glean.LocalOrRemote ( BackendKind(..),
   LocalOrRemote(..), serializeInventory )
 import Glean.Util.Service
+import qualified Glean.Interprocess.Worklist as Worklist
 
 import qualified Data.ByteString as BS
 import qualified Glean.Handler as GleanHandler
@@ -42,7 +47,9 @@ data Clang = Clang
   , clangDeriveBin    :: Maybe FilePath -- ^ path to @clang-derive@ binary
   , clangCompileDBDir :: Maybe FilePath
       -- ^ (optional) path to pre-existing @compile_commands.json@
+  , clangJobs         :: Int -- ^ number of indexers to run concurrently
   , clangVerbose      :: Bool -- ^ display debugging information
+  , clangProgress     :: Bool -- ^ display indexing progress
   } deriving Show
 
 options :: Parser Clang
@@ -56,10 +63,18 @@ options = do
   clangCompileDBDir <- optional $ strOption $
     long "cdb" <>
     help "path to a directory containing an existing compile_commands.json file"
+  clangJobs <- option auto $
+    short 'j' <>
+    long "jobs" <>
+    value 1 <>
+    help "run N indexers in parallel"
   clangVerbose <- switch $
     short 'v' <>
     long "verbose" <>
     help "Enable verbose logging from subprocesses"
+  clangProgress <- switch $
+    long "progress" <>
+    help "Display indexing progress even in verbose mode"
   return Clang{..}
 
 -- | Standard indexer, that also runs the deriver
@@ -77,99 +92,192 @@ indexerWith deriveToo = Indexer {
   indexerShortName = "cpp-cmake",
   indexerDescription = "Index C++ code with CMake (via Clang)",
   indexerOptParser = options,
-  indexerRun = \Clang{..} backend repo IndexerParams{..} -> do
+  indexerRun = \clang@Clang{..} backend repo IndexerParams{..} -> do
     -- indexing
     let tmpDir        = indexerOutput
         inventoryFile = tmpDir </> "inventory.data"
-        indexerData   = tmpDir </> "indexer.data"
     generateInventory backend repo inventoryFile
     compileDBDir <-
       case clangCompileDBDir of
         Nothing  -> cmake clangVerbose indexerRoot tmpDir >> return tmpDir
         Just dir -> return dir
-    index clangVerbose clangIndexBin inventoryFile
-      indexerRoot compileDBDir indexerData
+    indexerData <-
+      index clang inventoryFile indexerRoot compileDBDir indexerOutput
     writeToDB backend repo indexerData
 
     -- deriving
     when deriveToo $
       derive clangVerbose clangDeriveBin backend repo
   }
 
-  where generateInventory backend repo outFile =
-          serializeInventory backend repo >>= BS.writeFile outFile
+  where
+    generateInventory backend repo outFile =
+      serializeInventory backend repo >>= BS.writeFile outFile
+
+    cmake verbose srcDir tmpDir = withExe "cmake" Nothing $ \cmakeBin ->
+      spawnAndConcurrentLog verbose cmakeBin
+        [ "-DCMAKE_EXPORT_COMPILE_COMMANDS=1"
+        , "-S", srcDir
+        , "-B", tmpDir
+        ]
 
-        cmake verbose srcDir tmpDir = withExe "cmake" Nothing $ \cmakeBin ->
-          spawnAndConcurrentLog verbose cmakeBin
-            [ "-DCMAKE_EXPORT_COMPILE_COMMANDS=1"
-            , "-S", srcDir
-            , "-B", tmpDir
+    index Clang{..} inventory srcDir buildDir tmpDir =
+      withExe "clang-index" clangIndexBin $ \clangIndex -> do
+      let args =
+            [ "-cdb_dir", buildDir
+            , "-cdb_target", "all"
+            , "-root", srcDir
+            , "--inventory", inventory
+            , "-logtostderr"
             ]
 
-        index verbose indexBin inventory srcDir buildDir outFile =
-          withExe "clang-index" indexBin $ \clangIndex -> do
-            let args = [ "-cdb_dir", buildDir
-                       , "-cdb_target", "all"
-                       , "-root", srcDir
-                       , "-dump", outFile
-                       , "--inventory", inventory
-                       , "-logtostderr"
-                       ]
-            spawnAndConcurrentLog verbose clangIndex args
-
-        writeToDB backend repo dataFile = do
-          dat <- BS.readFile dataFile
-          case deserializeGen (Proxy :: Proxy Compact) dat of
-            Left parseError -> error parseError
-            Right batch     -> sendBatch backend repo batch
-
-        derive verbose deriveBin backend repo =
-          withExe "clang-derive" deriveBin $ \clangDerive -> do
-            let go service = spawnAndConcurrentLog verbose clangDerive
-                  [ "--repo", showRepo repo
-                  , "--service", service
-                  ]
-            case backendKind backend of
-              BackendEnv env -> do
-                fb303 <- newFb303 "gleandriver"
-                let state = GleanHandler.State fb303 env
-                withBackgroundFacebookService
-                  (GleanHandler.fb303State state)
-                  (GleanHandler.handler state)
-                  CppServer.defaultOptions
-                  $ \server ->
-                    go ("localhost:" <> show (CppServer.serverPort server))
-              BackendThrift thrift -> do
-                let clientConfig = thriftBackendClientConfig thrift
-                go $ serviceToString (clientConfig_serv clientConfig)
+      -- get the total number of source files
+      sources <- do
+        let pargs = args ++ ["--print_sources_count"]
+        s <- readProcess clangIndex pargs ""
+        case reads s of
+          [(sources,"")] -> return sources
+          _ -> error $ unwords (clangIndex:pargs)
+            ++ " produced unexpect output \"" ++ s ++ "\""
+
+      case sources of
+        0 -> do
+          -- TODO: should this be an error?
+          putStrLn "No source files to index"
+          return []
+        _ ->
+          -- set up worklist
+          let ranges =
+                map (\(i,n) -> Worklist.Range i (i+n)) $ chunk clangJobs sources
+              !workers = length ranges
+          in
+          Worklist.withTemp ranges $ \wfile worklist ->
+
+          -- progress and logging
+          (if clangProgress || not clangVerbose
+            then withProgress worklist clangJobs sources
+            else id) $
+          withLog clangVerbose (void . evaluate . length) $ \stream -> do
+
+          -- run workers
+          let dataFile i = tmpDir </> "indexer-" <> show i <> ".data"
+              workerargs i = args ++
+                [ "-dump", dataFile i
+                , "--work_file", wfile
+                , "--worker_index", show i
+                , "--worker_count", show workers
+                ]
+          forConcurrently_ [0 .. workers-1] $ \i -> bracket
+            -- createProcess_ because we don't want the stdout/stderr handles
+            -- to be closed
+            (createProcess_
+              "Cpp.index"
+              (proc clangIndex $ workerargs i)
+                {std_out = stream, std_err = stream})
+            cleanupProcess
+            $ \(_, _, _, ph) -> do
+              ex <- waitForProcess ph
+              case ex of
+                ExitSuccess -> return ()
+                ExitFailure i -> error $ unwords (clangIndex:workerargs i)
+                  ++ " returned exit code " ++ show i
+
+          -- return data file names
+          return $ map dataFile [0 .. workers-1]
+
+    writeToDB backend repo = mapM_ $ \dataFile -> do
+      dat <- BS.readFile dataFile
+      case deserializeGen (Proxy :: Proxy Compact) dat of
+        Left parseError -> error parseError
+        Right batch     -> sendBatch backend repo batch
+
+    derive verbose deriveBin backend repo =
+      withExe "clang-derive" deriveBin $ \clangDerive -> do
+        let go service = spawnAndConcurrentLog verbose clangDerive
+              [ "--repo", showRepo repo
+              , "--service", service
+              ]
+        case backendKind backend of
+          BackendEnv env -> do
+            fb303 <- newFb303 "gleandriver"
+            let state = GleanHandler.State fb303 env
+            withBackgroundFacebookService
+              (GleanHandler.fb303State state)
+              (GleanHandler.handler state)
+              CppServer.defaultOptions
+              $ \server ->
+                go ("localhost:" <> show (CppServer.serverPort server))
+          BackendThrift thrift -> do
+            let clientConfig = thriftBackendClientConfig thrift
+            go $ serviceToString (clientConfig_serv clientConfig)
+
+withProgress ::  Worklist.Worklist -> Int -> Int -> IO a -> IO a
+withProgress worklist jobs total action = do
+  terminal <- ANSI.hSupportsANSI stdout
+  flush $ start terminal
+  x <- withAsync (showProgress terminal 0) $ const action
+  flush $ finish terminal
+  return x
+  where
+    flush f = f >> hFlush stdout
+
+    tmsg n =
+      unwords ["Indexed", show (total - n), "of", show total, "source files"]
+
+    start True = putStrLn $ tmsg total
+    start False = putStr $ unwords ["Indexing", show total, "source files: 0%"]
+
+    finish True = do
+      ANSI.cursorUpLine 1
+      putStrLn $ tmsg 0
+    finish False =
+      putStrLn $ " ... 100%\nIndexed " ++ show total ++ "source files"
+
+    message True before now = when (before /= now) $ do
+      ANSI.cursorUpLine 1
+      putStrLn $ tmsg now
+    message False before now = do
+      let tenth n = ((total - n) * 10) `div` total
+          t = tenth now
+      when (t /= 0 && t /= 10 && t /= tenth before) $
+        putStr $ " ... " <> show (tenth now) <> "0%"
+
+    showProgress terminal before = do
+      ranges <- mapM (Worklist.get worklist) [0 .. jobs-1]
+      let remaining = max 0 $ sum $ map (\(Worklist.Range i k) -> k-i) ranges
+      flush $ message terminal before remaining
+      threadDelay 1000000
+      showProgress terminal remaining
+
+withLog :: Bool -> (String -> IO ()) -> (StdStream -> IO a) -> IO a
+withLog verbose log act
+  | verbose = act Inherit
+  | otherwise = bracket createPipe (\(r,w) -> hClose r >> hClose w)
+      $ \(outRead, outWrite) ->
+          withAsync (log =<< hGetContents outRead)
+          $ const $ act $ UseHandle outWrite
 
 -- | Simple concurrent logger. Spawn the process and asynchronously log
 -- concise or full contents to stdout. Should use a fancy progress bar really
 spawnAndConcurrentLog :: Bool -> FilePath -> [String] -> IO ()
-spawnAndConcurrentLog verbose exe args = do
-  (_, Just hout, Just herr, ph) <- createProcess (proc exe args)
-      { std_out = CreatePipe, std_err = CreatePipe }
-  ex <- withAsync (log hout) $ \asyncOut ->
-    withAsync (log herr) $ \asyncErr -> do
-      status <- waitForProcess ph
-      cancel asyncOut
-      cancel asyncErr
-      putStr "\n" >> hFlush stdout
-      return status
+spawnAndConcurrentLog verbose exe args = withLog verbose log $ \stream -> do
+  (_, _, _, ph) <- createProcess (proc exe args)
+    { std_out = stream, std_err = stream }
+  ex <- waitForProcess ph
   case ex of
     ExitSuccess -> return ()
     ExitFailure i -> error $
       unwords (exe:args) ++ " returned exit code " ++ show i
   where
-    log h = mapM_ draw . lines =<< hGetContents h
-    draw s
-      | verbose = putStrLn s
-      | otherwise = putChar '.' >> hFlush stdout
+    log s = mapM_ (const $ putChar '.' >> hFlush stdout) (lines s)
+      `finally` do
+        putStr "\n"
+        hFlush stdout
 
 --
 -- We need to find clang-index and clang-derive in $PATH or in-tree
 --
-withExe :: FilePath -> Maybe FilePath -> (FilePath -> IO ()) -> IO ()
+withExe :: FilePath -> Maybe FilePath -> (FilePath -> IO a) -> IO a
 withExe _ (Just exePath) f = do
   exeExists <- doesFileExist exePath
   if exeExists

glean/lang/clang/glean-clang.cabal

@@ -38,6 +38,7 @@ flag opt
 
 common deps
     build-depends:
+        ansi-terminal ^>= 0.11,
         array ^>=0.5.2.0,
         async ^>=2.2.1,
         base >=4.11.1 && <4.15,

glean/lang/clang/index.cpp

@@ -91,6 +91,12 @@ DEFINE_string(clang_resource_dir, "", "PATH to Clang resource dir");
 DEFINE_string(cdb_target, "", "Target name");
 DEFINE_string(cdb_dir, "", "Directory with compile_commands.json in it");
 
+// This is a hack to support parallel indexing in the Glean CLI
+DEFINE_bool(
+    print_sources_count,
+    false,
+    "Print the number source files and exit");
+
 static llvm::cl::OptionCategory indexerCategory("glean");
 
 // This file implements some plumbing and the main function for the
@@ -229,7 +235,7 @@ struct Config {
       // No logging when dumping to a file
       should_log = false;
       sender = fileWriter(FLAGS_dump);
-    } else {
+    } else if (!FLAGS_print_sources_count) {
       fail("missing --service or --dump");
     }
 
@@ -590,6 +596,11 @@ int main(int argc, char **argv) {
 
   Config config(argc, argv);
 
+  if (FLAGS_print_sources_count) {
+    std::cout << config.sources.size();
+    return 0;
+  }
+
   const auto work_counter = FLAGS_work_file.empty()
     ? worklist::serialCounter(0, config.sources.size())
     : worklist::stealingCounter(

glean/rts/ffi.cpp

@@ -7,11 +7,9 @@
  */
 
 #ifdef OSS
-#include <cpp/ffi.h> // @manual
 #include <cpp/memory.h> // @manual
 #include <cpp/wrap.h> // @manual
 #else
-#include <common/hs/util/cpp/ffi.h>
 #include <common/hs/util/cpp/memory.h>
 #include <common/hs/util/cpp/wrap.h>
 #endif