Hibernation API Memory Management

[helloimalastair]

It was pointed out to me today by @MellowYarker that using the Hibernatable WebSocket API does not preclude running out of memory due to WebSocket objects not being automatically deleted when no longer referenced in user code. While at a smaller scale, this wouldn't be much of an issue, when you get into the hundreds/thousands of WebSockets being hibernated, it could become a real problem, especially if your DO isn't evicted very often.

We ended on going on a discussion on possible fixes, here are two we could come up with:

• Runtime tracking of WebSocket object references. This would allow for the runtime to detect when there was no longer any user code referencing a WebSocket, then garbage collecting it.
• A state.evict API, which would allow user code to manually trigger an eviction. I had envisioned this as a non-blocking API, that would cause incoming events to be queued, rather than passed in to the DO. Then, once all currently processing events had completed, you could evict the DO from memory, allowing it to be restarted fresh. The evict function could also include the ability to pass a function that would be called right before eviction, allowing you to flush to storage(this would only apply for scheduled evictions called by user code, not unscheduled evictions caused by runtime/code updates).

Neither of these options are entirely ideal. The runtime tracking requires some finicky manipulation of object lifetimes, and the eviction API would have to be limited in what it can do to prevent it from holding up the DO from restarting in a timely manner.

I'm curious to hear if y'all have any other ideas for how management of WebSocket objects in regards to the Hibernation API could be handled efficiently.

[kentonv]

I don't quite understand the concern here. The WebSocket objects should go away when they are closed; if not, that is a bug we should fix. Assuming no such bug, can't the app just close connections if it needs to free up memory?

FWIW these objects are not very big, I think you'd need hundreds of thousands of them in order to impact the memory limit.

On another note, if you want to force eviction of your DO, you can do so by doing blockConcurrencyWhile() and throwing an exception inside the callback. This resets the object. We're also thinking of adding an abort() call to do this more directly. But these are really meant for situations where you've detected your object is in an inconsistent state (like, due to bugs) and want to turn it off and back on again to fix it. Not so much for memory management.

[helloimalastair]

This is not from personal inspection of the source code, but from my conversation with Milan in the Discord. Specifically

You still have to worry about the memory limit if, for example, your DO never gets evicted. If your DO is always active we won't get a chance to cleanup your unused JS WebSocket objects so (I would think) they kinda just stick around. If your DO has periods of inactivity then we'll only recreate the JS WebSocket as needed (but again it sticks around until your DO hibernates again).

- Milan, CF Dev Discord

As I understood it, when the runtime needs to make a Hibernated WebSocket object available to user code(for an incoming event, or a getWebSocket()), it constructs a WebSocket object, and passes it in. However, even when it goes out of scope in the user code, it is still stored somewhere else in the Runtime Code, and thus is not garbage-collected. This could be solved by performing a manual eviction/closing connections like you said. Though, if the objects really are that small, then I guess it wouldn't matter that much. Now I'm curious to try to measure the size of one of the objects, but unfortunately I have never tried running a debugger against wrangler/workerd, so I'm not even entirely sure that you can measure memory like that...

[kentonv]

Oh I see, you are talking about connection that are still live, and you don't want to disconnect them, you just want the JS wrapper objects to go away without the DO actually hibernating. Aborting the object as I described wouldn't accomplish that -- it would actually cause all the WebSocket connections to be disconnected.

I think if we're going to modify the runtime to accomplish what you want, the change we should make is to make these objects garbage-collectable.

Note that while only the JavaScript wrappers count against your memory limit, there are in fact native objects underneath them that don't go away even when your DO hibernates. We currently don't try to estimate the memory usage of these native objects because it'd be tedious to do so. But given that they are there, going out of our way to allow the relatively small JS wrappers to be GC'd seems weird. To the extent that a huge number of WebSockets filling up memory is a problem, the right answer is really to distribute your connections across more objects.

[helloimalastair]

Makes sense. I guess I had imagined that the wrappers had more of an impact on memory than they actually have, so how it is rn sounds fine