Implement Drop for handles / channels in Rust SDK

[tiziano88]

It would be useful to fix it at the Senderand Receiver level (or even better at the ReadHandle and WriteHandle level) if it is feasible. That will also solve a related issue with the router pattern. When the router receives a gRPC invocation, it must remember to close its handles to the request and response channels in the invocation after sending it to the target node to avoid potential channel leaks. In general, any node that receives a proto containing handles must remember to close all of the handles even if it does not do anything with the proto apart from sending it to another node.

Originally posted by @conradgrobler in #1676 (comment)

[tiziano88]

I guess the complication is that Sender and Receiver are defined in the oak_io crate, so I am not sure whether we can impl Drop for them from the oak SDK crate. Any suggestions? cc @project-oak/core

[conradgrobler]

In a java project the answer would have been dependency injection (something like adding an Option<Box<dyn ChannelDropper>> that is used by drop on the sender and receiver if it has a value). But that feels quite messy and not like idiomatic rust.

For #1676 the issue can be addressed by updating the generated client code to to close the invocation sender in drop

I am not sure how to address the more general issue when sending protos that contain handles. In principle we could always close the handles that are being extracted by HandleVisit, but this might be unexpected behaviour for an implementer.

[tiziano88]

I am slightly reluctant to add special handling for closing channels here and there as a best effort. Especially for gRPC clients, which I think are not really expected to be dropped throughout the execution of an application. Also the channel is already there (e.g. as returned from the node creation call), so it feels odd that it will only be correctly closed if it is actually used as part of a gRPC client, but if not it will be left to leak.

[tiziano88]

Looking some more into it, I think it may not be possible to impl Drop for oak_io::Sender, but perhaps we may create another wrapper type in oak::io (i.e. in the SDK) that impl Drop correctly and also impl Deref<oak_io::Sender> so that most methods will just keep working on the new type without having to manually forward them. Then we would make oak::io::channel_create and other methods return the more specific type. I am not sure how to handle protobuf though, apart from having two options for code generation, one for native, one for SDK. To clarify: this solution would only work for the Rust SDK.

[tiziano88]

@wildarch do you have any suggestions?

[wildarch]

I'm trying to understand the blocker here for a direct impl Drop for Sender, is the issue that we define Sender in oak_io, but channel_close in oak (which is built on top of oak_io), so that from oak_io we have no function available to close the channel?

[tiziano88]

Yes, Sender is used both in the Oak Runtime and also as part of the SDK, so there is no single impl of Drop that works for both. We only want to define it when it is used in the SDK.

[wildarch]

My first instinct was to impl Drop through some proxy trait, but that is not allowed by the compiler (for good reasons), Drop is too special for that.

Since we need to run a different destructor depending on the environment we use Sender in, I think the best way is probably a wrapper type, as @tiziano88 suggested earlier. That said, it looks to me like most of the code we have for Sender in oak_io does not actually leverage the type safety that it provides over WriteHandle, so maybe we could consider moving Receiver and Sender up to the SDK/runtime level, and add impl HandleVisit and impl Message on ReadHandle/WriteHandle instead?

[tiziano88]

The reason for having Sender and Receiver in oak_io was that the same generated code may be used in both the runtime and in the SDK, and then each of them uses extension traits to define operations on the same underlying types. I think what you are suggesting would require generating different code for runtime and SDK, right? (which is not necessarily a bad thing, maybe it would be simpler overall)

WDYT?

[wildarch]

Correct, yes. The logic that can be shared would live in the ReadHandle and WriteHandle logic, i.e. that is the HandleVisit and Message impls.

The logic that is currently in the extension methods would live in the Sender and Receiver types, for which we would have distinct types in the SDK and runtime. That way these types have access to the SDK and runtime functions respectively, and thus we can implement Drop for them. It does require that we create types with essentially the same layout, but that is only minimal duplication IMO.

[wildarch]

I've been experimenting with a few ideas here:

• Create a wrapper type in the SDK. Difficult because handles passed through protobufs cannot be converted automatically.
• Add a ChannelCloser trait as part of the Receiver and Sender definition. This causes problems in the generated code for protobufs
• Pass a channel close function in Receiver::new. Problematic because it is not compatible with Decodable
• Call oak_abi::channel_close in drop and implement this in oak_runtime, a bit weird but might work.

Now I have noticed that ReadHandle, WriteHandle, Receiver and Sender all #[derive(Copy, Clone)], which raises questions as to what should happen on Drop. if these handles are easily copied, closing the channel when just one is dropped probably does not make sense.

Do we want to consider doing more bookkeeping to keep track of individual copies of handles? We could add ABI calls like handle_clone and handle_drop to update refcounts. Then we can remove Copy, keep Clone and add Drop in a way that makes sense.

[tiziano88]

Yes I would expect these types should stop being Copy, and we may provide a manual impl of Clone that actually dups the handle via a new ABI call (though for the time being we will just disallow cloning, I guess).

[tiziano88]

handle_drop is essentially channel_close, right?

[wildarch]

Aah yes, good point, we can just reuse channel_close.

[wildarch]

Here is my game plan for making handles linear types:

1. Add handle_clone syscall (Add handle_clone to ABI and runtime. #1782)
2. Allow SDK functions (channel_read, channel_write) to use borrowed handles instead of owned ones (Allow passing borrowed handles to SDK functions. #1843)
3. Add linear-handles feature to oak_io that optionally makes handle types defined there linear (Add linear-handles feature to oak_io #1850)
4. Add linear-handles feature to SDK for compatibility with the oak_io feature
5. Migrate individual example modules to use the new feature
6. Migrate runtime to use linear handles
7. Remove the features and make their behaviour the default

Of course we could also do all of this at once (e.g. https://github.com/wildarch/oak/tree/handle_clone_drop_impls), but I suspect that would be a very tough PR to land, because it touches so much of the codebase.