forkmon.nelua

--[[
forkmon by edubart
Eduardo Bart - edub4rt@gmail.com
https://github.com/edubart/forkmon

Watch for file changes and auto restart an application using fork checkpoints to continue.
Intended for quick live development.

See end of file for license.
]]

require 'sys'
require 'string'
require 'hashmap'
require 'io'
require 'string'
require 'C.stdarg'

local restart_delay: uint64 = 50 * 1000
local ignore_delay: uint64 = 200 * 1000
local initialized: boolean = false
local handling_signal: boolean = false
local quiet: boolean = false
local no_colors: boolean = false
local preload_disabled: boolean = false
local parent_inotify_fd: cint = -1
local child_pid: cint = 0
local root_pid: cint = getpid()
local last_restart_ticks: uint64
local filter_patts_count: cint = 0
local filter_patts: [32]string
local time_checkpoint: timespec
local tracked_files: hashmap(string, boolean)
local FopenFunc: type = @function(cstring, cstring): *FILE
local OpenFunc: type = @function(cstring, cint, cint): cint
local orig_fopen: FopenFunc
local orig_fopen64: FopenFunc
local orig_open: OpenFunc

-- Exit process and children processes.
local function killapp(code: cint)
  killwait(root_pid, SIGUSR2)
  exit(code)
end

-- Format info message and output to stderr.
local function logf(fmt: string, ...: varargs)
  if not quiet then
    if not no_colors then io.stderr:write(Colors.Yellow) end
    io.stderr:write('[forkmon] ')
    io.stderr:writef(fmt, ...)
    if not no_colors then io.stderr:write(Colors.Reset) end
    io.stderr:write('\n')
    io.stderr:flush()
  end
end

-- Format error message, output to stderr and kill the application.
local function errorf(fmt: string, ...: varargs)
  if not no_colors then io.stderr:write(Colors.Red) end
  io.stderr:write('[forkmon] ')
  io.stderr:writef(fmt, ...)
  if not no_colors then io.stderr:write(Colors.Reset) end
  io.stderr:write('\n')
  io.stderr:flush()
  killapp(-1)
end

-- Kill the application with an error message if `cond` is false.
local function assertkill(cond: boolean, msg: string)
  if not cond then
    errorf(msg)
  end
end

-- Check weather a filename matches a filter pattern.
local function filter_filename(filename: string): boolean
  for i=0,<filter_patts_count do
    local patt = filter_patts[i]
    if filename:find(patt) ~= 0 then
      return true
    end
  end
  return false
end

-- Create a new inotify instance and add file to watch.
local function inotify_init_watch(filename: cstring, mask: cint): (cint, cint)
  local fd = inotify_init()
  assertkill(fd >= 0, 'failed to initialize inotify, maybe increase "user.max_inotify_instances" in sysctl?')
  local wd = inotify_add_watch(fd, filename, mask) -- watch the file
  assertkill(wd >= 0, 'failed to watch file after init')
  return fd, wd
end

