Short-Lived, Tick-Bound Memoization in JavaScript

I'm always looking for excuses to use queueMicrotask(), a relatively sparsely used function for inserting tasks at the end of the current call stack. I might've stumbled across one dealing with short-lived memoization. Humor me for a bit.

You've likely seen some variation of this memoization function in JavaScript. Pass in an expensive function, and you get back a different function that'll cache the result on first execution. This particular example even caches according to the parameters you use:

function memoize(fn) {
  const cache = new Map();

  return function (...args) {
    // Key created from arguments.
    // Note: only works for serializable values.
    const key = JSON.stringify(args);
    
    if (cache.has(key)) {
      return cache.get(key);
    }

    const result = fn.apply(this, args);
    cache.set(key, result);

    return result;
  };
}

const memoizedExpensiveFunction = memoize(expensiveFunction);

const freshResult = memoizedExpensiveFunction();
const cachedResult = memoizedExpensiveFunction();

That approach will cover a good share of use cases (probably most), and it doesn't pose many problems when you have clear control over the inputs and you know when they need to change.

But JavaScript runs on an event loop, and a lot can change between different lines of code in the same file, depending on how you write it. This will all run in a single turn ("tick") of the event loop:

doThing();
doAnotherThing();
doOneMoreThing();

The browser can do nothing until all of those functions are run through the call stack and executed. But you can schedule tasks to run on future ticks of the event loop too.

// executed on the current stack:
doThing();

// queued ASAP:
scheduler.postTask(() => doItAgain());

// queued at some point soon in the future:
setTimeout(() => doAnotherThing());

// queued just before the next repaint:
requestAnimationFrame(() => doOneMoreThing());

The browser is doing a lot more than just executing your code here, meaning the entire world could change between each of those lines. A repaint could've occurred, the DOM could have been reflowed, user events could've been handled, timer callbacks could've been executed, and more.

If you're memoizing a value based on parameters that could change across ticks, you might need more fine-grained control over when your cached is considered "stale."

Frequently Stale Memoization

That particular memoization function above works just fine across multiple ticks because no external inputs change and everything fires in the same execution context. These will all return the same random number:

function getRandom() {
  return Math.random();
}

const memoizedGetRandom = memoize(getRandom);

console.log('Task #1:', memoizedGetRandom());

setTimeout(() => {
  console.log('Task #2:', memoizedGetRandom());

  setTimeout(() => {
    console.log('Task #3:', memoizedGetRandom());
  });
});

But binding that cache to the execution context becomes a problem when variables start to change between ticks, like in the case below. This example memoizes a perimeter calculation based on screen size, which will change as soon as the user resizes the window. But you'll get the wrong answer because the cache has become stale.

function calculateWindowPerimeter() {
  return 2 * (window.innerWidth + window.innerHeight);
}

const memoizedCalculateWindowPerimeter = memoize(calculateWindowPerimeter);

setInterval(() => {
  console.log('Window Perimeter:', memoizedCalculateWindowPerimeter());
}, 1000);

Cases like that are simple enough to fix. Either clear the cache when the window resizes, or use the dimensions as your cache key (although you will need to worry about an incrementally growing memoization cache... different problem). Let's dream up a crazier example.

When Everything's Always Changing

You're making a Chrome extension that calculates a bunch of metrics about the current state of the DOM on any page. You're crawling through a lot of nodes and doing some expensive math:

const {
  maxDepth,
  averageDepth,
  totalElements,
  totalPageArea,
  largestElement,
  deepestElement,
  visibleElements,
  oversizedImages,
  offscreenElements,
  largestElementArea,
  overflowingElements,
} = calculatePerformanceMetrics();

These metrics are being recalculated every so often, and then being visualized in a series of ways – charts, graphs, text, etc. That calculation is ripe for memoization, but to keep it accurate, it'll need to be invalidated at the appropriate times.

Because there are so many factors in play, it's reasonable is to assume that key attributes of the DOM may change at any chance they get. With each tick of the event loop, you might be getting entirely different DOM altogether. You just can't know.

You want that invalidation process to be simple and self-contained. Needing to reach in and wipe it externally would just be another hassle.

Building Tick-Bound Memoization

There are a few options for blowing away a short-lived cache around the event loop. But most won't adequately do the job. Let's camp out on one of them – setTimeout(). You could use this to schedule a cache wipe right away in the future:

function memoizeForTick(fn) {
  const cache = new Map();

  // On the next tick, wipe the cache.
  setTimeout(() => {
    cache.clear();
  });

  return function (...args) {
    const key = JSON.stringify(args);
    
    if (cache.has(key)) {
      return cache.get(key);
    }

    const result = fn.apply(this, args);
    cache.set(key, result);

    return result;
  };
}

But that task isn't necessarily scheduled for the top of the next tick. Depending on the browser and tab conditions, there may be a minimum delay, even if you if explicitly pass 0. That leaves plenty of time for other tasks to be scheduled before your timer's callback, risking a stale cache. Another contrived example:

