feat: add a way to cache state between restarts of Neovim (#624)

Co-authored-by: Kristijan Husak <[email protected]>

Author: PriceHiller

lua/orgmode/state/state.lua

@@ -0,0 +1,168 @@
+local utils = require('orgmode.utils')
+local Promise = require('orgmode.utils.promise')
+
+---@class OrgState
+local OrgState = { data = {}, _ctx = { loaded = false, saved = false, curr_loader = nil, savers = 0 } }
+
+local cache_path = vim.fs.normalize(vim.fn.stdpath('cache') .. '/org-cache.json', { expand_env = false })
+
+---Returns the current OrgState singleton
+---@return OrgState
+function OrgState.new()
+  -- This is done so we can later iterate the 'data'
+  -- subtable cleanly and shove it into a cache
+  setmetatable(OrgState, {
+    __index = function(tbl, key)
+      return tbl.data[key]
+    end,
+    __newindex = function(tbl, key, value)
+      tbl.data[key] = value
+    end,
+  })
+  local self = OrgState
+  -- Start trying to load the state from cache as part of initializing the state
+  self:load()
+  return self
+end
+
+---Save the current state to cache
+---@return Promise
+function OrgState:save()
+  OrgState._ctx.saved = false
+  --- We want to ensure the state was loaded before saving.
+  self:load()
+  self._ctx.savers = self._ctx.savers + 1
+  return utils
+    .writefile(cache_path, vim.json.encode(OrgState.data))
+    :next(function()
+      self._ctx.savers = self._ctx.savers - 1
+      if self._ctx.savers == 0 then
+        OrgState._ctx.saved = true
+      end
+    end)
+    :catch(function(err_msg)
+      self._ctx.savers = self._ctx.savers - 1
+      vim.schedule_wrap(function()
+        utils.echo_warning('Failed to save current state! Error: ' .. err_msg)
+      end)
+    end)
+end
+
+---Synchronously save the state into cache
+---@param timeout? number How long to wait for the save operation
+function OrgState:save_sync(timeout)
+  OrgState._ctx.saved = false
+  self:save()
+  vim.wait(timeout or 500, function()
+    return OrgState._ctx.saved
+  end, 20)
+end
+
+---Load the state cache into the current state
+---@return Promise
+function OrgState:load()
+  --- If we currently have a loading operation already running, return that
+  --- promise. This avoids a race condition of sorts as without this there's
+  --- potential to have two OrgState:load operations occuring and whichever
+  --- finishes last sets the state. Not desirable.
+  if self._ctx.curr_loader ~= nil then
+    return self._ctx.curr_loader
+  end
+
+  --- If we've already loaded the state from cache we don't need to do so again
+  if self._ctx.loaded then
+    return Promise.resolve(self)
+  end
+
+  self._ctx.curr_loader = utils
+    .readfile(cache_path, { raw = true })
+    :next(function(data)
+      local success, decoded = pcall(vim.json.decode, data, {
+        luanil = { object = true, array = true },
+      })
+      self._ctx.curr_loader = nil
+      if not success then
+        local err_msg = vim.deepcopy(decoded)
+        vim.schedule(function()
+          utils.echo_warning('OrgState cache load failure, error: ' .. vim.inspect(err_msg))
+          -- Try to 'repair' the cache by saving the current state
+          self:save()
+        end)
+      end
+      -- Because the state cache repair happens potentially after the data has
+      -- been added to the cache, we need to ensure the decoded table is set to
+      -- empty if we got an error back on the json decode operation.
+      if type(decoded) ~= 'table' then
+        decoded = {}
+      end
+
+      -- It is possible that while the state was loading from cache values
+      -- were saved into the state. We want to preference the newer values in
+      -- the state and still get whatever values may not have been set in the
+      -- interim of the load operation.
+      self.data = vim.tbl_deep_extend('force', decoded, self.data)
+      return self
+    end)
+    :catch(function(err)
+      -- If the file didn't exist then go ahead and save
+      -- our current cache and as a side effect create the file
+      if type(err) == 'string' and err:match([[^ENOENT.*]]) then
+        self:save()
+        return self
+      end
+      -- If the file did exist, something is wrong. Kick this to the top
+      error(err)
+    end)
+    :finally(function()
+      self._ctx.loaded = true
+      self._ctx.curr_loader = nil
+    end)
+
+  return self._ctx.curr_loader
+end
+
+---Synchronously load the state from cache if it hasn't been loaded
+---@param timeout? number How long to wait for the cache load before erroring
+---@return OrgState
+function OrgState:load_sync(timeout)
+  local state
+  local err
+  self
+    :load()
+    :next(function(loaded_state)
+      state = loaded_state
+    end)
+    :catch(function(reject)
+      err = reject
+    end)
+
+  vim.wait(timeout or 500, function()
+    return state ~= nil or err ~= nil
+  end, 20)
+
+  if err then
+    error(err)
+  end
+
+  if err == nil and state == nil then
+    error('Did not load OrgState in time')
+  end
+
+  return state
+end
+
+---Reset the current state to empty
+---@param overwrite? boolean Whether or not the cache should also be wiped
+function OrgState:wipe(overwrite)
+  overwrite = overwrite or false
+
+  self.data = {}
+  self._ctx.curr_loader = nil
+  self._ctx.loaded = false
+  self._ctx.saved = false
+  if overwrite then
+    state:save_sync()
+  end
+end
+
+return OrgState.new()

lua/orgmode/utils/init.lua

@@ -6,7 +6,8 @@ local debounce_timers = {}
 local query_cache = {}
 local tmp_window_augroup = vim.api.nvim_create_augroup('OrgTmpWindow', { clear = true })
 
-function utils.readfile(file)
+function utils.readfile(file, opts)
+  opts = opts or {}
   return Promise.new(function(resolve, reject)
     uv.fs_open(file, 'r', 438, function(err1, fd)
       if err1 then
@@ -24,6 +25,9 @@ function utils.readfile(file)
             if err4 then
               return reject(err4)
             end
+            if opts.raw then
+              return resolve(data)
+            end
             local lines = vim.split(data, '\n')
             table.remove(lines, #lines)
             return resolve(lines)
@@ -34,6 +38,32 @@ function utils.readfile(file)
   end)
 end
 
+function utils.writefile(file, data)
+  return Promise.new(function(resolve, reject)
+    uv.fs_open(file, 'w', 438, function(err1, fd)
+      if err1 then
+        return reject(err1)
+      end
+      uv.fs_fstat(fd, function(err2, stat)
+        if err2 then
+          return reject(err2)
+        end
+        uv.fs_write(fd, data, nil, function(err3, bytes)
+          if err3 then
+            return reject(err3)
+          end
+          uv.fs_close(fd, function(err4)
+            if err4 then
+              return reject(err4)
+            end
+            return resolve(bytes)
+          end)
+        end)
+      end)
+    end)
+  end)
+end
+
 function utils.open(target)
   if vim.fn.executable('xdg-open') == 1 then
     return vim.fn.system(string.format('xdg-open %s', target))

tests/minimal_init.lua

@@ -65,6 +65,16 @@ function M.setup(plugins)
   vim.env.XDG_STATE_HOME = M.root('xdg/state')
   vim.env.XDG_CACHE_HOME = M.root('xdg/cache')
 
+  local std_paths = {
+    'cache',
+    'data',
+    'config',
+  }
+
+  for _, std_path in pairs(std_paths) do
+    vim.fn.mkdir(vim.fn.stdpath(std_path), 'p')
+  end
+
   -- NOTE: Cleanup the xdg cache on exit so new runs of the minimal init doesn't share any previous state, e.g. shada
   vim.api.nvim_create_autocmd('VimLeave', {
     callback = function()

tests/plenary/state/state_spec.lua

@@ -0,0 +1,133 @@
+local utils = require('orgmode.utils')
+---@type OrgState
+local state = nil
+local cache_path = vim.fs.normalize(vim.fn.stdpath('cache') .. '/org-cache.json', { expand_env = false })
+
+describe('State', function()
+  before_each(function()
+    -- Ensure the cache file is removed before each run
+    state = require('orgmode.state.state')
+    state:wipe()
+    vim.fn.delete(cache_path, 'rf')
+  end)
+
+  it("should create a state file if it doesn't exist", function()
+    local stat = vim.loop.fs_stat(cache_path)
+    if stat then
+      error('Cache file existed before it should! Ensure it is deleted before each test run!')
+    end
+
+    -- This creates the cache file on new instances of `State`
+    state:load()
+
+    -- wait until the state has been saved
+    vim.wait(50, function()
+      return state._ctx.saved
+    end, 10)
+
+    local stat, err, _ = vim.loop.fs_stat(cache_path)
+    if not stat then
+      error(err)
+    end
+  end)
+
+  it('should save the cache file as valid json', function()
+    local data = nil
+    local read_f_err = nil
+    state:save():next(function()
+      utils
+        .readfile(cache_path, { raw = true })
+        :next(function(state_data)
+          data = state_data
+        end)
+        :catch(function(err)
+          read_f_err = err
+        end)
+    end)
+
+    -- wait until the newly saved state file has been read
+    vim.wait(50, function()
+      return data ~= nil or read_f_err ~= nil
+    end, 10)
+    if read_f_err then
+      error(read_f_err)
+    end
+
+    local success, decoded = pcall(vim.json.decode, data, {
+      luanil = { object = true, array = true },
+    })
+    if not success then
+      error('Cache file did not contain valid json after saving! Error: ' .. vim.inspect(decoded))
+    end
+  end)
+
+  it('should be able to save and load state data', function()
+    -- Set a variable into the state object
+    state.my_var = 'hello world'
+    -- Save the state
+    state:save_sync()
+    -- Wipe the variable and "unload" the State
+    state.my_var = nil
+    state._ctx.loaded = false
+
+    -- Ensure the state can be loaded from the file now by ignoring the previous load
+    state:load_sync()
+    -- These should be the same after the wipe. We just loaded it back in from the state cache.
+    assert.are.equal('hello world', state.my_var)
+  end)
+
+  it('should be able to self-heal from an invalid state file', function()
+    state:save_sync()
+
+    -- Mangle the cache
+    vim.cmd.edit(cache_path)
+    vim.api.nvim_buf_set_lines(0, 0, -1, false, { '[ invalid json!' })
+    vim.cmd.write()
+
+    -- Ensure we reload the state from its cache file (this should also "heal" the cache)
+    state._ctx.loaded = false
+    state._ctx.saved = false
+    state:load_sync()
+    vim.wait(500, function()
+      return state._ctx.saved
+    end)
+    -- Wait a little longer to ensure the data is flushed into the cache
+    vim.wait(100)
+
+    -- Now attempt to read the file and check that it is, in fact, "healed"
+    local cache_data = nil
+    local read_f_err = nil
+    utils
+      .readfile(cache_path, { raw = true })
+      :next(function(state_data)
+        cache_data = state_data
+      end)
+      :catch(function(reject)
+        read_f_err = reject
+      end)
+      :finally(function()
+        read_file = true
+      end)
+
+    vim.wait(500, function()
+      return cache_data ~= nil or read_f_err ~= nil
+    end, 20)
+
+    if read_f_err then
+      error('Unable to read the healed state cache! Error: ' .. vim.inspect(read_f_err))
+    end
+
+    local success, decoded = pcall(vim.json.decode, cache_data, {
+      luanil = { object = true, array = true },
+    })
+
+    if not success then
+      error(
+        'Unable to self-heal from an invalid state! Error: '
+          .. vim.inspect(decoded)
+          .. '\n\t-> Got cache content as '
+          .. vim.inspect(cache_data)
+      )
+    end
+  end)
+end)