Merge pull request #2638 from emmericp/multi-threaded-check

Add multi-process support to --check.

3 files changed, 248 insertions, 111 deletions

diff --git a/script/cli/check.lua b/script/cli/check.lua
index 3902c4aa..8b314f24 100644
--- a/script/cli/check.lua
+++ b/script/cli/check.lua
@@ -1,130 +1,96 @@
-local lclient  = require 'lclient'()
-local furi     = require 'file-uri'
-local ws       = require 'workspace'
-local files    = require 'files'
-local diag     = require 'provider.diagnostic'
-local util     = require 'utility'
-local jsonb    = require 'json-beautify'
-local lang     = require 'language'
-local define   = require 'proto.define'
-local config   = require 'config.config'
-local fs       = require 'bee.filesystem'
-local provider = require 'provider'
+local lang       = require 'language'
+local platform   = require 'bee.platform'
+local subprocess = require 'bee.subprocess'
+local json       = require 'json'
+local jsonb      = require 'json-beautify'
+local util       = require 'utility'
 
-require 'plugin'
-require 'vm'
 
-lang(LOCALE)
+local numThreads = tonumber(NUM_THREADS or 1)
 
-if type(CHECK) ~= 'string' then
-    print(lang.script('CLI_CHECK_ERROR_TYPE', type(CHECK)))
-    return
+local exe = arg[-1]
+-- TODO: is this necessary? got it from the shell.lua helper in bee.lua tests
+if platform.os == 'windows' and not exe:match('%.[eE][xX][eE]$') then
+    exe = exe..'.exe'
 end
 
-local rootPath = fs.absolute(fs.path(CHECK)):string()
-local rootUri = furi.encode(rootPath)
-if not rootUri then
-    print(lang.script('CLI_CHECK_ERROR_URI', rootPath))
-    return
+local function logFileForThread(threadId)
+    return LOGPATH .. '/check-partial-' .. threadId .. '.json'
 end
-rootUri = rootUri:gsub("/$", "")
 