You have a function that counts elements in the DOM. When the page loads, you get that count. But there are also some other tasks scheduled to access the that same count (I'm using another setTimeout() to illustrate this, but it could be a number of things):

const memoizedCountElements = memoizeForTick(() => {
    return document.querySelectorAll('*').length;
});

// ⬇️ Other arbitrary, scheduled task needing
// a current count of the DOM:
setTimeout(() => {
    console.log('Second count:', memoizedCountElements());
});

// ⬇️ Sneaky task that modifies the DOM 
// at the beginning of the call stack:
scheduler.postTask(() => {
    document.body.appendChild(document.createElement('span'));
}, { priority: 'user-blocking' });

console.log('First count:', memoizedCountElements());

// First count: 8
// Second count: 8

In this case, there's no issue because there are no surprise DOM modifications getting in the way. Still, your scheduled cache.clear() is still technically happening at some point the future – not right away in the next tick.

That's not good when other tasks happen to get priority over your cache clean-up.

const memoizedCountElements = memoizeForTick(() => {
  return document.querySelectorAll('*').length;
});

// ⬇️ Other arbitrary, scheduled task needing
// a current count of the DOM:
setTimeout(() => {
    console.log('Second count:', memoizedCountElements());
});

+ // ⬇️ Sneaky task that modifies the DOM 
+ // at the beginning of the call stack:
+ scheduler.postTask(() => {
+    document.body.appendChild(document.createElement('span'));
+ }, { priority: 'user-blocking' });

console.log('First count:', memoizedCountElements());

// First count: 8
// Second count: 8

That's a high-priority task that's slipped in before your timeout's callback, free to do whatever it's instructed to do – like mess with the DOM.

And with that, the next count is incorrect. It's still relying on the previous tick's cache.

First count: 8
Second count: 8 # <-- should be 9! 

This same risk exists with other tools like requestAnimationFrame(). There's just too much going on to delay that invalidation past the current tick. If we want that cache to be reliably cleared, it'll need to be done before anything has a chance to control the event loop again.

Guaranteed Clean-Up

Cue the hero of this post: queueMicrotask(). Like I mentioned, this function allows us to tack on a task at the very end of the current call stack. It's the call stack telling the event loop to "hold up and take care of this one more thing" before moving on to other responsibilities. Let's drop it into our memoization function:

function memoizeForTick(fn) {
  const cache = new Map();

  return function (...args) {
    // Clear the cache right after the 
    // call stack is emptied.
    queueMicrotask(() => {
      cache.clear();
    });
    
    const key = JSON.stringify(args);
    
    if (cache.has(key)) {
      return cache.get(key);
    }

    const result = fn.apply(this, args);
    cache.set(key, result);

    return result;
  };
}

That little change prevents any other tasks from ever getting the stale cache when they shouldn't. Nothing can get at that stale cache because nothing else can happen until that call stack is completely empty.

Another Quick Verification

If you ran this updated code, you'd see that the count is now accurate after the DOM is sneakily changed. But for more verification, we can tweak that getRandom() example from earlier:

const memoizedGetRandom = memoizeForTick(getRandom);

console.log('Task #1:', memoizedGetRandom());
console.log('Task #1:', memoizedGetRandom());   

setTimeout(() => {
  console.log('Task #2:', memoizedGetRandom());
  console.log('Task #2:', memoizedGetRandom());

  setTimeout(() => {
    console.log('Task #3:', memoizedGetRandom());
    console.log('Task #3:', memoizedGetRandom());
    
    setTimeout(() => {
      console.log('Task #4:', memoizedGetRandom());
      console.log('Task #4:', memoizedGetRandom());
    });
  });
});

// Task #1: 0.34272666642062677
// Task #1: 0.34272666642062677

// Task #2: 0.8319017459368117
// Task #2: 0.8319017459368117

// Task #3: 0.6902789119580173
// Task #3: 0.6902789119580173

// Task #4: 0.031707814137348356
// Task #4: 0.031707814137348356

Exactly what we'd expect: memoized within the same tick, but not a microsecond loger.

Optional Micro-Optimization

You might've noticed that because queueMicrotask() is being called on every invocation of your memoized function, you're actually scheduling multiple cache clean-ups per tick. I'm hesitant to say that poses any serious downsides (.clear() is fast and side-effect free, given where it's placed), but you check for if you've already scheduled an invalidation as well:

function memoizeForTick(fn) {
  const cache = new Map();
+  let invalidationScheduled = false;

  return function (...args) {
+    if(!invalidationScheduled) {   
+      queueMicrotask(() => {
+        cache.clear();
+        invalidationScheduled = false;
+      });
+
+      invalidationScheduled = true;
+    }

    const key = JSON.stringify(args);
        
    if (cache.has(key)) {
      return cache.get(key);
    }

    const result = fn.apply(this, args);
    cache.set(key, result);

    return result;
  };
}

Regardless, you're ready to memoize the counting of elements or some gnarly performance calculations, regardless of what the browser throws at you. Confidence feels good.

Worth the Thinking

The circumstances of such short-lived, granular memoization being useful are undoubtedly rare, but wrestling with imagined use cases like this is a worthy thought experiment. You end feeling more comfortable with the many tools at your disposal, equipping you to more often reach for the right one at precisely the right time. Plus, having even the slightest better understanding of what's going on every time someone opens a browser pays dividends over time, and may even save your sanity when chasing down an obscure, inexplicable bug in the future.

Until next tick!