diff --git a/init.lua b/init.lua index fa7971c..967f158 100644 --- a/init.lua +++ b/init.lua @@ -1,9 +1,29 @@ -if vim.env["VIRTUAL_ENV"] ~= nil then - vim.g.python3_host_prog = vim.fn.system("which -a python3 | sed -n 2p | tr -d '\n'") -else - vim.g.python3_host_prog = vim.fn.system("which python3 | tr -d '\n'") +local should_profile = os.getenv("NVIM_PROFILE") +if should_profile then + require("profile").instrument_autocmds() + if should_profile:lower():match("^start") then + require("profile").start("*") + else + require("profile").instrument("*") + end end +local function toggle_profile() + local prof = require("profile") + if prof.is_recording() then + prof.stop() + vim.ui.input({ prompt = "Save profile to:", completion = "file", default = "profile.json" }, function(filename) + if filename then + prof.export(filename) + vim.notify(string.format("Wrote %s", filename)) + end + end) + else + prof.start("*") + end +end +vim.keymap.set("", "", toggle_profile) + -- lazy.nvim bootstrap local lazypath = vim.fn.stdpath("data") .. "/lazy/lazy.nvim" if not vim.loop.fs_stat(lazypath) then diff --git a/lua/profile.lua b/lua/profile.lua new file mode 100644 index 0000000..0fafc37 --- /dev/null +++ b/lua/profile.lua @@ -0,0 +1,143 @@ +local autocmd = require("profile.autocmd") +local clock = require("profile.clock") +local instrument = require("profile.instrument") +local util = require("profile.util") +local M = {} + +local event_defaults = { + pid = 1, + tid = 1, +} + +---Call this at the top of your init.vim to get durations for autocmds. If you +---don't, autocmds will show up as 'instant' events in the profile +M.instrument_autocmds = function() + autocmd.instrument_start() +end + +---Instrument matching modules +---@param name string Name of module or glob pattern (e.g. "telescope*") +M.instrument = function(name) + instrument(name) +end + +---Mark matching modules to be ignored by profiling +---@param name string Name of module or glob pattern (e.g. "telescope*") +M.ignore = function(name) + instrument.ignore(name) +end + +---@param sample_rate number Float between 0 and 1 +M.set_sample_rate = function(sample_rate) + instrument.set_sample_rate(sample_rate) +end + +---Start collecting data for a profile +---@param ... string Names or patterns of modules to instrument (if instrument() not called before this) +M.start = function(...) + for _, pattern in ipairs({ ... }) do + instrument(pattern) + end + autocmd.instrument_auto() + instrument.clear_events() + clock.reset() + instrument.recording = true + vim.api.nvim_exec_autocmds("User", { pattern = "ProfileStart", modeline = false }) +end + +---@return boolean +M.is_recording = function() + return instrument.recording +end + +---@param filename? string If present, write the profile data to this file +M.stop = function(filename) + instrument.recording = false + vim.api.nvim_exec_autocmds("User", { pattern = "ProfileStop", modeline = false }) + if filename then + M.export(filename) + end +end + +---@private +---@param name string Name of the function +---@param ... any Arguments to the function +M.log_start = function(name, ...) + if not instrument.recording then + return + end + instrument.add_event({ + name = name, + args = util.format_args(...), + cat = "function,manual", + ph = "B", + ts = clock(), + }) +end + +---@private +---@param name string Name of the function +---@param ... any Arguments to the function +M.log_end = function(name, ...) + if not instrument.recording then + return + end + instrument.add_event({ + name = name, + args = util.format_args(...), + cat = "function,manual", + ph = "E", + ts = clock(), + }) +end + +---@private +---@param name string Name of the function +---@param ... any Arguments to the function +M.log_instant = function(name, ...) + if not instrument.recording then + return + end + instrument.add_event({ + name = name, + args = util.format_args(...), + cat = "", + ph = "i", + ts = clock(), + s = "g", + }) +end + +---Print out a list of all modules that have been instrumented +M.print_instrumented_modules = function() + instrument.print_modules() +end + +---Write the trace to a file +---@param filename string +M.export = function(filename) + local file = assert(io.open(filename, "w")) + local events = instrument.get_events() + file:write("[") + for i, event in ipairs(events) do + local e = vim.tbl_extend("keep", event, event_defaults) + local ok, jse = pcall(vim.json.encode, e) + if not ok and e.args then + e.args = nil + ok, jse = pcall(vim.json.encode, e) + end + if ok then + file:write(jse) + if i < #events then + file:write(",\n") + end + else + local err = string.format("Could not encode event: %s\n%s", jse, vim.inspect(e)) + vim.api.nvim_echo({ { err, "Error" } }, true, {}) + end + end + file:write("]") + file:close() +end + +return M diff --git a/lua/profile/autocmd.lua b/lua/profile/autocmd.lua new file mode 100644 index 0000000..f7a5ee1 --- /dev/null +++ b/lua/profile/autocmd.lua @@ -0,0 +1,154 @@ +local M = {} + +local autocmds = { + "BufAdd", + "BufDelete", + "BufEnter", + "BufFilePost", + "BufFilePre", + "BufHidden", + "BufLeave", + "BufModifiedSet", + "BufNew", + "BufNewFile", + "BufRead", + -- "BufReadCmd", + "BufReadPre", + "BufUnload", + "BufWinEnter", + "BufWinLeave", + "BufWipeout", + "BufWrite", + -- "BufWriteCmd", + "BufWritePost", + "ChanInfo", + "ChanOpen", + "CmdUndefined", + "CmdlineChanged", + "CmdlineEnter", + "CmdlineLeave", + "CmdwinEnter", + "CmdwinLeave", + "ColorScheme", + "ColorSchemePre", + "CompleteChanged", + "CompleteDonePre", + "CompleteDone", + "CursorHold", + "CursorHoldI", + "CursorMoved", + "CursorMovedI", + "DiffUpdated", + "DirChanged", + -- "FileAppendCmd", + "FileAppendPost", + "FileAppendPre", + "FileChangedRO", + "ExitPre", + "FileChangedShell", + "FileChangedShellPost", + -- "FileReadCmd", + "FileReadPost", + "FileReadPre", + "FileType", + -- "FileWriteCmd", + "FileWritePost", + "FileWritePre", + "FilterReadPost", + "FilterReadPre", + "FilterWritePost", + "FilterWritePre", + "FocusGained", + "FocusLost", + "FuncUndefined", + "UIEnter", + "UILeave", + "InsertChange", + "InsertCharPre", + "TextYankPost", + "InsertEnter", + "InsertLeavePre", + "InsertLeave", + "LspAttach", + "LspDetach", + "LspTokenUpdate", + "MenuPopup", + "OptionSet", + "QuickFixCmdPre", + "QuickFixCmdPost", + "QuitPre", + "RemoteReply", + "SessionLoadPost", + "ShellCmdPost", + "Signal", + "ShellFilterPost", + "SourcePre", + "SourcePost", + -- "SourceCmd", + "SpellFileMissing", + "StdinReadPost", + "StdinReadPre", + "SwapExists", + "Syntax", + "TabEnter", + "TabLeave", + "TabNew", + "TabNewEntered", + "TabClosed", + "TermOpen", + "TermEnter", + "TermLeave", + "TermClose", + "TermResponse", + "TextChanged", + "TextChangedI", + "TextChangedP", + "User", + "VimEnter", + "VimLeave", + "VimLeavePre", + "VimResized", + "VimResume", + "VimSuspend", + "WinClosed", + "WinEnter", + "WinLeave", + "WinNew", + "WinScrolled", +} + +---@param groupname string +---@param fn string +local function create(groupname, fn) + local aug = vim.api.nvim_create_augroup(groupname, {}) + for _, autocmd in ipairs(autocmds) do + vim.api.nvim_create_autocmd(autocmd, { + desc = "profile.nvim " .. fn, + pattern = "*", + group = aug, + callback = function(args) + require("profile")[fn](autocmd, { match = args.match }) + end, + }) + end +end + +M.instrument_start = function() + create("lua_profile_start", "log_start") +end + +M.instrument_auto = function() + if vim.fn.exists("#lua_profile_start") ~= 0 then + create("lua_profile_end", "log_end") + else + create("lua_profile", "log_instant") + end +end + +M.clear = function() + vim.api.nvim_create_augroup("lua_profile", {}) + vim.api.nvim_create_augroup("lua_profile_start", {}) + vim.api.nvim_create_augroup("lua_profile_end", {}) +end + +return M diff --git a/lua/profile/clock.lua b/lua/profile/clock.lua new file mode 100644 index 0000000..0b53156 --- /dev/null +++ b/lua/profile/clock.lua @@ -0,0 +1,14 @@ +local hrtime = vim.loop.hrtime + +local start = hrtime() + +return setmetatable({ + reset = function() + start = hrtime() + end, +}, { + __call = function() + -- Microseconds + return (hrtime() - start) / 1e3 + end, +}) diff --git a/lua/profile/instrument.lua b/lua/profile/instrument.lua new file mode 100644 index 0000000..a0c0b69 --- /dev/null +++ b/lua/profile/instrument.lua @@ -0,0 +1,209 @@ +local clock = require("profile.clock") +local util = require("profile.util") +local M = {} + +local rawrequire = require +local events = {} +local ignore_list = { + "^_G$", + "^bit$", + "^coroutine$", + "^debug$", + "^ffi$", + "^io$", + "^jit.*$", + "^luv$", + "^math$", + "^os$", + "^package$", + "^string$", + "^table$", + "^vim%.inspect$", + "^profile.*$", + "^lspconfig%.util%.script_path$", + "^plenary%.async_lib.*$", +} + +local instrument_list = {} +local wrapped_modules = {} +local wrapped_functions = {} +M.recording = false +M.sample_rate = 1 +local exposed_globals = { + ["vim"] = vim, + ["vim.fn"] = vim.fn, + ["vim.api"] = vim.api, +} + +local function all_modules() + return vim.tbl_extend("keep", package.loaded, exposed_globals) +end + +local function should_instrument(name) + -- Don't double-wrap + if wrapped_functions[name] then + return false + elseif wrapped_modules[name] then + return false + else + for _, pattern in ipairs(ignore_list) do + if string.match(name, pattern) then + return false + end + end + return true + end +end + +local function wrap_function(name, fn) + return function(...) + if M.sample_rate < 1 and math.random() > M.sample_rate then + return fn(...) + end + local arg_string = util.format_args(...) + local start = clock() + local function handle_result(...) + local delta = clock() - start + M.add_event({ + name = name, + args = arg_string, + cat = "function", + ph = "X", + ts = start, + dur = delta, + }) + return ... + end + return handle_result(fn(...)) + end +end + +local function patch_function(modname, mod, name) + local fn = mod[name] + if type(fn) == "function" then + local fnname = string.format("%s.%s", modname, name) + if should_instrument(fnname) then + wrapped_functions[fnname] = fn + mod[name] = wrap_function(fnname, fn) + end + end +end + +local function maybe_wrap_function(name) + local parent_mod_name, fn = util.split_path(name) + local mod = M.get_module(parent_mod_name) + if mod then + patch_function(parent_mod_name, mod, fn) + end +end + +local function wrap_module(name, mod) + name = util.normalize_module_name(name) + if type(mod) ~= "table" or not should_instrument(name) then + return + end + wrapped_modules[name] = true + for k in pairs(mod) do + -- Do not wrap module functions that start with '_' + -- Those have to be explicitly passed to instrument() + if type(k) == "string" and not util.startswith(k, "_") then + patch_function(name, mod, k) + end + end +end + +local function should_profile_module(name) + for _, pattern in ipairs(instrument_list) do + if string.match(name, pattern) then + return true + end + end + return false +end + +M.hook_require = function(module_name) + if rawrequire ~= require then + return + end + local wrapped_require = wrap_function("require", rawrequire) + _G.require = function(name) + -- Don't time the require if the file is already loaded + if package.loaded[name] or not should_profile_module(name) then + return rawrequire(name) + end + local mod = wrapped_require(name) + wrap_module(name, mod) + return mod + end +end + +M.clear_events = function() + events = {} +end + +M.add_event = function(event) + if M.recording then + table.insert(events, event) + end +end + +M.get_module = function(name) + local ok, mod = pcall(require, name) + if ok then + if type(mod) == "table" then + return mod + else + return nil + end + else + mod = _G + local paths = vim.split(name, ".", { plain = true }) + for _, token in ipairs(paths) do + mod = mod[token] + if not mod then + break + end + end + return type(mod) == "table" and mod or nil + end +end + +M.get_events = function() + return events +end + +M.ignore = function(name) + table.insert(ignore_list, util.path_glob_to_regex(name)) +end + +M.print_modules = function() + for module, _ in pairs(wrapped_modules) do + print(module) + end +end + +local function instrument(_, name) + local pattern = util.path_glob_to_regex(name) + if not vim.tbl_contains(instrument_list, pattern) then + table.insert(instrument_list, pattern) + end + M.hook_require(name) + for modname, mod in pairs(all_modules()) do + if string.match(modname, pattern) then + wrap_module(modname, mod) + end + end + maybe_wrap_function(name) +end + +---@param sample_rate number Float between 0 and 1 +M.set_sample_rate = function(sample_rate) + if sample_rate <= 0 or sample_rate > 1 then + error("sample_rate must be between 0 (exclusive) and 1 (inclusive)") + end + M.sample_rate = sample_rate +end + +return setmetatable(M, { + __call = instrument, +}) diff --git a/lua/profile/util.lua b/lua/profile/util.lua new file mode 100644 index 0000000..b1648d3 --- /dev/null +++ b/lua/profile/util.lua @@ -0,0 +1,83 @@ +local M = {} + +local MAX_ARG_LEN = 200 +local tbl_isarray = vim.isarray or vim.tbl_isarray or vim.tbl_islist +local pack_len = vim.F.pack_len +local split = vim.split + +---@param glob string +---@return string +M.path_glob_to_regex = function(glob) + local pattern = string.gsub(glob, "%.", "[%./]") + pattern = string.gsub(pattern, "*", ".*") + return "^" .. pattern .. "$" +end + +---@param name string +---@return string +M.normalize_module_name = function(name) + local ret = string.gsub(name, "/", ".") + return ret +end + +---@param haystack string +---@param prefix string +---@return boolean +M.startswith = function(haystack, prefix) + return string.find(haystack, prefix) == 1 +end + +---@param path string +---@return string module +---@return string tail +M.split_path = function(path) + local pieces = split(path, ".", { plain = true }) + if #pieces == 1 then + return "_G", path + end + local mod = table.concat(pack_len(unpack(pieces, 1, #pieces - 1)), ".") + return mod, pieces[#pieces] +end + +local function sanitize(table) + local clean = {} + local iterfn + if tbl_isarray(table) then + iterfn = ipairs + else + iterfn = pairs + end + for i, v in iterfn(table) do + local t = type(v) + if t == "string" then + if string.len(v) > MAX_ARG_LEN then + clean[tostring(i)] = string.sub(v, 1, MAX_ARG_LEN - 3) .. "..." + else + clean[tostring(i)] = v + end + elseif t == "nil" or t == "boolean" or t == "number" then + clean[tostring(i)] = v + end + end + -- If no args, then return nil + if next(clean) == nil then + return nil + else + return clean + end +end + +---@param ... any[] +---@return any +M.format_args = function(...) + local args = pack_len(...) + if args.n == 0 then + return nil + elseif args.n == 1 and type(args[1]) == "table" then + return sanitize(args[1]) + else + return sanitize(args) + end +end + +return M