-if CHECKLEVEL then
-    if not define.DiagnosticSeverity[CHECKLEVEL] then
-        print(lang.script('CLI_CHECK_ERROR_LEVEL', 'Error, Warning, Information, Hint'))
-        return
+local function buildArgs(threadId)
+    local args = {exe}
+    local skipNext = false
+    for i = 1, #arg do
+        local arg = arg[i]
+        -- --check needs to be transformed into --check_worker
+        if arg:lower():match('^%-%-check$') or arg:lower():match('^%-%-check=') then
+            args[#args + 1] = arg:gsub('%-%-%w*', '--check_worker')
+        -- --check_out_path needs to be removed if we have more than one thread
+        elseif arg:lower():match('%-%-check_out_path') and numThreads > 1 then
+            if not arg:match('%-%-%w*=') then
+                skipNext = true
+            end
+        else
+            if skipNext then
+                skipNext = false
+            else
+                args[#args + 1] = arg
+            end
+        end
+    end
+    args[#args + 1] = '--thread_id'
+    args[#args + 1] = tostring(threadId)
+    if numThreads > 1 then
+        args[#args + 1] = '--quiet'
+        args[#args + 1] = '--check_out_path'
+        args[#args + 1] = logFileForThread(threadId)
     end
+    return args
 end
-local checkLevel = define.DiagnosticSeverity[CHECKLEVEL] or define.DiagnosticSeverity.Warning
-
-util.enableCloseFunction()
 
-local lastClock   = os.clock()
-local results = {}
-
-local function errorhandler(err)
-	print(err)
-	print(debug.traceback())
+if numThreads > 1 then
+    print(lang.script('CLI_CHECK_MULTIPLE_WORKERS', numThreads))
 end
 
----@async
-xpcall(lclient.start, errorhandler, lclient, function (client)
-    client:registerFakers()
-
-    client:initialize {
-        rootUri = rootUri,
-    }
-
-    client:register('textDocument/publishDiagnostics', function (params)
-        results[params.uri] = params.diagnostics
-    end)
-
-    io.write(lang.script('CLI_CHECK_INITING'))
-
-    provider.updateConfig(rootUri)
-
-    ws.awaitReady(rootUri)
-
-    local disables = util.arrayToHash(config.get(rootUri, 'Lua.diagnostics.disable'))
-    for name, serverity in pairs(define.DiagnosticDefaultSeverity) do
-        serverity = config.get(rootUri, 'Lua.diagnostics.severity')[name] or 'Warning'
-        if serverity:sub(-1) == '!' then
-            serverity = serverity:sub(1, -2)
-        end
-        if define.DiagnosticSeverity[serverity] > checkLevel then
-            disables[name] = true
-        end
+local procs = {}
+for i = 1, numThreads do
+    local process, err = subprocess.spawn({buildArgs(i)})
+    if err then
+        print(err)
     end
-    config.set(rootUri, 'Lua.diagnostics.disable', util.getTableKeys(disables, true))
-
-    local uris = files.getChildFiles(rootUri)
-    local max  = #uris
-    for i, uri in ipairs(uris) do
-        files.open(uri)
-        diag.doDiagnostic(uri, true)
-		-- Print regularly but always print the last entry to ensure that logs written to files don't look incomplete.
-        if os.clock() - lastClock > 0.2 or i == #uris then
-            lastClock = os.clock()
-			client:update()
-            local output = '\x0D'
-                        .. ('>'):rep(math.ceil(i / max * 20))
-                        .. ('='):rep(20 - math.ceil(i / max * 20))
-                        .. ' '
-                        .. ('0'):rep(#tostring(max) - #tostring(i))
-                        .. tostring(i) .. '/' .. tostring(max)
-            io.write(output)
-			local filesWithErrors = 0
-			local errors = 0
-			for _, diags in pairs(results) do
-				filesWithErrors = filesWithErrors + 1
-				errors = errors + #diags
-			end
-			if errors > 0 then
-				local errorDetails = ' [' .. lang.script('CLI_CHECK_PROGRESS', errors, filesWithErrors) .. ']'
-				io.write(errorDetails)
-			end
-            io.flush()
-        end
+    if process then
+        procs[#procs + 1] = process
     end
-    io.write('\x0D')
-end)
+end
 
-local count = 0
-for uri, result in pairs(results) do
-    count = count + #result
-    if #result == 0 then
-        results[uri] = nil
-    end
+for _, process in ipairs(procs) do
+    process:wait()
 end
 
-if count == 0 then
-    print(lang.script('CLI_CHECK_SUCCESS'))
-else
-    local outpath = CHECK_OUT_PATH
-    if outpath == nil then
-        outpath = LOGPATH .. '/check.json'
-    end
-    util.saveFile(outpath, jsonb.beautify(results))
+local outpath = CHECK_OUT_PATH
+if outpath == nil then
+    outpath = LOGPATH .. '/check.json'
+end
 
-    print(lang.script('CLI_CHECK_RESULTS', count, outpath))
+if numThreads > 1 then
+    local mergedResults = {}
+    local count = 0
+    for i = 1, numThreads do
+        local result = json.decode(util.loadFile(logFileForThread(i)) or '[]')
+        for k, v in pairs(result) do
+            local entries = mergedResults[k] or {}
+            mergedResults[k] = entries
+            for _, entry in ipairs(v) do
+                entries[#entries + 1] = entry
+                count = count + 1
+            end
+        end
+    end
+    util.saveFile(outpath, jsonb.beautify(mergedResults))
+    if count == 0 then
+        print(lang.script('CLI_CHECK_SUCCESS'))
+    else
+        print(lang.script('CLI_CHECK_RESULTS', count, outpath))
+    end
 end
diff --git a/script/cli/check_worker.lua b/script/cli/check_worker.lua
new file mode 100644
index 00000000..f8be88d6
--- /dev/null
+++ b/script/cli/check_worker.lua
@@ -0,0 +1,166 @@
+local lclient   = require 'lclient'()
+local furi      = require 'file-uri'
+local ws        = require 'workspace'
+local files     = require 'files'
+local diag      = require 'provider.diagnostic'
+local util      = require 'utility'
+local jsonb     = require 'json-beautify'
+local lang      = require 'language'
+local define    = require 'proto.define'
+local protoDiag = require 'proto.diagnostic'
+local config    = require 'config.config'
+local fs        = require 'bee.filesystem'
+local provider  = require 'provider'
+require 'plugin'
+require 'vm'
+
+lang(LOCALE)
+
+local numThreads = tonumber(NUM_THREADS or 1)
+local threadId = tonumber(THREAD_ID or 1)
+
+if type(CHECK_WORKER) ~= 'string' then
+    print(lang.script('CLI_CHECK_ERROR_TYPE', type(CHECK_WORKER)))
+    return
+end
+
+local rootPath = fs.absolute(fs.path(CHECK_WORKER)):string()
+local rootUri = furi.encode(rootPath)
+if not rootUri then
+    print(lang.script('CLI_CHECK_ERROR_URI', rootPath))
+    return
+end
+rootUri = rootUri:gsub("/$", "")
+
+if CHECKLEVEL then
+    if not define.DiagnosticSeverity[CHECKLEVEL] then
+        print(lang.script('CLI_CHECK_ERROR_LEVEL', 'Error, Warning, Information, Hint'))
+        return
+    end
+end
+local checkLevel = define.DiagnosticSeverity[CHECKLEVEL] or define.DiagnosticSeverity.Warning
+
+util.enableCloseFunction()
+
+-- Hash function used to distribute work.
+local function hashString(str)
+    local hash = 0
+    for i = 1, #str do
+        hash = (hash * 37 & 0xFFFFFFFF) + str:byte(i, i)
+    end
+    return hash
+end
+
+local lastClock = os.clock()
+local results = {}
+
+local function errorhandler(err)
+    print(err)
+    print(debug.traceback())
+end
+
+---@async
+xpcall(lclient.start, errorhandler, lclient, function (client)
+    client:registerFakers()
+
+    client:initialize {
+        rootUri = rootUri,
+    }
+
+    client:register('textDocument/publishDiagnostics', function (params)
+        results[params.uri] = params.diagnostics
+    end)
+
+    if not QUIET then
+        io.write(lang.script('CLI_CHECK_INITING'))
+    end
+
+    provider.updateConfig(rootUri)
+
+    ws.awaitReady(rootUri)
+
+    local disables = util.arrayToHash(config.get(rootUri, 'Lua.diagnostics.disable'))
+    for name, serverity in pairs(define.DiagnosticDefaultSeverity) do
+        serverity = config.get(rootUri, 'Lua.diagnostics.severity')[name] or 'Warning'
+        if serverity:sub(-1) == '!' then
+            serverity = serverity:sub(1, -2)
+        end
+        if define.DiagnosticSeverity[serverity] > checkLevel then
+            disables[name] = true
+        end
+    end
+    config.set(rootUri, 'Lua.diagnostics.disable', util.getTableKeys(disables, true))
+
+    -- Downgrade file opened status to Opened for everything to avoid reporting during compilation on files that do not belong to this thread
+    local diagStatus = config.get(rootUri, 'Lua.diagnostics.neededFileStatus')
+    for diag, status in pairs(diagStatus) do
+        if status == 'Any' or status == 'Any!' then
+            diagStatus[diag] = 'Opened!'
+        end
+    end
+    for diag, status in pairs(protoDiag.getDefaultStatus()) do
+        if status == 'Any' or status == 'Any!' then
+            diagStatus[diag] = 'Opened!'
+        end
+    end
+    config.set(rootUri, 'Lua.diagnostics.neededFileStatus', diagStatus)
+
+    local uris = files.getChildFiles(rootUri)
+    local max  = #uris
+    for i, uri in ipairs(uris) do
+        local hash = hashString(uri) % numThreads + 1
+        if hash == threadId then
+            files.open(uri)
+            diag.doDiagnostic(uri, true)
+            -- Print regularly but always print the last entry to ensure that logs written to files don't look incomplete.
+            if (os.clock() - lastClock > 0.2 or i == #uris) and not QUIET then
+                lastClock = os.clock()
+                client:update()
+                local output = '\x0D'
+                            .. ('>'):rep(math.ceil(i / max * 20))
+                            .. ('='):rep(20 - math.ceil(i / max * 20))
+                            .. ' '
+                            .. ('0'):rep(#tostring(max) - #tostring(i))
+                            .. tostring(i) .. '/' .. tostring(max)
+                io.write(output)
+                local filesWithErrors = 0
+                local errors = 0
+                for _, diags in pairs(results) do
+                    filesWithErrors = filesWithErrors + 1
+                    errors = errors + #diags
+                end
+                if errors > 0 then
+                    local errorDetails = ' [' .. lang.script('CLI_CHECK_PROGRESS', errors, filesWithErrors) .. ']'
+                    io.write(errorDetails)
+                end
+                io.flush()
+            end
+        end
+    end
+    if not QUIET then
+        io.write('\x0D')
+    end
+end)
+
+local count = 0
+for uri, result in pairs(results) do
+    count = count + #result
+    if #result == 0 then
+        results[uri] = nil
+    end
+end
+
+local outpath = CHECK_OUT_PATH
+if outpath == nil then
+    outpath = LOGPATH .. '/check.json'
+end
+-- Always write result, even if it's empty to make sure no one accidentally looks at an old output after a successful run.
+util.saveFile(outpath, jsonb.beautify(results))
+
+if not QUIET then
+    if count == 0 then
+        print(lang.script('CLI_CHECK_SUCCESS'))
+    else
+        print(lang.script('CLI_CHECK_RESULTS', count, outpath))
+    end
+end
diff --git a/script/cli/init.lua b/script/cli/init.lua
index d37c50ae..65e7e102 100644
--- a/script/cli/init.lua
+++ b/script/cli/init.lua
@@ -8,6 +8,11 @@ if _G['CHECK'] then
     os.exit(0, true)
 end
 
+if _G['CHECK_WORKER'] then
+    require 'cli.check_worker'
+    os.exit(0, true)
+end
+
 if _G['DOC_UPDATE'] then
     require 'cli.doc' .runCLI()
     os.exit(0, true)