-- Monitor a file for changes, returns true when the filed changed.
local function monitor_file(filename: string): boolean
  local inotify_mask: cint = IN_MODIFY | IN_CREATE | IN_DELETE_SELF | IN_MOVE_SELF
  if parent_inotify_fd >= 0 then -- a parent process restarted due to a file change
    -- lets just add this file to track in the parent inotify fd
    -- so we can avoid lot's of forks when changing the same file
    -- logf("watch dependency '%s'\n", filename)
    local inotify_wd = inotify_add_watch(parent_inotify_fd, filename, inotify_mask) -- watch the file
    assertkill(inotify_wd >= 0, 'failed to watch a file')
    return false
  end
  if not quiet then
    logf("watch '%s'", filename)
  end
  local pid = fork() -- fork in two processes
  assertkill(pid >= 0, 'fork failed')
  if pid == 0 then -- let the child process continue
    setpgid(0, 0) -- set a new process group for the child
    return false
  end
  -- parent process will only monitor the file
  child_pid = pid
  local inotify_fd, inotify_wd = inotify_init_watch(filename, inotify_mask)
  local buf: [1024]byte -- buffer for inotify events
  while true do -- keep waiting the file to change indefinitely
    local res = read(inotify_fd, nilptr, 0) -- wait for an event
    if res == -1 and errno == EINTR then -- interrupted by a signal (like Ctrl+C on terminal)
      killapp(0) -- recursive kill all children processes
    end
    usleep(restart_delay) -- wait some time so all file changes can flush
    -- must rewatch the file, because the old wd may be invalid
    inotify_rm_watch(inotify_fd, inotify_wd)
    for i=1,10 do -- try to watch 10 times
      inotify_wd = inotify_add_watch(inotify_fd, filename, inotify_mask)
      if inotify_wd >= 0 then break end -- watch succeeded
      usleep(restart_delay) -- wait more time
    end
    assertkill(inotify_wd >= 0, 'failed to rewatch file, does the file still exist?')
    local len = read(inotify_fd, &buf, #buf) -- wait an inotify event
    if len > 0 then -- got an inotify event, must be a file change
      local now: uint64 = uticks()
      if now - last_restart_ticks <= ignore_delay then
        continue -- changed too soon, just ignore
      end
      last_restart_ticks = now
      local wdchange = false
      local pos = 0
      while pos < len do
        local event: *inotify_event = (@*inotify_event)(&buf[pos])
        if event.wd == inotify_wd then
          wdchange = true
          break
        end
        pos = pos + #inotify_event + event.len
      end
      if wdchange then -- watched file for this process changed
        logf("file '%s' changed, resuming from it ...", filename)
        parent_inotify_fd = inotify_fd
      else -- watched file for a child process changed
        logf("some file changed, resuming from '%s' ...", filename)
        parent_inotify_fd = -1
        -- reinitialize inotify to cleanup all old watches
        close(inotify_fd)
        inotify_fd, inotify_wd = inotify_init_watch(filename, inotify_mask)
      end
      killwait(pid, SIGUSR2) -- propagate children kill and wait them to finish
      pid = fork() -- fork in two processes
      assertkill(pid >= 0, 'fork failed')
      if pid == 0 then -- let the child process continue
        setpgid(0, 0) -- set a new process group for the child
        local nowtext = tostring(now)
        setenv("FORKMON_RESTART_TICKS", nowtext, 1)
        nowtext:destroy()
        return true
      end
      child_pid = pid
    end
  end
  return false -- actually unreachable
end

-- Called once a file is opened.
local function track_file_open(name: cstring): boolean
  if not initialized then -- ignore when the library was not initialized yet
    return false
  end
  local buf: [4097]cchar
  local filename = resolvpath(name, &buf, #buf)
  if filename == '' or -- skip failed path resolutions
     filename:find('^/tmp') ~= 0 or -- skip temporary files
     not filter_filename(filename) or -- skip files not in the filter
     tracked_files:peek(filename) then -- skip files already tracked
    return false
  end
  if not preload_disabled then
    -- unset LD_PRELOAD because we don't want to monitor children processes
    unsetenv("LD_PRELOAD")
    preload_disabled = true
  end
  tracked_files[string.copy(filename)] = true
  return monitor_file(filename)
end

-- Signal handler.
local function signal_handler(signum: cint): void
  if handling_signal then -- signal already being handled, this should hardly happen
    return
  end
  handling_signal = true
  if child_pid ~= 0 then -- propagate kill on child process
    killwait(child_pid, SIGUSR2)
  end
  killwait(-getpid(), SIGKILL) -- kill all spawned sub processes
  _exit(0)
  return
end

local function parse_env()
  -- quiet
  local quiet_env: cstring = getenv('FORKMON_QUIET')
  if quiet_env then
    quiet = true
  end
  -- no colors
  local no_colors_env: cstring = getenv('FORKMON_NO_COLORS')
  if no_colors_env then
    no_colors = true
  end
  -- filter patterns
  local filter_env: cstring = getenv('FORKMON_FILTER')
  if not filter_env then
    errorf('the environment variable FORKMON_FILTER is not set, please set one (example FORKMON_FILTER="%.lua$")\n')
  end
  for patt in string.gmatchview(filter_env, '[^;]+') do
    if filter_patts_count >= #filter_patts then
      errorf('max number of filter patterns reached')
    end
    filter_patts[filter_patts_count] = patt
    filter_patts_count = filter_patts_count + 1
  end
  -- restart delay
  local restart_delay_env: cstring = getenv('FORKMON_RESTART_DELAY')
  if restart_delay_env then
    restart_delay = tointeger(restart_delay_env) * 1000
  end
  -- ignore delay
  local ignore_delay_env: cstring = getenv('FORKMON_IGNORE_DELAY')
  if ignore_delay_env then
    ignore_delay = tointeger(ignore_delay_env) * 1000
  end
end

do -- initialize
  -- setup signal handler, used to recursively propagate kill signal on children processes
  signal(SIGUSR2, signal_handler)
  signal(SIGTERM, signal_handler)
  signal(SIGINT, signal_handler)
  signal(SIGSEGV, signal_handler)
  signal(SIGBUS, signal_handler)
  signal(SIGABRT, signal_handler)
  signal(SIGFPE, signal_handler)

  parse_env()
  initialized = true
end

-- Hook `fopen`.
local function fopen(filename: cstring <const, restrict>, mode: cstring <const, restrict>): *FILE <cexport, codename 'fopen'>
  if not orig_fopen then
    orig_fopen = (@FopenFunc)(dlsym(RTLD_NEXT, "fopen"))
  end
  if (mode == 'r' or mode == 'rb') and filexists(filename) then
    track_file_open(filename)
  end
  local fp: *FILE = orig_fopen(filename, mode)
  return fp
end

-- Hook `fopen64`.
local function fopen64(filename: cstring <const, restrict>, mode: cstring <const, restrict>): *FILE <cexport, codename 'fopen64'>
  if not orig_fopen64 then
    orig_fopen64 = (@FopenFunc)(dlsym(RTLD_NEXT, "fopen64"))
  end
  if (mode == 'r' or mode == 'rb') and filexists(filename) then
    track_file_open(filename)
  end
  local fp: *FILE = orig_fopen64(filename, mode)
  return fp
end

-- Hook `open`.
local function open(path: cstring <const, restrict>, flags: cint, ...: cvarargs): cint <cexport, codename 'open'>
  if not orig_open then
    orig_open = (@OpenFunc)(dlsym(RTLD_NEXT, "open"))
  end
  local writeflags: cint = (O_WRONLY|O_RDWR|O_CREAT|O_APPEND)
  if flags & writeflags == 0 and filexists(path) then
    track_file_open(path)
  end
  local args: cvalist
  C.va_start(args, flags)
  local mode: cint = C.va_arg(args, @cint)
  C.va_end(args)
  local fd: cint = orig_open(path, flags, mode)
  return fd
end

-- Shared library entry point.
local function setup() <entrypoint, cattribute 'constructor'>
  local function nelua_main(argc: cint, argv: *cstring): cint <cimport,nodecl> end
  nelua_main(0, nilptr)
end

--[[
MIT License

Copyright (c) 2021 Eduardo Bart (https://github.com/edubart)

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
]]